serve gen peek

python/ray/_private/worker.py

@@ -35,6 +35,7 @@
     Type,
     TypeVar,
     Union,
+    cast,
     overload,
 )
 from urllib.parse import urlparse
@@ -1402,6 +1403,54 @@ def disconnect(self):
 per worker process.
 """
 
+
+class _CoreWorkerAPI(Protocol):
+    """Structural type for :func:`get_core_worker`.
+
+    This Protocol is intentionally incomplete and is meant to grow as more
+    call sites move off ``Worker.core_worker`` direct access. Only methods
+    used through the helper need to be declared here.
+
+    ``Worker.core_worker`` is attached dynamically on connect (and deleted on
+    shutdown), so static checkers do not see it on ``Worker``.
+    """
+
+    def wait_async(
+        self,
+        object_refs_or_generators: List[Any],
+        num_returns: int,
+        timeout_ms: int,
+        fetch_local: bool,
+        callback: Callable,
+    ) -> int:
+        ...
+
+    def cancel_wait_async(self, handle: int) -> None:
+        ...
+
+
+def get_core_worker() -> _CoreWorkerAPI:
+    """Return the connected worker's CoreWorker.
+
+    The returned value is typed as :class:`_CoreWorkerAPI`, a structural
+    Protocol that currently declares only the wait-async helpers; extend that
+    Protocol when new call sites need additional CoreWorker methods.
+
+    Returns:
+        The process-global CoreWorker attached on ``ray.init()``.
+
+    Raises:
+        RaySystemError: If Ray is not initialized or the core worker is
+            unavailable (same condition as ``Worker.check_connected``).
+    """
+    core_worker = getattr(global_worker, "core_worker", None)
+    if core_worker is None:
+        raise RaySystemError(
+            "Ray has not been started yet. You can start Ray with 'ray.init()'."
+        )
+    return cast(_CoreWorkerAPI, core_worker)
+
+
 _global_node = None
 """ray._private.node.Node: The global node object that is created by ray.init()."""
 
@@ -3220,6 +3269,136 @@ def wait(
         return ready_ids, remaining_ids
 
 
+async def _wait_async(
+    ray_waitables: List[Union[ObjectRef, ObjectRefGenerator]],
+    *,
+    num_returns: int = 1,
+    timeout: Optional[float] = None,
+    fetch_local: bool = True,
+) -> Tuple[
+    List[Union[ObjectRef, ObjectRefGenerator]],
+    List[Union[ObjectRef, ObjectRefGenerator]],
+]:
+    """Async-friendly wait that does not block the event loop.
+
+    Private API. Argument shape and ``(ready, remaining)`` return value match
+    :func:`ray.wait`, but this helper does not auto-init Ray or honor client
+    mode. Unlike ``await obj_ref``, ``fetch_local=False`` waits for readiness
+    without pulling the object to the local node.
+
+    Cancelling the awaiting task cancels the underlying C++ wait so the
+    Python callback can be released promptly.
+
+    Args:
+        ray_waitables: List of :class:`~ObjectRef` or
+            :class:`~ObjectRefGenerator` to wait on. Must be unique.
+        num_returns: Number of waitables that should become ready before
+            returning.
+        timeout: Max seconds to wait, or ``None`` to wait indefinitely.
+            ``timeout=0`` reports waitables already present in the in-process
+            memory store and does not start plasma pulls (unlike
+            :func:`ray.wait`, which may begin fetches on a zero timeout).
+        fetch_local: If True, wait until objects are present on the local
+            node. If False, return as soon as each object exists anywhere in
+            the cluster (no local pull).
+
+    Returns:
+        A pair ``(ready, remaining)`` of waitable lists, preserving input
+        order within each list.
+    """
+    import asyncio
+
+    worker = global_worker
+    worker.check_connected()
+
+    if isinstance(ray_waitables, ObjectRef) or isinstance(
+        ray_waitables, ObjectRefGenerator
+    ):
+        raise TypeError(
+            "_wait_async() expected a list of ray.ObjectRef or "
+            "ray.ObjectRefGenerator, got a single ray.ObjectRef or "
+            f"ray.ObjectRefGenerator {ray_waitables}"
+        )
+
+    if not isinstance(ray_waitables, list):
+        raise TypeError(
+            "_wait_async() expected a list of ray.ObjectRef or "
+            "ray.ObjectRefGenerator, "
+            f"got {type(ray_waitables)}"
+        )
+
+    if timeout is not None and timeout < 0:
+        raise ValueError(
+            "The 'timeout' argument must be nonnegative. " f"Received {timeout}"
+        )
+
+    for ray_waitable in ray_waitables:
+        if not isinstance(ray_waitable, ObjectRef) and not isinstance(
+            ray_waitable, ObjectRefGenerator
+        ):
+            raise TypeError(
+                "_wait_async() expected a list of ray.ObjectRef or "
+                "ray.ObjectRefGenerator, "
+                f"got list containing {type(ray_waitable)}"
+            )
+
+    if len(ray_waitables) == 0:
+        return [], []
+
+    if len(ray_waitables) != len(set(ray_waitables)):
+        raise ValueError("_wait_async requires a list of unique ray_waitables.")
+    if num_returns <= 0:
+        raise ValueError("Invalid number of objects to return %d." % num_returns)
+    if num_returns > len(ray_waitables):
+        raise ValueError(
+            "num_returns cannot be greater than the number "
+            "of ray_waitables provided to _wait_async."
+        )
+
+    # timeout=None -> wait forever (-1). Explicit timeout uses milliseconds.
+    timeout_milliseconds = -1 if timeout is None else int(timeout * 1000)
+
+    loop = asyncio.get_running_loop()
+    fut: asyncio.Future = loop.create_future()
+
+    def _on_complete(exc, ready_bits):
+        if fut.done():
+            return
+
+        def _set():
+            if fut.done():
+                return
+            if exc is not None:
+                fut.set_exception(exc)
+                return
+            ready_ids = []
+            remaining_ids = []
+            for i, ray_waitable in enumerate(ray_waitables):
+                if ready_bits[i]:
+                    ready_ids.append(ray_waitable)
+                else:
+                    remaining_ids.append(ray_waitable)
+            fut.set_result((ready_ids, remaining_ids))
+
+        loop.call_soon_threadsafe(_set)
+
+    core_worker = get_core_worker()
+    handle = core_worker.wait_async(
+        ray_waitables,
+        num_returns,
+        timeout_milliseconds,
+        fetch_local,
+        _on_complete,
+    )
+
+    def _cancel_cpp_wait(f: asyncio.Future) -> None:
+        if f.cancelled() and handle != 0:
+            core_worker.cancel_wait_async(handle)
+
+    fut.add_done_callback(_cancel_cpp_wait)
+    return await fut
+
+
 @client_mode_hook
 def _wait_generators_bulk(
     ray_generators: List[Tuple[ObjectRefGenerator, List[bool]]],

python/ray/_raylet.pyx

@@ -44,6 +44,7 @@ from libc.stdint cimport (
     uint64_t,
     uint8_t,
 )
+from libc.stddef cimport size_t
 from libcpp cimport bool as c_bool, nullptr
 from libcpp.memory cimport (
     dynamic_pointer_cast,
@@ -3767,6 +3768,69 @@ cdef class CoreWorker:
 
         return ready, not_ready
 
+    def wait_async(self,
+                   object_refs_or_generators,
+                   int num_returns,
+                   int64_t timeout_ms,
+                   c_bool fetch_local,
+                   callback: Callable):
+        """Register an async wait that invokes ``callback`` once.
+
+        Args:
+            object_refs_or_generators: List of ObjectRef or ObjectRefGenerator
+                to wait on.
+            num_returns: Number of waitables that should become ready.
+            timeout_ms: Timeout in milliseconds; negative means wait forever.
+            fetch_local: Whether ready objects must be present locally.
+            callback: Called as ``callback(exc, ready_bits)``. ``ready_bits``
+                is a list[bool] parallel to ``object_refs_or_generators``.
+                On error, ``exc`` is set and ``ready_bits`` is None.
+
+        Returns:
+            A non-zero handle for ``cancel_wait_async``, or 0 if ``callback``
+            already ran synchronously (validation error / immediate completion).
+        """
+        cdef:
+            c_vector[CObjectID] wait_ids
+            uint64_t handle = 0
+
+        for ref_or_generator in object_refs_or_generators:
+            if isinstance(ref_or_generator, ObjectRef):
+                wait_ids.push_back((<ObjectRef>ref_or_generator).native())
+            elif isinstance(ref_or_generator, ObjectRefGenerator):
+                wait_ids.push_back(
+                    CObjectID.FromBinary(
+                        ref_or_generator._get_next_object_id_binary()))
+            else:
+                raise TypeError(
+                    "wait_async() expected a list of ray.ObjectRef "
+                    "or ObjectRefGenerator, "
+                    f"got list containing {type(ref_or_generator)}"
+                )
+
+        # Keep the callback alive until the C++ side invokes it (or cancel).
+        # WaitAsync invokes the callback exactly once on every path (including
+        # sync validation / immediate completion); the callback's finally
+        # DECREFs. Do not DECREF again if WaitAsync returns or throws after
+        # that invocation — that would underflow the refcount.
+        cpython.Py_INCREF(callback)
+        with nogil:
+            handle = CCoreWorkerProcess.GetCoreWorker().WaitAsync(
+                wait_ids,
+                num_returns,
+                timeout_ms,
+                fetch_local,
+                wait_async_callback_impl,
+                <void*>callback,
+            )
+        # handle == 0 means the callback already ran synchronously.
+        return handle
+
+    def cancel_wait_async(self, uint64_t handle):
+        """Cancel an in-flight ``wait_async`` identified by ``handle``."""
+        with nogil:
+            CCoreWorkerProcess.GetCoreWorker().CancelWaitAsync(handle)
+
     def get_local_ongoing_lineage_reconstruction_tasks(self):
         cdef:
             unordered_map[CLineageReconstructionTask, uint64_t] tasks
@@ -5435,6 +5499,41 @@ cdef void async_callback(shared_ptr[CRayObject] obj,
         cpython.Py_DECREF(user_callback)
 
 
+cdef void wait_async_callback_impl(CRayStatus status,
+                                   const uint8_t *ready,
+                                   size_t n,
+                                   void *user_callback_ptr) with gil:
+    user_callback = <object>user_callback_ptr
+    try:
+        if not status.ok():
+            # Mirror check_status exception types without calling check_status
+            # (nogil raises across this C++ callback boundary are unsafe).
+            # Notably Status::Invalid falls through to RaySystemError there,
+            # not ValueError — keep the same here.
+            message = status.message().decode()
+            if (status.IsInvalidArgument() or status.IsNotFound() or
+                    status.IsObjectNotFound() or status.IsObjectUnknownOwner()):
+                exc = ValueError(message)
+            elif status.IsTimedOut():
+                exc = GetTimeoutError(message)
+            elif status.IsObjectRefEndOfStream():
+                exc = ObjectRefStreamEndOfStreamError(message)
+            elif status.IsIOError():
+                exc = IOError(message)
+            else:
+                # Includes Status::Invalid (duplicate ids, bad num_objects).
+                exc = RaySystemError(message)
+            user_callback(exc, None)
+            return
+        ready_bits = [ready[i] != 0 for i in range(n)]
+        user_callback(None, ready_bits)
+    except BaseException:
+        # Must not skip DECREF or leave the awaiting future hung.
+        logger.exception("failed to run wait_async callback (user func)")
+    finally:
+        cpython.Py_DECREF(user_callback)
+
+
 # Note this deletes keys with prefix `RAY{key_prefix}@`
 # Example: with key_prefix = `default`, we remove all `RAYdefault@...` keys.
 def del_key_prefix_from_storage(host, port, username, password, use_ssl, key_prefix):

python/ray/includes/libcoreworker.pxd

@@ -2,7 +2,8 @@
 # distutils: language = c++
 # cython: embedsignature = True
 
-from libc.stdint cimport int64_t, uint64_t
+from libc.stdint cimport int64_t, uint64_t, uint8_t
+from libc.stddef cimport size_t
 from libcpp cimport bool as c_bool
 from libcpp.functional cimport function
 from libcpp.memory cimport shared_ptr, unique_ptr
@@ -66,6 +67,9 @@ ctypedef void (*ray_callback_function) \
 ctypedef void (*plasma_callback_function) \
     (CObjectID object_id, int64_t data_size, int64_t metadata_size)
 
+ctypedef void (*wait_async_callback) \
+    (CRayStatus status, const uint8_t *ready, size_t n, void *user_data)
+
 # NOTE: This ctypedef is needed, because Cython doesn't compile
 # "pair[shared_ptr[const CActorHandle], CRayStatus]".
 # This is a bug of cython: https://github.com/cython/cython/issues/3967.
@@ -330,6 +334,10 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         CRayStatus Wait(const c_vector[CObjectID] &object_ids, int num_objects,
                         int64_t timeout_ms, c_vector[c_bool] *results,
                         c_bool fetch_local)
+        uint64_t WaitAsync(const c_vector[CObjectID] &object_ids, int num_objects,
+                           int64_t timeout_ms, c_bool fetch_local,
+                           wait_async_callback callback, void *user)
+        void CancelWaitAsync(uint64_t handle)
         CRayStatus GetLocalObjectLocations(
                 const c_vector[CObjectID] &object_ids,
                 c_vector[optional[CObjectLocation]] *results)