wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-03 02:47:51 +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/worker.py
+22 -6
View File
@@ -233,6 +233,9 @@ class Worker(object):
# The number of threads Plasma should use when putting an object in the
# object store.
self.memcopy_threads = 12
# When the worker is constructed. Record the original value of the
# CUDA_VISIBLE_DEVICES environment variable.
self.original_gpu_ids = ray.utils.get_cuda_visible_devices()
def set_mode(self, mode):
"""Set the mode of the worker.
@@ -868,8 +871,7 @@ class Worker(object):
self.actor_checkpoint_failed = False
# Automatically restrict the GPUs available to this task.
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
[str(i) for i in ray.get_gpu_ids()])
ray.utils.set_cuda_visible_devices(ray.get_gpu_ids())
return task
@@ -889,15 +891,29 @@ class Worker(object):
def get_gpu_ids():
"""Get the IDs of the GPU that are available to the worker.
"""Get the IDs of the GPUs that are available to the worker.
Each ID is an integer in the range [0, NUM_GPUS - 1], where NUM_GPUS is the
number of GPUs that the node has.
If the CUDA_VISIBLE_DEVICES environment variable was set when the worker
started up, then the IDs returned by this method will be a subset of the
IDs in CUDA_VISIBLE_DEVICES. If not, the IDs will fall in the range
[0, NUM_GPUS - 1], where NUM_GPUS is the number of GPUs that the node has.
Returns:
A list of GPU IDs.
"""
if _mode() == PYTHON_MODE:
raise Exception("ray.get_gpu_ids() currently does not work in PYTHON "
"MODE.")
return global_worker.local_scheduler_client.gpu_ids()
assigned_ids = global_worker.local_scheduler_client.gpu_ids()
# If the user had already set CUDA_VISIBLE_DEVICES, then respect that (in
# the sense that only GPU IDs that appear in CUDA_VISIBLE_DEVICES should be
# returned).
if global_worker.original_gpu_ids is not None:
assigned_ids = [global_worker.original_gpu_ids[gpu_id]
for gpu_id in assigned_ids]
return assigned_ids
def _webui_url_helper(client):