[Core] Add option to override environment variables for tasks and actors (#11619)

This commit is contained in:
architkulkarni
2020-10-29 14:22:44 -05:00
committed by GitHub
parent e82ff08b0c
commit 4175569d96
16 changed files with 286 additions and 50 deletions
+19 -11
View File
@@ -1001,28 +1001,32 @@ cdef class CoreWorker:
int max_retries,
PlacementGroupID placement_group_id,
int64_t placement_group_bundle_index,
c_bool placement_group_capture_child_tasks):
c_bool placement_group_capture_child_tasks,
override_environment_variables):
cdef:
unordered_map[c_string, double] c_resources
CTaskOptions task_options
CRayFunction ray_function
c_vector[unique_ptr[CTaskArg]] args_vector
c_vector[CObjectID] return_ids
CPlacementGroupID c_placement_group_id = \
placement_group_id.native()
unordered_map[c_string, c_string] \
c_override_environment_variables = \
override_environment_variables
with self.profile_event(b"submit_task"):
prepare_resources(resources, &c_resources)
task_options = CTaskOptions(
name, num_returns, c_resources)
ray_function = CRayFunction(
language.lang, function_descriptor.descriptor)
prepare_args(self, language, args, &args_vector)
with nogil:
CCoreWorkerProcess.GetCoreWorker().SubmitTask(
ray_function, args_vector, task_options, &return_ids,
max_retries, c_pair[CPlacementGroupID, int64_t](
ray_function, args_vector, CTaskOptions(
name, num_returns, c_resources,
c_override_environment_variables),
&return_ids, max_retries,
c_pair[CPlacementGroupID, int64_t](
c_placement_group_id, placement_group_bundle_index),
placement_group_capture_child_tasks)
@@ -1043,7 +1047,8 @@ cdef class CoreWorker:
PlacementGroupID placement_group_id,
int64_t placement_group_bundle_index,
c_bool placement_group_capture_child_tasks,
c_string extension_data
c_string extension_data,
override_environment_variables
):
cdef:
CRayFunction ray_function
@@ -1054,6 +1059,9 @@ cdef class CoreWorker:
CActorID c_actor_id
CPlacementGroupID c_placement_group_id = \
placement_group_id.native()
unordered_map[c_string, c_string] \
c_override_environment_variables = \
override_environment_variables
with self.profile_event(b"submit_task"):
prepare_resources(resources, &c_resources)
@@ -1072,7 +1080,8 @@ cdef class CoreWorker:
c_pair[CPlacementGroupID, int64_t](
c_placement_group_id,
placement_group_bundle_index),
placement_group_capture_child_tasks),
placement_group_capture_child_tasks,
c_override_environment_variables),
extension_data,
&c_actor_id))
@@ -1134,7 +1143,6 @@ cdef class CoreWorker:
cdef:
CActorID c_actor_id = actor_id.native()
unordered_map[c_string, double] c_resources
CTaskOptions task_options
CRayFunction ray_function
c_vector[unique_ptr[CTaskArg]] args_vector
c_vector[CObjectID] return_ids
@@ -1142,7 +1150,6 @@ cdef class CoreWorker:
with self.profile_event(b"submit_task"):
if num_method_cpus > 0:
c_resources[b"CPU"] = num_method_cpus
task_options = CTaskOptions(name, num_returns, c_resources)
ray_function = CRayFunction(
language.lang, function_descriptor.descriptor)
prepare_args(self, language, args, &args_vector)
@@ -1151,7 +1158,8 @@ cdef class CoreWorker:
CCoreWorkerProcess.GetCoreWorker().SubmitActorTask(
c_actor_id,
ray_function,
args_vector, task_options, &return_ids)
args_vector, CTaskOptions(name, num_returns, c_resources),
&return_ids)
return VectorToObjectRefs(return_ids)