wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-13 05:16:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make ray.get_gpu_ids() respect existing CUDA_VISIBLE_DEVICES. (#1499)

Browse Source 
* Make ray.get_gpu_ids() respect existing CUDA_VISIBLE_DEVICES.

* Comment out failing GPUID check.

* Add import.

* Fix test.

* Remove test.

* Factor out environment variable setting/getting into utils.
This commit is contained in:
Robert Nishihara
2018-02-01 21:29:14 -08:00
committed by Philipp Moritz
parent a5b00a545e
commit ed77a4c415
4 changed files with 106 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/services.py
+16
View File
@@ -709,9 +709,25 @@ def start_local_scheduler(redis_address,
# By default, use the number of hardware execution threads for the
# number of cores.
resources["CPU"] = psutil.cpu_count()
# See if CUDA_VISIBLE_DEVICES has already been set.
gpu_ids = ray.utils.get_cuda_visible_devices()
# Check that the number of GPUs that the local scheduler wants doesn't
# excede the amount allowed by CUDA_VISIBLE_DEVICES.
if ("GPU" in resources and gpu_ids is not None and
resources["GPU"] > len(gpu_ids)):
raise Exception("Attempting to start local scheduler with {} GPUs, "
"but CUDA_VISIBLE_DEVICES contains {}.".format(
resources["GPU"], gpu_ids))
if "GPU" not in resources:
# Try to automatically detect the number of GPUs.
resources["GPU"] = _autodetect_num_gpus()
# Don't use more GPUs than allowed by CUDA_VISIBLE_DEVICES.
if gpu_ids is not None:
resources["GPU"] = min(resources["GPU"], len(gpu_ids))
print("Starting local scheduler with the following resources: {}."
.format(resources))
local_scheduler_name, p = ray.local_scheduler.start_local_scheduler(