mirror of
https://github.com/wassname/ray.git
synced 2026-06-28 12:28:10 +08:00
Support metadata for passing by value task arguments (#5527)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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>"]
|
||||
|
||||
Reference in New Issue
Block a user