From fb8bdb8039fe8dde4b3eea0a740c859f9dde7f87 Mon Sep 17 00:00:00 2001 From: Hong Xu Date: Wed, 15 Sep 2021 08:09:49 -0700 Subject: [PATCH] When test set_affinity, don't hardcode the CPU ID (#65042) Summary: The setaffinity test always fails when the number of CPUs is smaller than 3. Changed the test to be dynamically based on the number of CPUs of the system. Pull Request resolved: https://github.com/pytorch/pytorch/pull/65042 Reviewed By: jbschlosser Differential Revision: D30960554 Pulled By: ejguan fbshipit-source-id: 55ac12714b4b0964b48c3617b79a7a345d40ebce --- test/test_dataloader.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/test/test_dataloader.py b/test/test_dataloader.py index e04c2b2..07628f6 100644 --- a/test/test_dataloader.py +++ b/test/test_dataloader.py @@ -1,6 +1,7 @@ import math import sys import errno +import multiprocessing import os import ctypes import faulthandler @@ -2353,7 +2354,7 @@ class SetAffinityDataset(IterableDataset): def worker_set_affinity(_): - os.sched_setaffinity(0, [2]) + os.sched_setaffinity(0, [multiprocessing.cpu_count() - 1]) @unittest.skipIf( @@ -2366,7 +2367,7 @@ class TestSetAffinity(TestCase): dataloader = torch.utils.data.DataLoader( dataset, num_workers=2, worker_init_fn=worker_set_affinity) for sample in dataloader: - self.assertEqual(sample, [2]) + self.assertEqual(sample, [multiprocessing.cpu_count() - 1]) class ConvDataset(Dataset): def __init__(self): -- 2.7.4