[core][state] make get_task() accept ObjectRef (#43507)

Signed-off-by: max-509 <123456vershinin@gmail.com>
Signed-off-by: hongchaodeng <hongchaodeng1@gmail.com>
Co-authored-by: max-509 <123456vershinin@gmail.com>

python/ray/tests/test_state_api.py

@@ -2340,9 +2340,14 @@ def g(dep):
     def impossible():
         pass
 
-    out = [f.options(name=f"f_{i}").remote() for i in range(2)]  # noqa
-    g_out = g.remote(f.remote())  # noqa
-    im = impossible.remote()  # noqa
+    f_refs = [f.options(name=f"f_{i}").remote() for i in range(2)]  # noqa
+    g_ref = g.remote(f.remote())  # noqa
+    im_ref = impossible.remote()  # noqa
+
+    def verify_task_from_objectref(task, job_id, tasks):
+        assert task["job_id"] == job_id
+        assert task["actor_id"] is None
+        assert any(task["task_id"] == t["task_id"] for t in tasks)
 
     def verify():
         tasks = list_tasks()
@@ -2352,6 +2357,12 @@ def verify():
         for task in tasks:
             assert task["actor_id"] is None
 
+        # Test get_task by objectRef
+        for ref in f_refs:
+            verify_task_from_objectref(get_task(ref), job_id, tasks)
+        verify_task_from_objectref(get_task(g_ref), job_id, tasks)
+        verify_task_from_objectref(get_task(im_ref), job_id, tasks)
+
         waiting_for_execution = len(
             list(
                 filter(

python/ray/util/state/api.py

@@ -8,6 +8,7 @@
 
 import requests
 
+import ray
 from ray.dashboard.modules.dashboard_sdk import SubmissionClient
 from ray.dashboard.utils import (
     get_address_for_submission_client,
@@ -710,15 +711,15 @@ def get_worker(
 
 @DeveloperAPI
 def get_task(
-    id: str,
+    id: Union[str, "ray.ObjectRef"],
     address: Optional[str] = None,
     timeout: int = DEFAULT_RPC_TIMEOUT,
     _explain: bool = False,
 ) -> Optional[TaskState]:
     """Get task attempts of a task by id.
 
     Args:
-        id: Id of the task
+        id: String id of the task or ObjectRef that corresponds to task
         address: Ray bootstrap address, could be `auto`, `localhost:6379`.
             If None, it will be resolved automatically from an initialized ray.
         timeout: Max timeout value for the state APIs requests made.
@@ -734,9 +735,14 @@ def get_task(
         Exceptions: :class:`RayStateApiException <ray.util.state.exception.RayStateApiException>` if the CLI
             failed to query the data.
     """  # noqa: E501
+    str_id: str
+    if isinstance(id, str):
+        str_id = id
+    else:
+        str_id = id.task_id().hex()
     return StateApiClient(address=address).get(
         StateResource.TASKS,
-        id,
+        str_id,
         GetApiOptions(timeout=timeout),
         _explain=_explain,
     )