wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-06-28 03:34:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

Automatically set CUDA_VISIBLE_DEVICES when worker gets task. (#1044)

Browse Source 
* Automatically set CUDA_VISIBLE_DEVICES when worker gets task.

* Add test.
This commit is contained in:
Robert Nishihara
2017-10-06 18:38:08 -07:00
committed by Philipp Moritz
parent 4669c59fa8
commit a52a1e893f
2 changed files with 59 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/worker.py
+5
View File
@@ -868,6 +868,11 @@ class Worker(object):
"""
with log_span("ray:get_task", worker=self):
task = self.local_scheduler_client.get_task()
# Automatically restrict the GPUs available to this task.
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
[str(i) for i in ray.get_gpu_ids()])
return task
def main_loop(self):