diff --git a/tests/experimental/distributed/runtime/contexts/k8s_context_test.py b/tests/experimental/distributed/runtime/contexts/k8s_context_test.py index b677b2678..dd8eedd8c 100644 --- a/tests/experimental/distributed/runtime/contexts/k8s_context_test.py +++ b/tests/experimental/distributed/runtime/contexts/k8s_context_test.py @@ -50,7 +50,7 @@ def test_resolve_self_hostname_missing_env(self): with self.assertRaises(ValueError): k8s_context.resolve_self_hostname() - def test_k8s_jax_context_pathways(self): + def test_jax_context_pathways(self): envs = { "JAX_PLATFORMS": "proxy", "JAX_BACKEND_TARGET": "0.0.0.0:8000", @@ -61,7 +61,7 @@ def test_k8s_jax_context_pathways(self): k8s_context.K8sJaxContext().initialize() mock_pw.initialize.assert_called_once() - def test_k8s_jax_context_mcjax(self): + def test_jax_context_mcjax(self): envs = {} mock_jax = mock.MagicMock() with mock.patch.dict(os.environ, envs, clear=True): @@ -69,7 +69,7 @@ def test_k8s_jax_context_mcjax(self): k8s_context.K8sJaxContext().initialize() mock_jax.distributed.initialize.assert_called_once() - def test_k8s_discovery_context_register(self): + def test_discovery_context_register(self): envs = { "JOBSET_NAME": "myjobset", "REPLICATED_JOB_NAME": "worker", @@ -95,7 +95,68 @@ def test_k8s_discovery_context_register(self): b"pod-meta", ) - def test_k8s_process_context(self): + @mock.patch( + "tunix.experimental.distributed.runtime.contexts.k8s_context.discovery.connect" + ) + @mock.patch( + "tunix.experimental.distributed.runtime.contexts.k8s_context.discovery.DiscoveryServer" + ) + def test_discovery_context_connect(self, mock_server_cls, mock_connect): + envs = { + "JOBSET_NAME": "myjobset", + "REPLICATED_JOB_NAME": "worker", + "JOB_INDEX": "0", + "POD_INDEX": "1", + } + port = 8888 + args = argparse.Namespace( + discovery_port=port, + discovery_addrs=f"door:{port}", + discovery_id="door", + ) + + mock_server = mock_server_cls.return_value + + with mock.patch.dict(os.environ, envs): + with k8s_context.K8sDiscoveryContext(args) as disc_ctx: + on_client_connected = lambda cid, h, p, m, rec: None + on_client_disconnected = lambda cid, h, p, r: None + + disc_ctx.on_connect( + on_client_connected=on_client_connected, + on_client_disconnected=on_client_disconnected, + ) + mock_server.on_connect.assert_called_once_with( + on_client_connected=on_client_connected, + on_client_disconnected=on_client_disconnected, + ) + mock_server.start.assert_called_once_with(port) + + on_connected = lambda epoch, rec: None + on_disconnected = lambda epoch, r: None + + client = disc_ctx.connect( + b"pod-meta", + client_id="door", + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + mock_connect.assert_called_once_with( + "door-proc-0-0.door:8888", + "myjobset-worker-0-1.myjobset", + port, + b"pod-meta", + client_id="door", + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + self.assertEqual(disc_ctx._client, mock_connect.return_value) + + mock_connect.return_value.stop.assert_called_once() + mock_server.stop.assert_called_once() + self.assertIsNone(disc_ctx._client) + + def test_process_context(self): args = argparse.Namespace( discovery_port=portpicker.pick_unused_port(), discovery_addrs="door:8888", diff --git a/tests/experimental/distributed/runtime/contexts/local_context_test.py b/tests/experimental/distributed/runtime/contexts/local_context_test.py index acabb775d..5aac2cf3c 100644 --- a/tests/experimental/distributed/runtime/contexts/local_context_test.py +++ b/tests/experimental/distributed/runtime/contexts/local_context_test.py @@ -32,9 +32,7 @@ def test_resolve_discovery_address(self): @mock.patch( "tunix.experimental.distributed.runtime.discovery.discovery.grpc.server" ) - def test_local_discovery_context_lifecycle_and_registration( - self, mock_grpc_server - ): + def test_discovery_context_register(self, mock_grpc_server): port = portpicker.pick_unused_port() args = argparse.Namespace( discovery_port=port, @@ -56,7 +54,61 @@ def test_local_discovery_context_lifecycle_and_registration( self.assertFalse(disc_ctx._server.is_started()) - def test_local_process_context(self): + @mock.patch( + "tunix.experimental.distributed.runtime.contexts.local_context.discovery.connect" + ) + @mock.patch( + "tunix.experimental.distributed.runtime.contexts.local_context.discovery.DiscoveryServer" + ) + def test_discovery_context_connect(self, mock_server_cls, mock_connect): + port = 8888 + args = argparse.Namespace( + discovery_port=port, + discovery_addrs=f"leader:{port}", + discovery_id="worker-0", + ) + + mock_server = mock_server_cls.return_value + + with local_context.LocalDiscoveryContext(args) as disc_ctx: + on_client_connected = lambda cid, h, p, m, rec: None + on_client_disconnected = lambda cid, h, p, r: None + + disc_ctx.on_connect( + on_client_connected=on_client_connected, + on_client_disconnected=on_client_disconnected, + ) + mock_server.on_connect.assert_called_once_with( + on_client_connected=on_client_connected, + on_client_disconnected=on_client_disconnected, + ) + mock_server.start.assert_called_once_with(port) + + on_connected = lambda epoch, rec: None + on_disconnected = lambda epoch, r: None + + client = disc_ctx.connect( + b"my-metadata", + client_id="worker-0", + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + mock_connect.assert_called_once_with( + "localhost:8888", + "localhost", + port, + b"my-metadata", + client_id="worker-0", + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + self.assertEqual(disc_ctx._client, mock_connect.return_value) + + mock_connect.return_value.stop.assert_called_once() + mock_server.stop.assert_called_once() + self.assertIsNone(disc_ctx._client) + + def test_process_context(self): args = argparse.Namespace( discovery_port=portpicker.pick_unused_port(), discovery_addrs="leader:9999", diff --git a/tests/experimental/distributed/runtime/discovery/discovery_test.py b/tests/experimental/distributed/runtime/discovery/discovery_test.py index 81efc9106..2aa788e14 100644 --- a/tests/experimental/distributed/runtime/discovery/discovery_test.py +++ b/tests/experimental/distributed/runtime/discovery/discovery_test.py @@ -25,19 +25,26 @@ class DiscoveryTest(absltest.TestCase): + def test_start_unconfigured_mode_raises(self): + server = discovery.DiscoveryServer() + with self.assertRaises(RuntimeError): + server.start(8888) + def test_start_with_zero_port_raises(self): server = discovery.DiscoveryServer() + server.on_register(lambda h, p, m: None) with self.assertRaises(ValueError): - server.start(0, lambda h, p, m: None) + server.start(0) @mock.patch.object(grpc, "server") def test_start_twice_raises(self, mock_grpc_server): server = discovery.DiscoveryServer() port = 8888 - server.start(port, lambda h, p, m: None) + server.on_register(lambda h, p, m: None) + server.start(port) try: with self.assertRaises(RuntimeError): - server.start(port, lambda h, p, m: None) + server.start(port) finally: server.stop() @@ -52,7 +59,8 @@ def callback(hostname, p, metadata): received["port"] = p received["metadata"] = metadata - server.start(port, callback) + server.on_register(callback) + server.start(port) try: self.assertTrue(server.is_started()) mock_grpc_server.return_value.add_insecure_port.assert_called_once_with( @@ -102,6 +110,149 @@ def test_register_non_retryable_error_raises(self, mock_channel): with self.assertRaises(RuntimeError): discovery.register("localhost:9999", "node-0", 1234, b"meta") + def test_connect_initial_connection(self): + port = portpicker.pick_unused_port() + server = discovery.DiscoveryServer(heartbeat_sec=1) + server_connected_events = [] + + server.on_connect( + on_client_connected=lambda cid, h, p, m, rec: server_connected_events.append( + (cid, h, p, m, rec) + ) + ) + server.start(port, heartbeat_sec=1) + + client_connected_events = [] + try: + client = discovery.connect( + f"localhost:{port}", + "node-0", + 1234, + b"meta-data", + client_id="node-0", + on_connected=lambda epoch, rec: client_connected_events.append( + (epoch, rec) + ), + ) + + self.assertEqual(len(client_connected_events), 1) + epoch, is_reconnect = client_connected_events[0] + self.assertFalse(is_reconnect) + + self.assertEqual(len(server_connected_events), 1) + cid, h, p, m, is_rec = server_connected_events[0] + self.assertEqual(cid, "node-0") + self.assertFalse(is_rec) + + client.stop() + finally: + server.stop() + + def test_connect_reconnect_on_server_restart(self): + port = portpicker.pick_unused_port() + server = discovery.DiscoveryServer(heartbeat_sec=1) + server_connected_events = [] + + server.on_connect( + on_client_connected=lambda cid, h, p, m, rec: server_connected_events.append( + (cid, h, p, m, rec) + ) + ) + server.start(port, heartbeat_sec=1) + + client_connected_events = [] + client_disconnected_events = [] + reconnected_event = threading.Event() + + def on_connected(epoch, rec): + client_connected_events.append((epoch, rec)) + if rec: + reconnected_event.set() + + try: + client = discovery.connect( + f"localhost:{port}", + "node-0", + 1234, + b"meta-data", + client_id="node-0", + on_connected=on_connected, + on_disconnected=lambda epoch, reason: client_disconnected_events.append( + (epoch, reason) + ), + ) + + initial_epoch = client_connected_events[0][0] + + # Simulate server restart by changing servicer epoch + server._servicer._server_epoch = "rebooted-epoch-1234" + + # Wait for heartbeat loop to detect epoch mismatch and reconnect + self.assertTrue(reconnected_event.wait(timeout=5.0)) + + self.assertGreaterEqual(len(client_disconnected_events), 1) + self.assertEqual(client_disconnected_events[0][0], initial_epoch) + self.assertEqual(client_disconnected_events[0][1], "epoch_mismatch") + + self.assertGreaterEqual(len(client_connected_events), 2) + new_epoch, is_reconnected = client_connected_events[1] + self.assertTrue(is_reconnected) + self.assertEqual(new_epoch, "rebooted-epoch-1234") + + self.assertGreaterEqual(len(server_connected_events), 2) + self.assertTrue(server_connected_events[1][4]) # is_reconnect=True + + client.stop() + finally: + server.stop() + + def test_server_lease_eviction_on_heartbeat_timeout(self): + port = portpicker.pick_unused_port() + server = discovery.DiscoveryServer(heartbeat_sec=1) + server_disconnected_events = [] + evicted_event = threading.Event() + + def on_disconnected(cid, h, p, reason): + server_disconnected_events.append((cid, h, p, reason)) + evicted_event.set() + + server.on_connect(on_client_disconnected=on_disconnected) + server.start(port, heartbeat_sec=1) + + try: + client = discovery.connect( + f"localhost:{port}", + "node-0", + 1234, + b"meta-data", + client_id="node-0", + ) + # Stop client heartbeat thread prematurely to simulate crashed/dead client + client._stop_event.set() + + # Wait for server eviction loop (threshold = 3 * heartbeat_sec) + self.assertTrue(evicted_event.wait(timeout=5.0)) + + self.assertGreaterEqual(len(server_disconnected_events), 1) + cid, h, p, reason = server_disconnected_events[0] + self.assertEqual(cid, "node-0") + self.assertEqual(reason, "heartbeat_timeout") + + client.stop() + finally: + server.stop() + + def test_mode_mutual_exclusion(self): + server = discovery.DiscoveryServer() + server.on_register(lambda h, p, m: None) + with self.assertRaises(RuntimeError): + server.on_connect(lambda cid, h, p, m, rec: None) + + server2 = discovery.DiscoveryServer() + server2.on_connect(lambda cid, h, p, m, rec: None) + with self.assertRaises(RuntimeError): + server2.on_register(lambda h, p, m: None) + if __name__ == "__main__": absltest.main() diff --git a/tunix/experimental/distributed/examples/README.md b/tunix/experimental/distributed/examples/README.md index a6e40971f..2e75b9ea1 100644 --- a/tunix/experimental/distributed/examples/README.md +++ b/tunix/experimental/distributed/examples/README.md @@ -12,9 +12,10 @@ The Tunix distributed process runtime provides a **platform-agnostic execution f - [Example 2: Process with CLI Flags](#example-2-process-with-cli-flags) - [Example 3: Process with TPUs](#example-3-process-with-tpus) - [Example 4: Peer Discovery and Inter-Process Communication](#example-4-peer-discovery-and-inter-process-communication) -- [Example 5: Simulated Distributed RL Workload (Local)](#example-5-simulated-distributed-rl-workload-local) -- [Example 6: Simulated Distributed RL Workload on Kubernetes](#example-6-simulated-distributed-rl-workload-on-kubernetes) -- [Example 7: Distributed RL Generation with vLLM Workers](#example-7-distributed-rl-generation-with-vllm-workers) +- [Example 5: Persistent Reconnecting Discovery with Heartbeats](#example-5-persistent-reconnecting-discovery-with-heartbeats) +- [Example 6: Simulated Distributed RL Workload (Local)](#example-6-simulated-distributed-rl-workload-local) +- [Example 7: Simulated Distributed RL Workload on Kubernetes](#example-7-simulated-distributed-rl-workload-on-kubernetes) +- [Example 8: Distributed RL Generation with vLLM Workers](#example-8-distributed-rl-generation-with-vllm-workers) --- @@ -28,7 +29,7 @@ Before running any examples, generate the required protobuf Python stubs from th python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. \ tunix/experimental/distributed/runtime/discovery/discovery_service.proto -# Compile the RL simulation service proto (required for Examples 5 & 6) +# Compile the RL simulation service proto (required for Examples 6 & 7) python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. \ tunix/experimental/distributed/examples/rl/service.proto ``` @@ -222,25 +223,86 @@ python -m tunix.experimental.distributed.runtime.main \ --say="open the door" ``` +--- + +## Example 5: Persistent Reconnecting Discovery with Heartbeats + +Building on the basic peer discovery concept introduced in [Example 4](#example-4-peer-discovery-and-inter-process-communication), this example demonstrates how to establish a persistent reconnecting session with automatic server restart recovery and heartbeat leases (using a WiFi hotspot and connecting phone metaphor). + +### 1. WiFi Hotspot Server (`wifi.py`) +The `wifi` process starts a persistent discovery server on port `12345` (`--discovery_port=12345`) and registers callbacks for client connection and disconnection events: + +```python +def main(argv: list[str], context: ProcessContext | None) -> None: + context.ipc.discovery.on_connect( + on_client_connected=lambda client_id, h, p, m, is_rec: logging.info( + f"Phone {client_id} connected (reconnect={is_rec})" + ), + on_client_disconnected=lambda client_id, h, p, reason: logging.warning( + f"Phone {client_id} disconnected ({reason})" + ), + ) +``` + +### 2. WiFi Phone Client (`phone.py`) +The `phone` process establishes a persistent reconnecting session with the WiFi hotspot. If the hotspot restarts or its network signal is briefly lost, the phone automatically re-registers: + +```python +def main(argv: list[str], context: ProcessContext | None) -> None: + client = context.ipc.discovery.connect( + metadata=None, + client_id="Pixel 9 Pro", + on_connected=lambda epoch, is_rec: logging.info(f"Connected (epoch: {epoch})"), + on_disconnected=lambda epoch, reason: logging.warning(f"Disconnected ({reason})"), + ) +``` + +### Run Locally + +```shell +# Terminal 1: Start the persistent WiFi hotspot discovery server +python -m tunix.experimental.distributed.runtime.main \ + --process_main=tunix.experimental.distributed.examples.basics.wifi.main \ + --discovery_id=wifi \ + --discovery_port=12345 +# Press Ctrl+C to stop the server. +# Rerun the command to restart the server. + +# Terminal 2: Start the persistent phone WiFi client +python -m tunix.experimental.distributed.runtime.main \ + --process_main=tunix.experimental.distributed.examples.basics.phone.main \ + --discovery_addrs=wifi:12345 \ + --model="Pixel 9 Pro" +``` + ### Expected Output -#### Door Terminal +#### WiFi Hotspot Terminal (Terminal 1) ``` -this is door! +this is wifi! +discovery server started on port 12345 +Phone Pixel 9 Pro connected (reconnect=False) +^Cdiscovery server stopped + + +this is wifi! discovery server started on port 12345 -localhost knocked and said: open the door -discovery server stopped +Phone Pixel 9 Pro connected (reconnect=False) ``` -#### Knocker Terminal +#### Phone Terminal (Terminal 2) ``` -this is knocker! -registered to discovery server at localhost:12345 +this is phone! +connecting to discovery server at localhost:12345 +Connected (epoch: uuid-1) +connected to discovery server at localhost:12345 +Disconnected (rpc_error) +Connected (epoch: uuid-2) ``` --- -## Example 5: Simulated Distributed RL Workload (Local) +## Example 6: Simulated Distributed RL Workload (Local) This example simulates a distributed reinforcement learning (RL) training workflow across **4 collaborating processes**: @@ -293,7 +355,7 @@ python -m tunix.experimental.distributed.runtime.main \ --- -## Example 6: Simulated Distributed RL Workload on Kubernetes +## Example 7: Simulated Distributed RL Workload on Kubernetes You can execute the exact same distributed RL simulation on a Kubernetes cluster using the `K8sExecutor` and JobSet deployment templates. @@ -314,7 +376,7 @@ bash tunix/experimental/distributed/examples/rl/launcher.sh --role=trainer --- -## Example 7: Distributed RL Generation with vLLM Workers +## Example 8: Distributed RL Generation with vLLM Workers This example demonstrates a distributed reinforcement learning generation pipeline using **vLLM** and the Tunix remote execution framework (`remote_execution.GrpcRemoteExecutionServer` / `ActorHandle`). @@ -367,6 +429,3 @@ Sample Response: To solve the problem, we need to add the two numbers 123 and 45 ------------------------ ``` - - - diff --git a/tunix/experimental/distributed/examples/basics/phone.py b/tunix/experimental/distributed/examples/basics/phone.py new file mode 100644 index 000000000..a510ef9e4 --- /dev/null +++ b/tunix/experimental/distributed/examples/basics/phone.py @@ -0,0 +1,46 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import pickle +import time + +from tunix.experimental.distributed.runtime.context import ProcessContext + + +def main(argv: list[str], context: ProcessContext | None) -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--message", type=str, default="this is phone!", help="") + args = parser.parse_args(argv) + + logging.info(args.message) + + assert context is not None + client = context.ipc.discovery.connect( + metadata=b"", + client_id="Pixel 9 Pro", + on_connected=lambda epoch, is_rec: logging.info( + f"Connected (epoch: {epoch})" + ), + on_disconnected=lambda epoch, reason: logging.warning( + f"Disconnected ({reason})" + ), + ) + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + client.stop() diff --git a/tunix/experimental/distributed/examples/basics/wifi.py b/tunix/experimental/distributed/examples/basics/wifi.py new file mode 100644 index 000000000..6b74a7ddf --- /dev/null +++ b/tunix/experimental/distributed/examples/basics/wifi.py @@ -0,0 +1,44 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import pickle +import time + +from tunix.experimental.distributed.runtime.context import ProcessContext + + +def main(argv: list[str], context: ProcessContext | None) -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--message", type=str, default="this is wifi!", help="") + args = parser.parse_args(argv) + + logging.info(args.message) + + assert context is not None + context.ipc.discovery.on_connect( + on_client_connected=lambda client_id, h, p, m, is_rec: logging.info( + f"Phone {client_id} connected (reconnect={is_rec})" + ), + on_client_disconnected=lambda client_id, h, p, reason: logging.warning( + f"Phone {client_id} disconnected ({reason})" + ), + ) + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + pass diff --git a/tunix/experimental/distributed/runtime/context.py b/tunix/experimental/distributed/runtime/context.py index 628dae263..985ce780c 100644 --- a/tunix/experimental/distributed/runtime/context.py +++ b/tunix/experimental/distributed/runtime/context.py @@ -45,6 +45,42 @@ def register(self, metadata: bytes) -> None: """ pass + def on_connect( + self, + on_client_connected: ( + Callable[[str, str, int, bytes, bool], None] | None + ) = None, + *, + on_client_disconnected: ( + Callable[[str, str, int, str], None] | None + ) = None, + ) -> None: + """Configures server-side handlers for managed client connections and eviction. + + Args: + on_client_connected: Invoked when a client connects or re-connects. + on_client_disconnected: Invoked when a connected client misses heartbeats. + """ + pass + + def connect( + self, + metadata: bytes, + *, + client_id: str, + on_connected: Callable[[str, bool], None] | None = None, + on_disconnected: Callable[[str, str], None] | None = None, + ) -> Any: + """Establishes a persistent reconnecting discovery session with heartbeats. + + Args: + metadata: Serialized metadata bytes describing this worker node. + client_id: Unique client identifier. + on_connected: Invoked when connection succeeds (initial or reconnect). + on_disconnected: Invoked when server session breaks or epoch changes. + """ + pass + class IpcContext: """Abstract interface providing inter-process communication contexts.""" diff --git a/tunix/experimental/distributed/runtime/contexts/k8s_context.py b/tunix/experimental/distributed/runtime/contexts/k8s_context.py index 35c961e4c..aae951684 100644 --- a/tunix/experimental/distributed/runtime/contexts/k8s_context.py +++ b/tunix/experimental/distributed/runtime/contexts/k8s_context.py @@ -104,6 +104,7 @@ def __init__(self, args: argparse.Namespace) -> None: """ self._args = args self._server = discovery.DiscoveryServer() + self._client: discovery.DiscoveryClient | None = None def __enter__(self) -> "K8sDiscoveryContext": """Enters the discovery context manager scope.""" @@ -115,7 +116,10 @@ def __exit__( exc: Any | None, tb: Any | None, ) -> None: - """Stops the discovery server if started.""" + """Stops the discovery client and server if started.""" + if self._client is not None: + self._client.stop() + self._client = None if self._server.is_started(): self._server.stop() logging.info("discovery server stopped") @@ -126,7 +130,28 @@ def on_register(self, callback: Callable[[str, int, bytes], None]) -> None: Args: callback: Invoked when a peer registers with this server. """ - self._server.start(self._args.discovery_port, callback) + self._server.on_register(callback) + self._server.start(self._args.discovery_port) + logging.info( + "discovery server started on port %s", self._args.discovery_port + ) + + def on_connect( + self, + on_client_connected: ( + Callable[[str, str, int, bytes, bool], None] | None + ) = None, + *, + on_client_disconnected: ( + Callable[[str, str, int, str], None] | None + ) = None, + ) -> None: + """Configures handlers and starts the discovery server on the configured port.""" + self._server.on_connect( + on_client_connected=on_client_connected, + on_client_disconnected=on_client_disconnected, + ) + self._server.start(self._args.discovery_port) logging.info( "discovery server started on port %s", self._args.discovery_port ) @@ -147,6 +172,32 @@ def register(self, metadata: bytes) -> None: ) logging.info("registered to discovery server at %s", server_address) + def connect( + self, + metadata: bytes, + *, + client_id: str, + on_connected: Callable[[str, bool], None] | None = None, + on_disconnected: Callable[[str, str], None] | None = None, + ) -> discovery.DiscoveryClient: + """Establishes a persistent reconnecting discovery session with heartbeats.""" + server_address = resolve_discovery_address(self._args.discovery_addrs) + + hostname = resolve_self_hostname() + + logging.info("connecting to discovery server at %s", server_address) + self._client = discovery.connect( + server_address, + hostname, + self._args.discovery_port, + metadata, + client_id=client_id, + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + logging.info("connected to discovery server at %s", server_address) + return self._client + class K8sIpcContext(context.IpcContext): """Kubernetes inter-process communication context.""" diff --git a/tunix/experimental/distributed/runtime/contexts/local_context.py b/tunix/experimental/distributed/runtime/contexts/local_context.py index e4b838b67..a89b7b664 100644 --- a/tunix/experimental/distributed/runtime/contexts/local_context.py +++ b/tunix/experimental/distributed/runtime/contexts/local_context.py @@ -47,6 +47,7 @@ def __init__(self, args: argparse.Namespace) -> None: """ self._args = args self._server = discovery.DiscoveryServer() + self._client: discovery.DiscoveryClient | None = None def __enter__(self) -> "LocalDiscoveryContext": """Enters the discovery context manager scope.""" @@ -58,7 +59,10 @@ def __exit__( exc: Any | None, tb: Any | None, ) -> None: - """Stops the local discovery server if started.""" + """Stops the local discovery client and server if started.""" + if self._client is not None: + self._client.stop() + self._client = None if self._server.is_started(): self._server.stop() logging.info("discovery server stopped") @@ -69,7 +73,28 @@ def on_register(self, callback: Callable[[str, int, bytes], None]) -> None: Args: callback: Invoked when a peer registers with this server. """ - self._server.start(self._args.discovery_port, callback) + self._server.on_register(callback) + self._server.start(self._args.discovery_port) + logging.info( + "discovery server started on port %s", self._args.discovery_port + ) + + def on_connect( + self, + on_client_connected: ( + Callable[[str, str, int, bytes, bool], None] | None + ) = None, + *, + on_client_disconnected: ( + Callable[[str, str, int, str], None] | None + ) = None, + ) -> None: + """Configures handlers and starts the local discovery server on the configured port.""" + self._server.on_connect( + on_client_connected=on_client_connected, + on_client_disconnected=on_client_disconnected, + ) + self._server.start(self._args.discovery_port) logging.info( "discovery server started on port %s", self._args.discovery_port ) @@ -90,6 +115,32 @@ def register(self, metadata: bytes) -> None: ) logging.info("registered to discovery server at %s", server_address) + def connect( + self, + metadata: bytes, + *, + client_id: str, + on_connected: Callable[[str, bool], None] | None = None, + on_disconnected: Callable[[str, str], None] | None = None, + ) -> discovery.DiscoveryClient: + """Establishes a persistent reconnecting discovery session with heartbeats.""" + server_address = resolve_discovery_address(self._args.discovery_addrs) + + hostname = "localhost" + + logging.info("connecting to discovery server at %s", server_address) + self._client = discovery.connect( + server_address, + hostname, + self._args.discovery_port, + metadata, + client_id=client_id, + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + logging.info("connected to discovery server at %s", server_address) + return self._client + class LocalIpcContext(context.IpcContext): """Inter-process communication context for local execution.""" diff --git a/tunix/experimental/distributed/runtime/discovery/discovery.py b/tunix/experimental/distributed/runtime/discovery/discovery.py index 243fc15f4..b73ff1437 100644 --- a/tunix/experimental/distributed/runtime/discovery/discovery.py +++ b/tunix/experimental/distributed/runtime/discovery/discovery.py @@ -15,89 +15,302 @@ """gRPC-based peer discovery server and client helper functions.""" from concurrent import futures +import dataclasses +import logging +import threading import time -from typing import Callable +from typing import Any, Callable +import uuid import grpc from tunix.experimental.distributed.runtime.discovery import discovery_service_pb2 as pb2 from tunix.experimental.distributed.runtime.discovery import discovery_service_pb2_grpc as pb2_grpc +class _RegistryServicer(pb2_grpc.DiscoveryServiceServicer): + """gRPC servicer for one-shot fire-and-forget worker registration.""" + + def __init__( + self, callback: Callable[[str, int, bytes], None] | None = None + ) -> None: + self._callback = callback + + def Register( + self, request: pb2.RegisterRequest, context: grpc.ServicerContext + ) -> pb2.RegisterResponse: + if self._callback is not None: + try: + self._callback(request.hostname, request.port, request.metadata) + except Exception as e: # pylint: disable=broad-except + logging.exception("Error in discovery server callback: %s", e) + return pb2.RegisterResponse() + + def Connect( + self, request: pb2.ConnectRequest, context: grpc.ServicerContext + ) -> pb2.ConnectResponse: + context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "Discovery server is running in register mode. Use Register" + " RPC instead of Connect.", + ) + return pb2.ConnectResponse() + + def Heartbeat( + self, request: pb2.HeartbeatRequest, context: grpc.ServicerContext + ) -> pb2.HeartbeatResponse: + context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "Discovery server is running in register mode. Heartbeats not" + " supported.", + ) + return pb2.HeartbeatResponse() + + +class _ConnectionServicer(pb2_grpc.DiscoveryServiceServicer): + """gRPC servicer for managed persistent client connections with heartbeats and lease eviction.""" + + @dataclasses.dataclass + class ClientConnection: + client_id: str + hostname: str + port: int + metadata: bytes + last_seen: float + + def __init__( + self, + on_client_connected: ( + Callable[[str, str, int, bytes, bool], None] | None + ) = None, + on_client_disconnected: ( + Callable[[str, str, int, str], None] | None + ) = None, + heartbeat_sec: int = 5, + ) -> None: + self._on_client_connected = on_client_connected + self._on_client_disconnected = on_client_disconnected + self._heartbeat_sec = heartbeat_sec + self._server_epoch = str(uuid.uuid4()) + self._connected_clients: dict[str, _ConnectionServicer.ClientConnection] = ( + {} + ) + self._lock = threading.Lock() + self._stop_event = threading.Event() + self._evictor_thread: threading.Thread | None = None + evictor_thread = threading.Thread( + target=self._run_evictor_loop, daemon=True + ) + evictor_thread.start() + self._evictor_thread = evictor_thread + + def Register( + self, request: pb2.RegisterRequest, context: grpc.ServicerContext + ) -> pb2.RegisterResponse: + context.abort( + grpc.StatusCode.UNIMPLEMENTED, + "Discovery server is running in connect mode. Use Connect" + " RPC instead of Register.", + ) + return pb2.RegisterResponse() + + def Connect( + self, request: pb2.ConnectRequest, context: grpc.ServicerContext + ) -> pb2.ConnectResponse: + client_id = request.client_id or f"{request.hostname}:{request.port}" + now = time.time() + + with self._lock: + is_reconnect = client_id in self._connected_clients + self._connected_clients[client_id] = _ConnectionServicer.ClientConnection( + client_id=client_id, + hostname=request.hostname, + port=request.port, + metadata=request.metadata, + last_seen=now, + ) + + if self._on_client_connected is not None: + try: + self._on_client_connected( + client_id, + request.hostname, + request.port, + request.metadata, + is_reconnect, + ) + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Error in discovery server on_client_connected callback: %s", e + ) + + return pb2.ConnectResponse( + server_epoch=self._server_epoch, + heartbeat_sec=self._heartbeat_sec, + ) + + def Heartbeat( + self, request: pb2.HeartbeatRequest, context: grpc.ServicerContext + ) -> pb2.HeartbeatResponse: + with self._lock: + if ( + request.server_epoch != self._server_epoch + or request.client_id not in self._connected_clients + ): + return pb2.HeartbeatResponse( + action=pb2.HEARTBEAT_ACTION_RE_REGISTER, + server_epoch=self._server_epoch, + heartbeat_sec=self._heartbeat_sec, + ) + + self._connected_clients[request.client_id].last_seen = time.time() + return pb2.HeartbeatResponse( + action=pb2.HEARTBEAT_ACTION_OK, + server_epoch=self._server_epoch, + heartbeat_sec=self._heartbeat_sec, + ) + + def _run_evictor_loop(self) -> None: + while not self._stop_event.is_set(): + self._stop_event.wait(timeout=float(self._heartbeat_sec)) + if self._stop_event.is_set(): + break + + now = time.time() + evicted: list[_ConnectionServicer.ClientConnection] = [] + + with self._lock: + timeout_threshold = 3 * self._heartbeat_sec + stale_ids = [ + cid + for cid, reg in self._connected_clients.items() + if now - reg.last_seen > timeout_threshold + ] + for cid in stale_ids: + evicted.append(self._connected_clients.pop(cid)) + + if self._on_client_disconnected is not None: + for reg in evicted: + try: + self._on_client_disconnected( + reg.client_id, reg.hostname, reg.port, "heartbeat_timeout" + ) + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Error in discovery server on_client_disconnected callback: %s", + e, + ) + + def stop(self, timeout: float | None = None) -> None: + """Stops the persistent servicer evictor thread.""" + self._stop_event.set() + if self._evictor_thread and self._evictor_thread.is_alive(): + self._evictor_thread.join(timeout=timeout) + self._evictor_thread = None + + class DiscoveryServer: - """Lightweight gRPC server for registering distributed worker nodes.""" + """Lightweight gRPC server for discovery, operating in either register or connect mode.""" - def __init__(self) -> None: + def __init__(self, heartbeat_sec: int = 5) -> None: """Initializes an unstarted discovery server instance.""" self._server: grpc.Server | None = None + self._executor: futures.ThreadPoolExecutor | None = None + self._servicer: _RegistryServicer | _ConnectionServicer | None = None + self._mode: str | None = None # "register" or "connect" + self._heartbeat_sec: int = heartbeat_sec + + # Registered configuration arguments + self._on_client_register: Callable[[str, int, bytes], None] | None = None + self._on_client_connected: ( + Callable[[str, str, int, bytes, bool], None] | None + ) = None + self._on_client_disconnected: ( + Callable[[str, str, int, str], None] | None + ) = None def is_started(self) -> bool: """Returns True if the discovery server is running.""" return self._server is not None - def start( - self, port: int, callback: Callable[[str, int, bytes], None] - ) -> None: - """Starts the discovery gRPC server on the given port. + def on_register(self, callback: Callable[[str, int, bytes], None]) -> None: + """Configures register mode for the server.""" + if self._mode == "connect": + raise RuntimeError( + "Cannot configure on_register when on_connect is already configured." + ) + self._mode = "register" + self._on_client_register = callback - Args: - port: Network port on which the gRPC discovery server listens. - callback: Function invoked when a peer node registers via RPC. + def on_connect( + self, + on_client_connected: ( + Callable[[str, str, int, bytes, bool], None] | None + ) = None, + *, + on_client_disconnected: ( + Callable[[str, str, int, str], None] | None + ) = None, + ) -> None: + """Configures connect mode for the server.""" + if self._mode == "register": + raise RuntimeError( + "Cannot configure on_connect when on_register is already configured." + ) + self._mode = "connect" + self._on_client_connected = on_client_connected + self._on_client_disconnected = on_client_disconnected - Raises: - ValueError: If `port` is zero or invalid. - RuntimeError: If the server has already been started. - """ + def start( + self, + port: int, + heartbeat_sec: int = 5, + ) -> None: + """Starts the discovery gRPC server on the given port.""" if not port: raise ValueError("port must be non-zero. did you set --discovery_port ?") if self._server is not None: raise RuntimeError("server already started") + if self._mode is None: + raise RuntimeError( + "Discovery server mode not configured. Call on_register() or" + " on_connect() before starting the server." + ) - server = grpc.server(futures.ThreadPoolExecutor(max_workers=10)) + if self._mode == "connect": + self._servicer = _ConnectionServicer( + on_client_connected=self._on_client_connected, + on_client_disconnected=self._on_client_disconnected, + heartbeat_sec=heartbeat_sec or self._heartbeat_sec, + ) + else: + self._servicer = _RegistryServicer(self._on_client_register) - # define and register handler - class _handler(pb2_grpc.DiscoveryServiceServicer): - - def Register( - self, request: pb2.RegisterRequest, context: grpc.ServicerContext - ): - callback(request.hostname, request.port, request.metadata) - return pb2.RegisterResponse() - - pb2_grpc.add_DiscoveryServiceServicer_to_server(_handler(), server) - - # start server + self._executor = futures.ThreadPoolExecutor(max_workers=10) + server = grpc.server(self._executor) + pb2_grpc.add_DiscoveryServiceServicer_to_server(self._servicer, server) server.add_insecure_port(f"[::]:{port}") server.start() self._server = server def stop(self, timeout: float | None = None) -> None: - """Stops the discovery gRPC server and waits for termination. + """Stops the discovery gRPC server and waits for termination.""" + if isinstance(self._servicer, _ConnectionServicer): + self._servicer.stop(timeout=timeout) - Args: - timeout: Grace period in seconds to wait for active RPCs to terminate. - """ if self._server: self._server.stop(timeout) self._server.wait_for_termination(timeout) self._server = None + if self._executor: + self._executor.shutdown(wait=True) + self._executor = None + self._servicer = None def register( server_address: str, hostname: str, port: int, metadata: bytes ) -> None: - """Registers a node with the remote discovery server using exponential backoff. - - Args: - server_address: Host and port of the target discovery server (e.g. - 'host:port'). - hostname: Hostname or address of the registering node. - port: Port of the registering node. - metadata: Custom serialized metadata bytes to pass to the server. - - Raises: - ValueError: If `server_address` is empty. - RuntimeError: If registration fails with a non-retryable gRPC error. - """ + """Registers a node with the remote discovery server using exponential backoff (one-shot).""" if not server_address: raise ValueError( "server_address must be non-empty. did you set --discovery_addrs ?" @@ -124,3 +337,186 @@ def register( raise RuntimeError( f"discovery register failed: {e.code()} - {e.details()}" # pytype: disable=attribute-error ) + + +class DiscoveryClient: + """Manages persistent worker registration and background heartbeat loop.""" + + def __init__( + self, + server_address: str, + hostname: str, + port: int, + metadata: bytes, + client_id: str, + *, + on_connected: Callable[[str, bool], None] | None = None, + on_disconnected: Callable[[str, str], None] | None = None, + ) -> None: + self._server_address = server_address + self._hostname = hostname + self._port = port + self._metadata = metadata + self._client_id = client_id + self._on_connected = on_connected + self._on_disconnected = on_disconnected + + self._server_epoch: str = "" + self._heartbeat_sec: int = 5 + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._channel: grpc.Channel | None = None + self._stub: pb2_grpc.DiscoveryServiceStub | None = None + + def start(self) -> None: + """Starts the discovery client, performs initial connect, and launches background heartbeats.""" + if not self._server_address: + raise ValueError( + "server_address must be non-empty. did you set --discovery_addrs ?" + ) + + self._channel = grpc.insecure_channel(self._server_address) + self._stub = pb2_grpc.DiscoveryServiceStub(self._channel) + + response = self._connect_with_backoff() + self._server_epoch = response.server_epoch + self._heartbeat_sec = response.heartbeat_sec or 5 + + if self._on_connected is not None: + try: + self._on_connected(self._server_epoch, False) + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Error in discovery client on_connected callback: %s", e + ) + + self._stop_event.clear() + self._thread = threading.Thread( + target=self._run_heartbeat_loop, daemon=True + ) + self._thread.start() + + def _connect_with_backoff(self) -> pb2.ConnectResponse: + request = pb2.ConnectRequest( + client_id=self._client_id, + hostname=self._hostname, + port=self._port, + metadata=self._metadata, + ) + delay = 1 + while not self._stop_event.is_set(): + try: + assert self._stub is not None + return self._stub.Connect(request) + except grpc.RpcError as e: + if e.code() == grpc.StatusCode.UNAVAILABLE: # pytype: disable=attribute-error + if self._stop_event.wait(delay): + break + delay = min(delay * 2, 60) + continue + else: + raise RuntimeError( + f"discovery connect failed: {e.code()} - {e.details()}" # pytype: disable=attribute-error + ) + raise RuntimeError("discovery client stopped during connect") + + def _run_heartbeat_loop(self) -> None: + while not self._stop_event.is_set(): + self._stop_event.wait(timeout=float(self._heartbeat_sec)) + if self._stop_event.is_set(): + break + + try: + assert self._stub is not None + req = pb2.HeartbeatRequest( + client_id=self._client_id, server_epoch=self._server_epoch + ) + resp = self._stub.Heartbeat(req) + + if resp.action == pb2.HEARTBEAT_ACTION_OK: + if resp.heartbeat_sec: + self._heartbeat_sec = resp.heartbeat_sec + elif resp.action == pb2.HEARTBEAT_ACTION_RE_REGISTER: + old_epoch = self._server_epoch + if self._on_disconnected is not None: + try: + self._on_disconnected(old_epoch, "epoch_mismatch") + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Error in discovery client on_disconnected callback: %s", e + ) + + new_resp = self._connect_with_backoff() + self._server_epoch = new_resp.server_epoch + if new_resp.heartbeat_sec: + self._heartbeat_sec = new_resp.heartbeat_sec + + if self._on_connected is not None: + try: + self._on_connected(self._server_epoch, True) + except Exception as e: # pylint: disable=broad-except + logging.exception( + "Error in discovery client on_connected callback: %s", e + ) + except grpc.RpcError as e: + old_epoch = self._server_epoch + if self._on_disconnected is not None: + try: + self._on_disconnected(old_epoch, "rpc_error") + except Exception as ex: # pylint: disable=broad-except + logging.exception( + "Error in discovery client on_disconnected callback: %s", ex + ) + + try: + new_resp = self._connect_with_backoff() + self._server_epoch = new_resp.server_epoch + if new_resp.heartbeat_sec: + self._heartbeat_sec = new_resp.heartbeat_sec + + if self._on_connected is not None: + try: + self._on_connected(self._server_epoch, True) + except Exception as ex: # pylint: disable=broad-except + logging.exception( + "Error in discovery client on_connected callback: %s", ex + ) + except Exception as ex: # pylint: disable=broad-except + logging.exception( + "Failed to re-connect with discovery server: %s", ex + ) + + def stop(self, timeout: float | None = None) -> None: + """Stops the discovery client and background heartbeat loop.""" + self._stop_event.set() + if self._thread and self._thread.is_alive(): + self._thread.join(timeout=timeout) + self._thread = None + if self._channel: + self._channel.close() + self._channel = None + self._stub = None + + +def connect( + server_address: str, + hostname: str, + port: int, + metadata: bytes, + client_id: str, + *, + on_connected: Callable[[str, bool], None] | None = None, + on_disconnected: Callable[[str, str], None] | None = None, +) -> DiscoveryClient: + """Establishes a persistent connection with heartbeats.""" + client = DiscoveryClient( + server_address=server_address, + hostname=hostname, + port=port, + metadata=metadata, + client_id=client_id, + on_connected=on_connected, + on_disconnected=on_disconnected, + ) + client.start() + return client diff --git a/tunix/experimental/distributed/runtime/discovery/discovery_service.proto b/tunix/experimental/distributed/runtime/discovery/discovery_service.proto index 510de886e..dc2c7b611 100644 --- a/tunix/experimental/distributed/runtime/discovery/discovery_service.proto +++ b/tunix/experimental/distributed/runtime/discovery/discovery_service.proto @@ -35,8 +35,53 @@ message RegisterRequest { // Response returned by the discovery service after successful registration. message RegisterResponse {} +// Request message sent by client nodes when establishing a persistent +// connection session. +message ConnectRequest { + // Unique client identifier (e.g. 'hostname:port' or UUID). + string client_id = 1; + // Hostname or reachable FQDN of the registering worker node. + string hostname = 2; + // Network port on which the registering node is listening. + uint32 port = 3; + // Custom serialized metadata bytes describing this node. + bytes metadata = 4; + // Optional requested heartbeat interval hint in seconds. + uint32 requested_heartbeat_sec = 5; +} + +// Response returned by the discovery service after successful connection. +message ConnectResponse { + // Unique server boot instance ID. + string server_epoch = 1; + // Authoritative heartbeat interval dictated by the server. + uint32 heartbeat_sec = 2; +} + +// Action returned by discovery service in response to a client heartbeat. +enum HeartbeatAction { + HEARTBEAT_ACTION_UNSPECIFIED = 0; + HEARTBEAT_ACTION_OK = 1; + HEARTBEAT_ACTION_RE_REGISTER = 2; +} + +message HeartbeatRequest { + string client_id = 1; + string server_epoch = 2; +} + +message HeartbeatResponse { + HeartbeatAction action = 1; + string server_epoch = 2; + uint32 heartbeat_sec = 3; +} + // Discovery service used by distributed workers to register their coordinates. service DiscoveryService { - // Registers a worker node with the discovery service. + // Registers a worker node with the discovery service (one-shot). rpc Register(RegisterRequest) returns (RegisterResponse); + // Establishes a persistent connection session with heartbeats. + rpc Connect(ConnectRequest) returns (ConnectResponse); + // Sends periodic heartbeats from worker node to discovery service. + rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse); }