[Core] Accelerator type API (#10561)

This commit is contained in:
Alex Wu
2020-09-06 20:58:40 -07:00
committed by GitHub
parent a699f6a4d8
commit d6a9f0e2e4
10 changed files with 133 additions and 30 deletions
+11 -3
View File
@@ -313,9 +313,10 @@ def set_cuda_visible_devices(gpu_ids):
def resources_from_resource_arguments(
default_num_cpus, default_num_gpus, default_memory,
default_object_store_memory, default_resources, runtime_num_cpus,
runtime_num_gpus, runtime_memory, runtime_object_store_memory,
runtime_resources):
default_object_store_memory, default_resources,
default_accelerator_type, runtime_num_cpus, runtime_num_gpus,
runtime_memory, runtime_object_store_memory, runtime_resources,
runtime_accelerator_type):
"""Determine a task's resource requirements.
Args:
@@ -376,6 +377,13 @@ def resources_from_resource_arguments(
resources["object_store_memory"] = ray_constants.to_memory_units(
object_store_memory, round_up=True)
if runtime_accelerator_type is not None:
resources[f"{ray_constants.RESOURCE_CONSTRAINT_PREFIX}"
f"{runtime_accelerator_type}"] = 0.001
elif default_accelerator_type is not None:
resources[f"{ray_constants.RESOURCE_CONSTRAINT_PREFIX}"
f"{default_accelerator_type}"] = 0.001
return resources