Objectref typestubbing

[minerharry]

Why are these changes needed?

Adds proper type-stubbing for ObjectRef (and its cousins, ActorID, TaskID, etc). See #53591. The ray/includes .pxi files (and the _raylet.pxi) are inaccessible to type-checkers, so core ray objects like ObjectRef don't have proper code-completion. While the issue extends to many ray objects defined in includes/* and _raylet.pxi, notably ObjectRefGenerator and similar, this PR focuses just on unique_ids.pxi and object_ref.pxi. By adding python stub files, .pyi, that correspond to the .pxi files, this PR allows type-checkers to understand the contents of the cython files. Each .pyi file contains the signature of every class and method in the .pxi files to expose (though I went ahead and added every method for completeness).

Also, I added a __class_getitem__ method to ObjectRef in object_ref.pxi which mirrors how typing.Generic works, as multiple inheritance (which would be required for generic subclassing of BaseID) is not well supported in cython. This allows for runtime type-subscripting (e.g. ObjectRef[int]) without error and returns a proper annotated type.

Related issue number

Limited-scope version of #55066, just focusing on unique_ids.pxi and object_ref.pxi. Partially satisfies #53591.

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 :(

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[github-actions]

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

You can always ask for help on our discussion forum or Ray's public slack channel.

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

[minerharry]

@can-anyscale can you review?

[jjyao]

@minerharry Thanks for doing this! Could you update the PR description with more details including the issue and solution.

[minerharry]

Sure, done

[can-anyscale]

hi yes, not stale, still in my list of doing, sorry about this

[pcmoritz]

This doesn't seem to be used any more, I made a PR to remove it #56184, can you also remove it from this PR?

[pcmoritz]

This PR looks overall good to me, I have a few formatting requests before we can merge it:

• In all of the # comments, can you add a space after the #, to be consistent with the rest of the codebase?
• In all of the typing annotations like call_site_data:str, can you add a space after the : like call_site_data: str to be consistent with the rest of the codebase?
• In all of the arrow expressions like __dealloc__(self)->None, can you add a space before and after the arrow like __dealloc__(self) -> None?

[minerharry]

thanks @pcmoritz! Great points about formatting, and thanks for updating them. I'll go ahead and do the same to the other PR as well. (what's the last step for merging this? new here)

[pcmoritz]

Sounds great, thanks! I'll just do a few final checks on the typing and then will merge the PR, aiming to get it done later today :)

[pcmoritz]

I performed some tests now:

import ray 

@ray.remote
def f() -> int:
    return 1

a = f.remote()

reveal_type(ray.get(a))

reveal_type(a)

reveal_type(a.redis_shard_hash())

reveal_type(a.task_id())

async def g():
    await a

print("path", ray.__path__, ray.__version__)

Running mypy <file> both with and without the PR:

Without:

t.py:9: note: Revealed type is "builtins.int"
t.py:11: note: Revealed type is "ray._raylet.ObjectRef[builtins.int]"
t.py:13: error: "ObjectRef[int]" has no attribute "redis_shard_hash"  [attr-defined]
t.py:13: note: Revealed type is "Any"
t.py:15: error: "ObjectRef[int]" has no attribute "task_id"  [attr-defined]
t.py:15: note: Revealed type is "Any"
Found 2 errors in 1 file (checked 1 source file)

With:

(base) pcmoritz@pcmoritz-DQ44HV60WX ~ % mypy t.py
t.py:9: note: Revealed type is "builtins.int"
t.py:11: note: Revealed type is "ray.includes.object_ref.ObjectRef[builtins.int]"
t.py:13: note: Revealed type is "builtins.int"
t.py:15: note: Revealed type is "ray.includes.unique_ids.TaskID"
Success: no issues found in 1 source file