[Core][1/3] Make streaming generator public API #38784

rkooo567 · 2023-08-23T17:46:24Z

Why are these changes needed?

This PR is doing
- [x] Introduce throw/send
- [x] Bring up RayWaitable type
- [x] Modify types for cancel, remote, ray.get, and ray.wait
- [x] Update doctoring for ray.wait & ray.cancel
- [x] Implement completed, is_next_ready, is_finished
- [x] Allow to not specify num_returns=“streaming”

There will be total 3 PRs;

All sort of implementation work <---- current
Rename StreamingObjectRefGenerator and annotate it as a public API (and move it to __init__.py)
Typing support

Related issue number

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

python/ray/_private/ray_option_utils.py

rkooo567 · 2023-08-23T17:47:47Z

python/ray/_private/worker.py

@@ -121,13 +121,17 @@
 DAGNode = TypeVar("DAGNode")


+class RayWaitable(Generic[R]):


This could be contraversial.

I found a bit weird to make ObjectRef and ObjectRefGenerator a subclass of RayWaitable because

they don't share any functionality

they don't share any interface.

So I found it a bit more natural to just define the Union types. Any thoughts?

using a union makes sense. but I think using an empty interface won't hurt either. whichever way should be fine.

2 related questions:

Should StreamingObjectRefGenerator also be a generic type?

is there any difference between defining RayWaitable in this way versus just RayWaitable[R] = Union["ObjectRef[R]", StreamingObjectRefGenerator]?

so in this case, it is the type of returned object ref?

StreamingObjectRefGenerator[T] == __next__(self) -> ObjectRef[T]?

I think this sgtm.

So I tried defining RayWaitable = Union["ObjectRef[R]", StreamingObjectRefGenerator], but it raised an exception, so I tried the current approach lol..

1 is addressed

Hmm actually I found out Cython doesn't support Generic... We can't add a type here unless we move this out of cython code (which I don't plan to)

So we cannot just do type alias RayWaitable = Union[]

class RayWaitable and Union["ObjectRef[R]", StreamingObjectRefGenerator] are fundamentally different? looks Union["ObjectRef[R]", StreamingObjectRefGenerator] are more preferrable

It looks like we can't if we'd like to support RayWaitable[R] -> this syntax. I originally tried RayWaitable = Union[...], but it failed when I declared RayWaitable[R]

Confirmed it doesn't work. We cannot make RayWaittable a Generic type without declaring a class that inherits from Generic

rkooo567 · 2023-08-23T17:48:08Z

python/ray/_private/worker.py

        if isinstance(object_refs, StreamingObjectRefGenerator):
+            # TODO(sang): We should raise an exception.


Raising an exception requires lots of test changes. I will fix it in the third PR.

raulchen · 2023-08-23T21:11:41Z

python/ray/_private/worker.py

@@ -121,13 +121,17 @@
 DAGNode = TypeVar("DAGNode")


+class RayWaitable(Generic[R]):


using a union makes sense. but I think using an empty interface won't hurt either. whichever way should be fine.

2 related questions:

Should StreamingObjectRefGenerator also be a generic type?

is there any difference between defining RayWaitable in this way versus just RayWaitable[R] = Union["ObjectRef[R]", StreamingObjectRefGenerator]?

python/ray/_private/worker.py

raulchen · 2023-08-23T21:25:24Z

python/ray/_raylet.pyx

@@ -262,23 +262,31 @@ class ObjectRefGenerator:


 class StreamingObjectRefGenerator:


not blocking this PR. Should we move this class to a different file? this file is pretty long. also, is there a perf improvement to define it in Cython? It seems to me that most code in this class is pure Python code.

I think we can definitely move out of it. I will handle it in the second PR together

rkooo567 · 2023-08-24T10:08:30Z

Btw, until @pcmoritz gives a green light, I will mark the feature as alpha. I will still proceed merging PRs needed for the proposal (and try to get it in by 2.7), so please review it! https://docs.google.com/document/d/1LzKPL7wxz5tUJJLMMZ7BlfdDlYclAzMdzolTezU06ck/edit#heading=h.ehfmi7oa003b

jjyao · 2023-08-24T16:40:22Z

python/ray/_private/ray_option_utils.py

        else "The keyword 'num_returns' only accepts None, a non-negative integer, or "
        '"dynamic" (for generators)',


This needs to be updated to mention streaming.

I think we want to hide "streaming". It will work without streaming if it is a generator. I will modify the comment when I hard deprecate the dynamic generator.

python/ray/_private/ray_option_utils.py

jjyao · 2023-08-24T16:55:29Z

python/ray/_private/worker.py

@@ -121,13 +121,17 @@
 DAGNode = TypeVar("DAGNode")


+class RayWaitable(Generic[R]):


So we cannot just do type alias RayWaitable = Union[]

jjyao · 2023-08-24T17:03:07Z

python/ray/_private/worker.py

-        __arg0: "Union[T0, ObjectRef[T0]]",
-        __arg1: "Union[T1, ObjectRef[T1]]",
-    ) -> "ObjectRef[R]":
+        __arg0: "Union[T0, RayWaitable[T0]]",


I think for the inputs, it should still be T0 or ObjectRef[T0]?

I originally did that. It doesn't work with mypi because remote methods return RayWaitable, and if we pass that into other method, it cannot be automatically cast to ObjectRef without runtime check (mypi check is done statically)

Yep updated. I overloaded the remote method so that it can return both generator and object ref. After that it is fixed

python/ray/_private/worker.py

jjyao · 2023-08-24T17:11:16Z

python/ray/_private/worker.py

@@ -2837,7 +2863,9 @@ def kill(actor: "ray.actor.ActorHandle", *, no_restart: bool = True):

 @PublicAPI
 @client_mode_hook
-def cancel(object_ref: "ray.ObjectRef", *, force: bool = False, recursive: bool = True):
+def cancel(
+    ray_waitable: RayWaitable[R], *, force: bool = False, recursive: bool = True


I guess it's both waitable and cancelable lol

jjyao · 2023-08-24T17:12:26Z

python/ray/_raylet.pyx

@@ -264,23 +264,31 @@ class ObjectRefGenerator:


 class StreamingObjectRefGenerator:


Annotate with @publicapi

Also do we want to start with alpha or beta or stable

This is the scope of the second PR!

We will scope it to beta

rkooo567 · 2023-08-29T07:26:07Z

We decided to make the API public in Ray 2.8 instead (we will still update the documentation for 2.7 as an alpha feature)

stale · 2023-10-15T08:10:20Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

This is a PR to write a streaming generator documentation. Unlike existing generator doc, streaming generator doc is under "advanced" usage, and I will add proper links from each task/actor/pattern pages to this doc. The doc is written with the assumption is is an "alpha" feature. It is important to merge this doc PR before proceeding with #38784 as it will need to refer docs here and there. Some of the contents (such as "API is subject to change") will be removed after we merge these PRs.

rkooo567 · 2023-11-20T14:07:28Z

@jjyao @scv119

Let's get back to merging this PR.

Re; typing.

I remove RayWaitable and used Union instead.
To make generator -> ObjectRefGenerator[int], it requires lots of repetitive code. I will do it in a follow up. I got it working in this commit, 28d5355. But I won't inclde this in this PR

In this PR, I just made sure it doesn't raise an error when using mypy

rkooo567 · 2023-11-20T14:08:51Z

python/ray/_private/worker.py

-        __arg0: "Union[T0, ObjectRef[T0]]",
-        __arg1: "Union[T1, ObjectRef[T1]]",
-    ) -> "ObjectRef[R]":
+        __arg0: "Union[T0, ObjectRefType[T0]]",


This is a bug I found. When I passed object ref, mypy couldn't find the correct type, but it found ray._raylet.ObjectRef

cc @pcmoritz

rkooo567 · 2023-11-20T14:10:03Z

python/ray/types.py

+
+
+@PublicAPI(stability="beta")
+class StreamingObjectRefGenerator(Generic[T]):


Note: Currently, Cython cannot define generic inherited classes. WE need to upgrade it to support it

jjyao · 2023-11-20T17:41:33Z

python/ray/_private/ray_option_utils.py

-        default_value=1,
+        else "Default None. The keyword 'num_returns' only accepts None, "
+        "a non-negative integer, or "
+        '"dynamic" (for generators). "dynamic" is deprecated from Ray 2.7.',


Suggested change

'"dynamic" (for generators). "dynamic" is deprecated from Ray 2.7.',

'"streaming" (for generators). "dynamic" is deprecated from Ray 2.7.',

?

why the default value is not 1 but None, what does None mean here?

Let me update the comment. None basically means default.

Default behavior is 1 for regular task, and streaming for a generator task

updated the comment. Please take a look and lmk what you think

jjyao · 2023-11-20T17:42:17Z

python/ray/_private/worker.py

@@ -59,6 +59,8 @@
    StreamingObjectRefGenerator,
    raise_sys_exit_with_custom_error_message,
 )
+from ray.types import ObjectRef as ObjectRefType


why this alias?

Because we also import ObjectRef. We currently have 2 object ref, one real one (ObjectRef in Cython), and the other just for typing (types.ObjectRef). I think we should fix this, but it is not in a scope of this PR.

When I don't import this type, I found mypy looks for the ObjectRef in Cython, and it fails a pretty basic test when I annotate the object ref type to test_typing_good.

jjyao · 2023-11-20T17:44:19Z

python/ray/_private/worker.py

@@ -2682,50 +2685,65 @@ def put(
 @PublicAPI
 @client_mode_hook
 def wait(
-    object_refs: List["ray.ObjectRef"],
+    ray_waitables: Union["ObjectRefType[R]", "StreamingObjectRefGeneratorType[R]"],


Does it make sense to make StreaingObjectRefGenerator generic since it can return multiple objects with different types?

Yeah you can do StreaingObjectRefGenerator[Union[...]]

in that case

I think this is important decision though. Normally generator has a type annotation Generator[yield type, send type, return type], so I think it makes sense to have it in generator too. But I created a sync meeting with @pcmoritz tomorrow to discuss a bit more details.

jjyao · 2023-11-20T17:46:20Z

python/ray/_raylet.pyx

+    def send(self, value):
+        raise NotImplementedError("`gen.send` is not supported.")
+
+    def throw(self, value):
+        raise NotImplementedError("`gen.throw` is not supported.")
+
+    def close(self):
+        raise NotImplementedError("`gen.close` is not supported.")


why do we need these?

These are standard generator APIs that are not supported in Ray

Also specified in https://docs.google.com/document/d/1LzKPL7wxz5tUJJLMMZ7BlfdDlYclAzMdzolTezU06ck/edit#heading=h.ehfmi7oa003b

jjyao · 2023-11-22T00:40:16Z

python/ray/tests/typing_files/check_typing_good.py

+
+
+"""
+TODO(sang): Enable it.


Why we need to disable it for now?

It is not working now

edoakes · 2023-11-23T06:03:39Z

python/ray/_raylet.pyx

@@ -4101,6 +4101,7 @@ cdef class CoreWorker:
                                         worker.current_session_and_job)
        else:
            return ray.actor.ActorHandle(language, actor_id,
+                                         {},  # method is_generator


better to make these kwargs

btw, are you talking about

def __init__(self, method is_generator=None ):

something like this?

yes, use kwargs instead of a long list of optional positionals with comments

edoakes · 2023-11-23T06:03:40Z

python/ray/_raylet.pyx

@@ -4101,6 +4101,7 @@ cdef class CoreWorker:
                                         worker.current_session_and_job)
        else:
            return ray.actor.ActorHandle(language, actor_id,
+                                         {},  # method is_generator


better to make these kwargs

This PR is doing - [x] Introduce throw/send - [x] Bring up RayWaitable type - [x] Modify types for cancel, remote, ray.get, and ray.wait - [x] Update doctoring for ray.wait & ray.cancel - [x] Implement completed, is_next_ready, is_finished - [x] Allow to not specify num_returns=“streaming” There will be total 3 PRs; All sort of implementation work <---- current Rename StreamingObjectRefGenerator and annotate it as a public API (and move it to __init__.py) Typing support

This is the second PR of 3 steps. All sort of implementation work [done] Rename StreamingObjectRefGenerator and annotate it as a public API (and move it to init.py) <---- this PR Typing support [future] This PR renames StreamingObjectRefGenerator - > ObjectRefGenerator. This also means the existing dynamic generator's (num_returns="dynamic") return type will become ObjectRefGenerator -> DynamicObjectRefGenerator, which is a breaking change, but we will anyway deprecate this API. I am open for a different name if we'd like to keep the perfect backward compatibility.

rkooo567 added 2 commits August 24, 2023 02:44

Done

239a63a

clean up.

4492b41

rkooo567 commented Aug 23, 2023

View reviewed changes

rkooo567 assigned ericl, scv119, jjyao, raulchen, edoakes and pcmoritz Aug 23, 2023

raulchen reviewed Aug 23, 2023

View reviewed changes

Merge branch 'master' into streaming-generator-public-api

bf36e5f

Addressed code review.

6d2a0bf

jjyao reviewed Aug 24, 2023

View reviewed changes

rkooo567 added 2 commits August 25, 2023 08:04

Merge branch 'master' into streaming-generator-public-api

4379f9f

Addressed code review.

2086104

rkooo567 mentioned this pull request Oct 11, 2023

[Doc] Streaming generator alpha doc #39914

Merged

8 tasks

stale bot added the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Oct 15, 2023

rkooo567 removed the stale The issue is stale. It will be closed within 7 days unless there are further conversation label Oct 16, 2023

SangBin Cho added 2 commits November 20, 2023 23:00

ip

1a7360d

remove typing

28d5355

rkooo567 commented Nov 20, 2023

View reviewed changes

streaming-generator-public-api

4400194

jjyao reviewed Nov 20, 2023

View reviewed changes

SangBin Cho added 2 commits November 21, 2023 13:33

Merge branch 'master' into streaming-generator-public-api

8f35359

Fix issues.

a9a78b7

Fixed.

bf26864

jjyao approved these changes Nov 22, 2023

View reviewed changes

SangBin Cho added 3 commits November 22, 2023 20:09

.

9211dc4

Remove typing changes

e05281b

done

0cfdfd8

edoakes reviewed Nov 23, 2023

View reviewed changes

SangBin Cho added 2 commits November 27, 2023 16:45

Merge branch 'master' into streaming-generator-public-api

7829e42

done

ebd8238

rkooo567 merged commit 7b33cd8 into ray-project:master Nov 27, 2023
16 of 17 checks passed

		@@ -121,13 +121,17 @@
		DAGNode = TypeVar("DAGNode")


		class RayWaitable(Generic[R]):

		if isinstance(object_refs, StreamingObjectRefGenerator):
		# TODO(sang): We should raise an exception.

		@@ -262,23 +262,31 @@ class ObjectRefGenerator:


		class StreamingObjectRefGenerator:

		else "The keyword 'num_returns' only accepts None, a non-negative integer, or "
		'"dynamic" (for generators)',

		@@ -264,23 +264,31 @@ class ObjectRefGenerator:


		class StreamingObjectRefGenerator:



		@PublicAPI(stability="beta")
		class StreamingObjectRefGenerator(Generic[T]):

	'"dynamic" (for generators). "dynamic" is deprecated from Ray 2.7.',
	'"streaming" (for generators). "dynamic" is deprecated from Ray 2.7.',

[Core][1/3] Make streaming generator public API #38784

[Core][1/3] Make streaming generator public API #38784

Conversation

rkooo567 commented Aug 23, 2023 • edited Loading

Why are these changes needed?

Related issue number

Checks

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 commented Aug 24, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 Aug 24, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 commented Aug 29, 2023

stale bot commented Oct 15, 2023

rkooo567 commented Nov 20, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 commented Aug 23, 2023 •

edited

Loading

rkooo567 Aug 24, 2023 •

edited

Loading

rkooo567 commented Nov 20, 2023 •

edited

Loading