[docker] Detect CPUs in container correctly (#10507)

Co-authored-by: simon-mo <simon.mo@hey.com>
Co-authored-by: Richard Liaw <rliaw@berkeley.edu>
Co-authored-by: Alex Wu <itswu.alex@gmail.com>

python/ray/tune/examples/cifar10_pytorch.py

@@ -12,6 +12,7 @@ import torch.optim as optim
 from torch.utils.data import random_split
 import torchvision
 import torchvision.transforms as transforms
+import ray
 from ray import tune
 from ray.tune import CLIReporter
 from ray.tune.schedulers import ASHAScheduler
@@ -231,6 +232,7 @@ if __name__ == "__main__":
     args, _ = parser.parse_known_args()
 
     if args.smoke_test:
+        ray.init(num_cpus=2)
         main(num_samples=1, max_num_epochs=1, gpus_per_trial=0)
     else:
         # Change this to activate training on GPUs

python/ray/tune/examples/horovod_simple.py

@@ -97,6 +97,9 @@ if __name__ == "__main__":
     parser.add_argument("--hosts-per-trial", type=int, default=1)
     parser.add_argument("--slots-per-host", type=int, default=2)
     args = parser.parse_args()
+    if args.smoke_test:
+        import ray
+        ray.init(num_cpus=2)
 
     # import ray
     # ray.init(address="auto")  # assumes ray is started with ray up

python/ray/tune/examples/pbt_convnet_example.py

@@ -76,7 +76,7 @@ if __name__ == "__main__":
         "--smoke-test", action="store_true", help="Finish quickly for testing")
     args, _ = parser.parse_known_args()
 
-    ray.init()
+    ray.init(num_cpus=2)
     datasets.MNIST("~/data", train=True, download=True)
 
     # check if PytorchTrainble will save/restore correctly before execution

python/ray/tune/tests/ext_pytorch.py

@@ -58,6 +58,7 @@ import torch.optim as optim
 from torch.utils.data import random_split
 import torchvision
 import torchvision.transforms as transforms
+import ray
 from ray import tune
 from ray.tune import CLIReporter
 from ray.tune.schedulers import ASHAScheduler
@@ -434,6 +435,7 @@ def main(num_samples=10, max_num_epochs=10, gpus_per_trial=2):
 
 if __name__ == "__main__":
     # You can change the number of GPUs per trial here:
+    ray.init(num_cpus=2)  # for testing purposes only
     main(num_samples=2, max_num_epochs=2, gpus_per_trial=0)
 
 

python/ray/tune/tests/test_dependency.py

@@ -11,7 +11,8 @@ def f(config, reporter):
 
 
 if __name__ == "__main__":
-    ray.init()
+    ray.init(num_cpus=2)
+
     register_trainable("my_class", f)
     run_experiments({
         "test": {

python/ray/tune/tests/test_ray_trial_executor.py

@@ -16,7 +16,7 @@ from ray.cluster_utils import Cluster
 class RayTrialExecutorTest(unittest.TestCase):
     def setUp(self):
         self.trial_executor = RayTrialExecutor(queue_trials=False)
-        ray.init(ignore_reinit_error=True)
+        ray.init(num_cpus=2, ignore_reinit_error=True)
         _register_all()  # Needed for flaky tests
 
     def tearDown(self):

python/ray/tune/tests/test_sync.py

@@ -17,7 +17,7 @@ from ray.tune.syncer import CommandBasedClient
 
 class TestSyncFunctionality(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()

python/ray/tune/tests/test_track.py

@@ -20,7 +20,7 @@ class TrackApiTest(unittest.TestCase):
     def testSoftDeprecation(self):
         """Checks that tune.track.log code does not break."""
         from ray.tune import track
-        ray.init()
+        ray.init(num_cpus=2)
 
         def testme(config):
             for i in range(config["iters"]):

python/ray/tune/tests/test_trial_runner.py

@@ -23,7 +23,8 @@ class TrialRunnerTest(unittest.TestCase):
         ray.shutdown()
 
     def testTrialStatus(self):
-        ray.init()
+        ray.init(num_cpus=2)
+
         trial = Trial("__fake")
         trial_executor = RayTrialExecutor()
         self.assertEqual(trial.status, Trial.PENDING)
@@ -35,7 +36,7 @@ class TrialRunnerTest(unittest.TestCase):
         self.assertEqual(trial.status, Trial.ERROR)
 
     def testExperimentTagTruncation(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
         def train(config, reporter):
             reporter(timesteps_total=1)

python/ray/tune/tests/test_trial_runner_3.py

@@ -540,7 +540,8 @@ class TrialRunnerTest3(unittest.TestCase):
         runner2.step()
 
     def testCheckpointWithFunction(self):
-        ray.init()
+        ray.init(num_cpus=2)
+
         trial = Trial(
             "__fake",
             config={"callbacks": {
@@ -565,7 +566,8 @@ class TrialRunnerTest3(unittest.TestCase):
                         and fname.endswith(".json"))
                        for fname in os.listdir(cdir))
 
-        ray.init()
+        ray.init(num_cpus=2)
+
         trial = Trial("__fake", checkpoint_freq=1)
         tmpdir = tempfile.mkdtemp()
         runner = TrialRunner(local_checkpoint_dir=tmpdir, checkpoint_period=0)

python/ray/tune/tests/test_trial_scheduler.py

@@ -38,7 +38,7 @@ def mock_trial_runner(trials=None):
 
 class EarlyStoppingSuite(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()
@@ -782,7 +782,7 @@ class _MockTrial(Trial):
 
 class PopulationBasedTestingSuite(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()
@@ -1802,7 +1802,7 @@ class E2EPopulationBasedTestingSuite(unittest.TestCase):
 
 class AsyncHyperBandSuite(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()

python/ray/tune/tests/test_trial_scheduler_pbt.py

@@ -104,7 +104,7 @@ class PopulationBasedTrainingSynchTest(unittest.TestCase):
 
 class PopulationBasedTrainingConfigTest(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()
@@ -145,7 +145,7 @@ class PopulationBasedTrainingConfigTest(unittest.TestCase):
 
 class PopulationBasedTrainingResumeTest(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()

python/ray/tune/tests/test_tune_restore.py

@@ -85,7 +85,7 @@ class TuneRestoreTest(unittest.TestCase):
 
 class TuneExampleTest(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()

python/ray/tune/tests/test_var.py

@@ -15,7 +15,7 @@ from ray.tune.suggest.variant_generator import (RecursiveDependencyError,
 
 class VariantGeneratorTest(unittest.TestCase):
     def setUp(self):
-        ray.init()
+        ray.init(num_cpus=2)
 
     def tearDown(self):
         ray.shutdown()