Allow scheduling with arbitrary user-defined resource labels. (#1236)

* Enable scheduling with custom resource labels.

* Fix.

* Minor fixes and ref counting fix.

* Linting

* Use .data() instead of .c_str().

* Fix linting.

* Fix ResourcesTest.testGPUIDs test by waiting for workers to start up.

* Sleep in test so that all tasks are submitted before any completes.
This commit is contained in:
Robert Nishihara
2017-12-01 11:41:40 -08:00
committed by Philipp Moritz
parent ac64631043
commit c21e189371
42 changed files with 1073 additions and 806 deletions
+4 -6
View File
@@ -109,9 +109,7 @@ def hex_to_binary(hex_identifier):
FunctionProperties = collections.namedtuple("FunctionProperties",
["num_return_vals",
"num_cpus",
"num_gpus",
"num_custom_resource",
"resources",
"max_calls"])
"""FunctionProperties: A named tuple storing remote functions information."""
@@ -131,7 +129,7 @@ def attempt_to_reserve_gpus(num_gpus, driver_id, local_scheduler,
"""
assert num_gpus != 0
local_scheduler_id = local_scheduler["DBClientID"]
local_scheduler_total_gpus = int(local_scheduler["NumGPUs"])
local_scheduler_total_gpus = int(local_scheduler["GPU"])
success = False
@@ -253,9 +251,9 @@ def select_local_scheduler(driver_id, local_schedulers, num_gpus,
# Loop through all of the local schedulers in a random order.
local_schedulers = np.random.permutation(local_schedulers)
for local_scheduler in local_schedulers:
if local_scheduler["NumCPUs"] < 1:
if local_scheduler["CPU"] < 1:
continue
if local_scheduler["NumGPUs"] < num_gpus:
if local_scheduler.get("GPU", 0) < num_gpus:
continue
if num_gpus == 0:
local_scheduler_id = hex_to_binary(local_scheduler["DBClientID"])