chore: import upstream snapshot with attribution

python/ray/data/_internal/execution/node_trackers/__init__.py

@@ -0,0 +1,3 @@
+from .actor_location import ActorLocationTracker, get_or_create_actor_location_tracker
+
+__all__ = ["get_or_create_actor_location_tracker", "ActorLocationTracker"]

python/ray/data/_internal/execution/node_trackers/actor_location.py

@@ -0,0 +1,39 @@
+import threading
+from typing import List
+
+import ray
+
+
+@ray.remote(num_cpus=0, max_restarts=-1, max_task_retries=-1)
+class ActorLocationTracker:
+    def __init__(self):
+        self._actor_locations = {}
+        self._actor_locations_lock = threading.Lock()
+
+    def update_actor_location(self, logical_actor_id: str, node_id: str):
+        with self._actor_locations_lock:
+            self._actor_locations[logical_actor_id] = node_id
+
+    def get_actor_locations(self, logical_actor_ids: List[str]):
+        return {
+            logical_actor_id: self._actor_locations.get(logical_actor_id, None)
+            for logical_actor_id in logical_actor_ids
+        }
+
+
+def get_or_create_actor_location_tracker():
+
+    # Pin the actor location tracker to the local node so it fate-shares with the driver.
+    # NOTE: for Ray Client, the ray.get_runtime_context().get_node_id() should
+    # point to the head node.
+    label_selector = {
+        ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
+    }
+    return ActorLocationTracker.options(
+        name="ActorLocationTracker",
+        namespace="ActorLocationTracker",
+        get_if_exists=True,
+        lifetime="detached",
+        label_selector=label_selector,
+        max_concurrency=8,
+    ).remote()