From fb074da7c374a69f04a1add9e62b96ec471d32c3 Mon Sep 17 00:00:00 2001 From: Ian Rodney Date: Tue, 30 Jun 2020 00:55:09 -0700 Subject: [PATCH] [core] Return all GPUs for Local Mode (#9108) --- python/ray/tests/test_advanced_2.py | 20 ++++++++++++++++++++ python/ray/worker.py | 4 ++++ 2 files changed, 24 insertions(+) diff --git a/python/ray/tests/test_advanced_2.py b/python/ray/tests/test_advanced_2.py index c3ef1bc90..e74816dff 100644 --- a/python/ray/tests/test_advanced_2.py +++ b/python/ray/tests/test_advanced_2.py @@ -661,6 +661,26 @@ def test_specific_gpus(save_gpu_ids_shutdown_only): ray.get([g.remote() for _ in range(100)]) +def test_local_mode_gpus(save_gpu_ids_shutdown_only): + allowed_gpu_ids = [4, 5, 6, 7, 8] + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( + [str(i) for i in allowed_gpu_ids]) + + from importlib import reload + reload(ray.worker) + + ray.init(num_gpus=3, local_mode=True) + + @ray.remote + def f(): + gpu_ids = ray.get_gpu_ids() + assert len(gpu_ids) == 3 + for gpu in gpu_ids: + assert gpu in allowed_gpu_ids + + ray.get([f.remote() for _ in range(100)]) + + def test_blocking_tasks(ray_start_regular): @ray.remote def f(i, j): diff --git a/python/ray/worker.py b/python/ray/worker.py index e7fa68574..6d3b195dc 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -395,6 +395,10 @@ def get_gpu_ids(): assigned_ids = [ global_worker.original_gpu_ids[gpu_id] for gpu_id in assigned_ids ] + # Give all GPUs in local_mode. + if global_worker.mode == LOCAL_MODE: + max_gpus = global_worker.node.get_resource_spec().num_gpus + return global_worker.original_gpu_ids[:max_gpus] return assigned_ids