chore: import upstream snapshot with attribution

python/ray/util/client/runtime_context.py

@@ -0,0 +1,77 @@
+from types import SimpleNamespace
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ray import JobID, NodeID, WorkerID
+    from ray.runtime_context import RuntimeContext
+
+
+class _ClientWorkerPropertyAPI:
+    """Emulates the properties of the ray._private.worker object for the client"""
+
+    def __init__(self, worker):
+        assert worker is not None
+        self.worker = worker
+
+    def build_runtime_context(self) -> "RuntimeContext":
+        """Creates a RuntimeContext backed by the properites of this API"""
+        # Defer the import of RuntimeContext until needed to avoid cycles
+        from ray.runtime_context import RuntimeContext
+
+        return RuntimeContext(self)
+
+    def _fetch_runtime_context(self):
+        import ray.core.generated.ray_client_pb2 as ray_client_pb2
+
+        return self.worker.get_cluster_info(
+            ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT
+        )
+
+    @property
+    def mode(self):
+        from ray._private.worker import SCRIPT_MODE
+
+        return SCRIPT_MODE
+
+    @property
+    def current_job_id(self) -> "JobID":
+        from ray import JobID
+
+        return JobID(self._fetch_runtime_context().job_id)
+
+    @property
+    def current_node_id(self) -> "NodeID":
+        from ray import NodeID
+
+        return NodeID(self._fetch_runtime_context().node_id)
+
+    @property
+    def worker_id(self) -> "WorkerID":
+        """Binary worker id for the Ray Client server process (cluster driver worker)."""
+        from ray import WorkerID
+
+        return WorkerID(self._fetch_runtime_context().worker_id)
+
+    @property
+    def namespace(self) -> str:
+        return self._fetch_runtime_context().namespace
+
+    @property
+    def should_capture_child_tasks_in_placement_group(self) -> bool:
+        return self._fetch_runtime_context().capture_client_tasks
+
+    @property
+    def runtime_env(self) -> str:
+        return self._fetch_runtime_context().runtime_env
+
+    def check_connected(self) -> bool:
+        return self.worker.ping_server()
+
+    @property
+    def gcs_client(self) -> str:
+        return SimpleNamespace(address=self._fetch_runtime_context().gcs_address)
+
+    @property
+    def node(self):
+        """Emulates the worker.node property for client mode"""
+        return SimpleNamespace(session_name=self._fetch_runtime_context().session_name)