mirror of
https://github.com/wassname/ray.git
synced 2026-06-28 01:16:06 +08:00
[core worker] Python core worker normal task submission (#5566)
This commit is contained in:
+66
-5
@@ -27,6 +27,8 @@ from ray.includes.common cimport (
|
||||
CRayObject,
|
||||
CRayStatus,
|
||||
CGcsClientOptions,
|
||||
CTaskArg,
|
||||
CRayFunction,
|
||||
LocalMemoryBuffer,
|
||||
LANGUAGE_CPP,
|
||||
LANGUAGE_JAVA,
|
||||
@@ -46,7 +48,7 @@ from ray.includes.unique_ids cimport (
|
||||
CObjectID,
|
||||
CClientID,
|
||||
)
|
||||
from ray.includes.libcoreworker cimport CCoreWorker
|
||||
from ray.includes.libcoreworker cimport CCoreWorker, CTaskOptions
|
||||
from ray.includes.task cimport CTaskSpec
|
||||
from ray.includes.ray_config cimport RayConfig
|
||||
from ray.exceptions import RayletError, ObjectStoreFullError
|
||||
@@ -96,6 +98,13 @@ cdef int check_status(const CRayStatus& status) nogil except -1:
|
||||
raise RayletError(message)
|
||||
|
||||
|
||||
cdef VectorToObjectIDs(const c_vector[CObjectID] &object_ids):
|
||||
result = []
|
||||
for i in range(object_ids.size()):
|
||||
result.append(ObjectID(object_ids[i].Binary()))
|
||||
return result
|
||||
|
||||
|
||||
cdef c_vector[CObjectID] ObjectIDsToVector(object_ids):
|
||||
"""A helper function that converts a Python list of object IDs to a vector.
|
||||
|
||||
@@ -244,6 +253,16 @@ cdef unordered_map[c_string, double] resource_map_from_dict(resource_map):
|
||||
return out
|
||||
|
||||
|
||||
cdef c_vector[c_string] string_vector_from_list(list string_list):
|
||||
cdef:
|
||||
c_vector[c_string] out
|
||||
for s in string_list:
|
||||
if not isinstance(s, bytes):
|
||||
raise TypeError("string_list elements must be bytes")
|
||||
out.push_back(s)
|
||||
return out
|
||||
|
||||
|
||||
cdef class RayletClient:
|
||||
cdef CRayletClient* client
|
||||
|
||||
@@ -400,6 +419,10 @@ cdef class CoreWorker:
|
||||
"outside _raylet. See __init__.py for "
|
||||
"details.")
|
||||
|
||||
def disconnect(self):
|
||||
with nogil:
|
||||
self.core_worker.get().Disconnect()
|
||||
|
||||
def get_objects(self, object_ids, TaskID current_task_id):
|
||||
cdef:
|
||||
c_vector[shared_ptr[CRayObject]] results
|
||||
@@ -538,6 +561,48 @@ cdef class CoreWorker:
|
||||
with nogil:
|
||||
self.core_worker.get().SetCurrentJobId(c_job_id)
|
||||
|
||||
def submit_task(self,
|
||||
function_descriptor,
|
||||
args,
|
||||
int num_return_vals,
|
||||
resources):
|
||||
cdef:
|
||||
unordered_map[c_string, double] c_resources
|
||||
CTaskOptions task_options
|
||||
CRayFunction ray_function
|
||||
c_vector[CTaskArg] args_vector
|
||||
c_vector[CObjectID] return_ids
|
||||
c_string pickled_str
|
||||
shared_ptr[CBuffer] arg_data
|
||||
shared_ptr[CBuffer] arg_metadata
|
||||
|
||||
c_resources = resource_map_from_dict(resources)
|
||||
task_options = CTaskOptions(num_return_vals, c_resources)
|
||||
ray_function = CRayFunction(
|
||||
LANGUAGE_PYTHON, string_vector_from_list(function_descriptor))
|
||||
|
||||
for arg in args:
|
||||
if isinstance(arg, ObjectID):
|
||||
args_vector.push_back(
|
||||
CTaskArg.PassByReference((<ObjectID>arg).native()))
|
||||
else:
|
||||
pickled_str = pickle.dumps(
|
||||
arg, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
arg_data = dynamic_pointer_cast[CBuffer, LocalMemoryBuffer](
|
||||
make_shared[LocalMemoryBuffer](
|
||||
<uint8_t*>(pickled_str.data()),
|
||||
pickled_str.size(),
|
||||
True))
|
||||
args_vector.push_back(
|
||||
CTaskArg.PassByValue(
|
||||
make_shared[CRayObject](arg_data, arg_metadata)))
|
||||
|
||||
with nogil:
|
||||
check_status(self.core_worker.get().Tasks().SubmitTask(
|
||||
ray_function, args_vector, task_options, &return_ids))
|
||||
|
||||
return VectorToObjectIDs(return_ids)
|
||||
|
||||
def set_object_store_client_options(self, c_string client_name,
|
||||
int64_t limit_bytes):
|
||||
with nogil:
|
||||
@@ -552,7 +617,3 @@ cdef class CoreWorker:
|
||||
message = self.core_worker.get().Objects().MemoryUsageString()
|
||||
|
||||
return message.decode("utf-8")
|
||||
|
||||
def disconnect(self):
|
||||
with nogil:
|
||||
self.core_worker.get().Disconnect()
|
||||
|
||||
@@ -2,11 +2,13 @@ from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport shared_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
|
||||
from libc.stdint cimport uint8_t
|
||||
from libc.stdint cimport uint8_t, uint64_t, int64_t
|
||||
from libcpp.unordered_map cimport unordered_map
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
|
||||
from ray.includes.unique_ids cimport (
|
||||
CActorID,
|
||||
CActorHandleID,
|
||||
CJobID,
|
||||
CWorkerID,
|
||||
CObjectID,
|
||||
@@ -138,15 +140,57 @@ cdef extern from "ray/common/buffer.h" namespace "ray" nogil:
|
||||
size_t Size() const
|
||||
|
||||
cdef cppclass LocalMemoryBuffer(CBuffer):
|
||||
LocalMemoryBuffer(uint8_t *data, size_t size)
|
||||
LocalMemoryBuffer(uint8_t *data, size_t size, c_bool copy_data)
|
||||
|
||||
cdef extern from "ray/core_worker/store_provider/store_provider.h" nogil:
|
||||
cdef extern from "ray/common/ray_object.h" nogil:
|
||||
cdef cppclass CRayObject "ray::RayObject":
|
||||
const shared_ptr[CBuffer] &GetData()
|
||||
const size_t DataSize() const
|
||||
const shared_ptr[CBuffer] &GetMetadata() const
|
||||
c_bool HasData() const
|
||||
c_bool HasMetadata() const
|
||||
const size_t DataSize() const
|
||||
const shared_ptr[CBuffer] &GetData()
|
||||
const shared_ptr[CBuffer] &GetMetadata() const
|
||||
|
||||
cdef extern from "ray/core_worker/common.h" nogil:
|
||||
cdef cppclass CRayFunction "ray::RayFunction":
|
||||
CRayFunction()
|
||||
CRayFunction(CLanguage language,
|
||||
const c_vector[c_string] function_descriptor)
|
||||
CLanguage GetLanguage()
|
||||
c_vector[c_string] GetFunctionDescriptor()
|
||||
|
||||
cdef cppclass CTaskArg "ray::TaskArg":
|
||||
@staticmethod
|
||||
CTaskArg PassByReference(const CObjectID &object_id)
|
||||
|
||||
@staticmethod
|
||||
CTaskArg PassByValue(const shared_ptr[CRayObject] &data)
|
||||
|
||||
cdef extern from "ray/core_worker/task_interface.h" nogil:
|
||||
cdef cppclass CTaskOptions "ray::TaskOptions":
|
||||
CTaskOptions()
|
||||
CTaskOptions(int num_returns,
|
||||
unordered_map[c_string, double] &resources)
|
||||
|
||||
cdef cppclass CActorCreationOptions "ray::ActorCreationOptions":
|
||||
CActorCreationOptions(uint64_t max_reconstructions,
|
||||
const unordered_map[c_string, double] &resources)
|
||||
|
||||
cdef cppclass CActorHandle "ray::ActorHandle":
|
||||
CActorHandle(
|
||||
const CActorID &actor_id, const CActorHandleID &actor_handle_id,
|
||||
const CLanguage actor_language,
|
||||
const c_vector[c_string] &actor_creation_task_function_descriptor)
|
||||
|
||||
CActorHandle(const CActorHandle &other)
|
||||
CActorID ActorID() const
|
||||
CActorHandleID ActorHandleID() const
|
||||
c_vector[c_string] ActorCreationTaskFunctionDescriptor() const
|
||||
CObjectID ActorCursor() const
|
||||
int64_t TaskCursor() const
|
||||
int64_t NumForks() const
|
||||
CActorHandle Fork()
|
||||
void Serialize(c_string *output)
|
||||
CActorHandle Deserialize(const c_string &data)
|
||||
|
||||
cdef extern from "ray/gcs/gcs_client_interface.h" nogil:
|
||||
cdef cppclass CGcsClientOptions "ray::gcs::GcsClientOptions":
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from libc.stdint cimport int64_t
|
||||
from libcpp cimport bool as c_bool
|
||||
from libcpp.memory cimport shared_ptr
|
||||
from libcpp.memory cimport shared_ptr, unique_ptr
|
||||
from libcpp.string cimport string as c_string
|
||||
from libcpp.vector cimport vector as c_vector
|
||||
|
||||
@@ -10,9 +10,14 @@ from ray.includes.unique_ids cimport (
|
||||
CObjectID,
|
||||
)
|
||||
from ray.includes.common cimport (
|
||||
CActorCreationOptions,
|
||||
CActorHandle,
|
||||
CBuffer,
|
||||
CRayStatus,
|
||||
CRayFunction,
|
||||
CRayObject,
|
||||
CRayStatus,
|
||||
CTaskArg,
|
||||
CTaskOptions,
|
||||
CWorkerType,
|
||||
CLanguage,
|
||||
CGcsClientOptions,
|
||||
@@ -20,6 +25,20 @@ from ray.includes.common cimport (
|
||||
from ray.includes.libraylet cimport CRayletClient
|
||||
|
||||
|
||||
cdef extern from "ray/core_worker/task_interface.h" namespace "ray" nogil:
|
||||
cdef cppclass CTaskSubmissionInterface "CoreWorkerTaskInterface":
|
||||
CRayStatus SubmitTask(
|
||||
const CRayFunction &function, const c_vector[CTaskArg] &args,
|
||||
const CTaskOptions &options, c_vector[CObjectID] *return_ids)
|
||||
CRayStatus CreateActor(
|
||||
const CRayFunction &function, const c_vector[CTaskArg] &args,
|
||||
const CActorCreationOptions &options,
|
||||
unique_ptr[CActorHandle] *handle)
|
||||
CRayStatus SubmitActorTask(
|
||||
CActorHandle &handle, const CRayFunction &function,
|
||||
const c_vector[CTaskArg] &args, const CTaskOptions &options,
|
||||
c_vector[CObjectID] *return_ids)
|
||||
|
||||
cdef extern from "ray/core_worker/object_interface.h" nogil:
|
||||
cdef cppclass CObjectInterface "ray::CoreWorkerObjectInterface":
|
||||
CRayStatus SetClientOptions(c_string client_name, int64_t limit)
|
||||
@@ -50,13 +69,13 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
|
||||
CWorkerType &GetWorkerType()
|
||||
CLanguage &GetLanguage()
|
||||
CObjectInterface &Objects()
|
||||
# CTaskSubmissionInterface &Tasks()
|
||||
CTaskSubmissionInterface &Tasks()
|
||||
# CTaskExecutionInterface &Execution()
|
||||
|
||||
# TODO(edoakes): remove this once the raylet client is no longer used
|
||||
# directly.
|
||||
CRayletClient &GetRayletClient()
|
||||
# TODO(edoakes): remove this once the Python core worker uses the task
|
||||
# TODO(edoakes): remove these once the Python core worker uses the task
|
||||
# interfaces
|
||||
void SetCurrentJobId(const CJobID &job_id)
|
||||
void SetCurrentTaskId(const CTaskID &task_id)
|
||||
|
||||
@@ -707,11 +707,14 @@ class Worker(object):
|
||||
self.current_job_id, self.current_task_id,
|
||||
self.task_context.task_index, actor_id)
|
||||
else:
|
||||
# This is a normal task.
|
||||
task_id = TaskID.for_normal_task(self.current_job_id,
|
||||
self.current_task_id,
|
||||
self.task_context.task_index)
|
||||
# Normal tasks are submitted through the core worker (in the
|
||||
# future, all tasks will be).
|
||||
return self.core_worker.submit_task(function_descriptor_list,
|
||||
args_for_raylet,
|
||||
num_return_vals, resources)
|
||||
|
||||
# Actor creation tasks and actor tasks are submitted directly to
|
||||
# the raylet.
|
||||
task = ray._raylet.TaskSpec(
|
||||
task_id,
|
||||
job_id,
|
||||
|
||||
Reference in New Issue
Block a user