Support metadata for passing by value task arguments (#5527)

This commit is contained in:
Kai Yang
2019-09-08 11:07:48 +08:00
committed by Hao Chen
parent cb7102f31e
commit d8f5804690
27 changed files with 364 additions and 244 deletions
+5 -3
View File
@@ -54,8 +54,10 @@ cdef extern from "ray/common/task/task_spec.h" namespace "ray" nogil:
int ArgIdCount(uint64_t arg_index) const
CObjectID ArgId(uint64_t arg_index, uint64_t id_index) const
CObjectID ReturnId(uint64_t return_index) const
const uint8_t *ArgVal(uint64_t arg_index) const
size_t ArgValLength(uint64_t arg_index) const
const uint8_t *ArgData(uint64_t arg_index) const
size_t ArgDataSize(uint64_t arg_index) const
const uint8_t *ArgMetadata(uint64_t arg_index) const
size_t ArgMetadataSize(uint64_t arg_index) const
double GetRequiredResource(const c_string &resource_name) const
const ResourceSet GetRequiredResources() const
const ResourceSet GetRequiredPlacementResources() const
@@ -86,7 +88,7 @@ cdef extern from "ray/common/task/task_util.h" namespace "ray" nogil:
TaskSpecBuilder &AddByRefArg(const CObjectID &arg_id)
TaskSpecBuilder &AddByValueArg(const c_string &data)
TaskSpecBuilder &AddByValueArg(const c_string &data, const c_string &metadata)
TaskSpecBuilder &SetActorCreationTaskSpec(
const CActorID &actor_id, uint64_t max_reconstructions,
+10 -4
View File
@@ -12,6 +12,7 @@ from ray.includes.task cimport (
TaskSpecBuilder,
TaskTableData,
)
from ray.ray_constants import RAW_BUFFER_METADATA
from ray.utils import decode
@@ -68,10 +69,12 @@ cdef class TaskSpec:
for arg in arguments:
if isinstance(arg, ObjectID):
builder.AddByRefArg((<ObjectID>arg).native())
elif isinstance(arg, bytes):
builder.AddByValueArg(arg, RAW_BUFFER_METADATA)
else:
pickled_str = pickle.dumps(
arg, protocol=pickle.HIGHEST_PROTOCOL)
builder.AddByValueArg(pickled_str)
builder.AddByValueArg(pickled_str, b'')
if not actor_creation_id.is_nil():
# Actor creation task.
@@ -180,9 +183,12 @@ cdef class TaskSpec:
arg_list.append(
ObjectID(task_spec.ArgId(i, 0).Binary()))
else:
serialized_str = (
task_spec.ArgVal(i)[:task_spec.ArgValLength(i)])
obj = pickle.loads(serialized_str)
data = (task_spec.ArgData(i)[:task_spec.ArgDataSize(i)])
metadata = (task_spec.ArgMetadata(i)[:task_spec.ArgMetadataSize(i)])
if metadata == RAW_BUFFER_METADATA:
obj = data
else:
obj = pickle.loads(data)
arg_list.append(obj)
elif lang == <int32_t>LANGUAGE_JAVA:
arg_list = num_args * ["<java-argument>"]