From 1a489d951b2211d8b0435e5929e6edc4118eb00a Mon Sep 17 00:00:00 2001 From: Yangmu Jiang Date: Wed, 5 Aug 2026 11:24:08 -0700 Subject: [PATCH] test: add worker transport validation test. PiperOrigin-RevId: 959770319 --- .../worker_transport/worker_transport_test.py | 159 ++++++++++++++++++ .../examples/worker_transport/README.md | 93 ++++++++++ .../worker_transport/remote_worker_server.py | 66 ++++++++ .../examples/worker_transport/transport.py | 46 +++++ .../examples/worker_transport/worker.py | 25 +++ 5 files changed, 389 insertions(+) create mode 100644 tests/experimental/distributed/examples/worker_transport/worker_transport_test.py create mode 100644 tunix/experimental/distributed/examples/worker_transport/README.md create mode 100644 tunix/experimental/distributed/examples/worker_transport/remote_worker_server.py create mode 100644 tunix/experimental/distributed/examples/worker_transport/transport.py create mode 100644 tunix/experimental/distributed/examples/worker_transport/worker.py diff --git a/tests/experimental/distributed/examples/worker_transport/worker_transport_test.py b/tests/experimental/distributed/examples/worker_transport/worker_transport_test.py new file mode 100644 index 000000000..63c690c84 --- /dev/null +++ b/tests/experimental/distributed/examples/worker_transport/worker_transport_test.py @@ -0,0 +1,159 @@ +# 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. + +"""Validation demo and test for local and remote worker transport. + +This module demonstrates and tests unified execution transport parity across +local and remote workers using `tunix.experimental.worker.remote_execution`: + 1. Local Worker: Co-located in the orchestrator process via in-process + transport (`transport.local(...)`). + 2. Remote Worker: Running in a separate OS subprocess via gRPC transport + (`transport.remote(...)`), whose address is dynamically resolved using + Tunix + peer discovery (`context.ipc.discovery`). + +Both transport modes are validated using a uniform orchestrator test loop +(`validate_worker_transport`), ensuring identical method invocation behavior +regardless of the underlying transport layer. +""" + +import argparse +import asyncio +import os +import pickle +import subprocess +import sys +import threading + +from absl.testing import absltest +import portpicker +from tunix.experimental.distributed.examples.worker_transport import transport +from tunix.experimental.distributed.examples.worker_transport.worker import Worker +from tunix.experimental.distributed.runtime.contexts.local_context import LocalProcessContext +from tunix.experimental.worker import remote_execution + + +def start_remote_worker_process( + service_port: int, discovery_port: int, discovery_addrs: str +) -> subprocess.Popen: + """Dedicated helper method to spawn the remote worker server in a separate process.""" + remote_worker_bin = os.path.join( + os.path.dirname(sys.argv[0]), "remote_worker_server" + ) + if os.path.exists(remote_worker_bin): + cmd = [remote_worker_bin] + else: + python_bin = sys.executable or "python" + cmd = [ + python_bin, + "-m", + "tunix.experimental.distributed.examples.worker_transport.remote_worker_server", + ] + return subprocess.Popen( + cmd + + [ + f"--port={service_port}", + "--discovery_id=remote_worker_server", + f"--discovery_port={discovery_port}", + f"--discovery_addrs={discovery_addrs}", + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +async def validate_worker_transport( + handle: remote_execution.ActorHandle, expected_ack: str +) -> None: + """Connects Orchestrator to a worker via unified ActorHandle interface to validate transport.""" + res = None + for _ in range(50): + try: + res = await handle.asubmit("ping", msg="hello") + break + except Exception: # pylint: disable=broad-exception-caught + await asyncio.sleep(0.1) + if res is None: + res = await handle.asubmit("ping", msg="hello") + + assert res == expected_ack, f"Expected {expected_ack}, got {res}" + + +class WorkerTransportTest(absltest.TestCase): + + def test_local_worker(self): + local_handle = transport.local(Worker, name="local") + asyncio.run( + validate_worker_transport( + local_handle, expected_ack="[local] ack: hello" + ) + ) + + def test_remote_worker(self): + discovery_port = portpicker.pick_unused_port() + service_port = portpicker.pick_unused_port() + + args = argparse.Namespace( + discovery_id="orchestrator", + discovery_port=discovery_port, + discovery_addrs=f"orchestrator:{discovery_port}", + ) + with LocalProcessContext(args) as ctx: + discovered_addr = None + discovery_event = threading.Event() + + def on_register(hostname: str, _: int, metadata: bytes) -> None: + nonlocal discovered_addr + md = pickle.loads(metadata) + discovered_addr = f"grpc://{hostname}:{md['service_port']}" + discovery_event.set() + + ctx.ipc.discovery.on_register(on_register) + + proc = start_remote_worker_process( + service_port=service_port, + discovery_port=portpicker.pick_unused_port(), + discovery_addrs=f"orchestrator:{discovery_port}", + ) + try: + if not discovery_event.wait(timeout=5.0): + if proc.poll() is not None: + out, err = proc.communicate() + raise RuntimeError(f"Remote worker process exited: {err}") + raise RuntimeError( + "Failed to resolve worker via discovery (timed out)" + ) + + assert ( + discovered_addr is not None + ), "Failed to resolve worker via discovery" + remote_handle = transport.remote(Worker, address=discovered_addr) + asyncio.run( + validate_worker_transport( + remote_handle, expected_ack="[remote] ack: hello" + ) + ) + finally: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=2.0) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + + +if __name__ == "__main__": + absltest.main() diff --git a/tunix/experimental/distributed/examples/worker_transport/README.md b/tunix/experimental/distributed/examples/worker_transport/README.md new file mode 100644 index 000000000..0e0ded4cf --- /dev/null +++ b/tunix/experimental/distributed/examples/worker_transport/README.md @@ -0,0 +1,93 @@ +# Tunix Worker Transport Example + +This example demonstrates how to use the Tunix worker execution transport layer to seamlessly run worker tasks either locally (in-process) or remotely (over gRPC) using a unified actor handle API. + +--- + +## 1. Overview & API + +With Tunix worker transport, an orchestrator can submit tasks to local or remote workers through identical `ActorHandle` interfaces. Code written against an actor handle works transparently regardless of where the worker is running. + +### Key API Functions (`transport.py`) + +- **`transport.local(cls, *args, **kwargs)`**: Creates an actor handle for a local, in-process instance of target class `cls` initialized with `*args, **kwargs`. Ideal for debugging, testing, or zero-serialization execution. +- **`transport.remote(cls, address: str)`**: Creates an actor handle connecting to a remote worker daemon for class `cls` at network address `address`. + +```python +from tunix.experimental.distributed.examples.worker_transport import transport +from tunix.experimental.distributed.examples.worker_transport.worker import Worker + +# Local in-process worker handle +local_handle = transport.local(Worker, name="local") + +# Remote gRPC worker handle +remote_handle = transport.remote(Worker, address="grpc://worker-host:12345") +``` + +--- + +## 2. Distributed Workflow + +### Step 1: Define a Worker Class + +Define a worker class containing the business logic or methods you wish to execute: + +```python +class Worker: + def __init__(self, name: str): + self.name = name + + def ping(self, msg: str) -> str: + return f"[{self.name}] ack: {msg}" +``` + +### Step 2: Choose a Worker Deployment Sub-workflow + +Tunix supports three worker deployment models depending on your execution environment: + +1. **Same-process Worker (In-process)**: + - Co-locates the worker instance within the orchestrator process. + - Ideal for debugging, testing, or zero-serialization execution. + ```python + handle = transport.local(Worker, name="same-process-worker") + ``` + +2. **Same-host Worker (Separate local process)**: + - Runs the worker inside an independent OS process on the same machine communicating via gRPC. + - Useful for isolating memory/GIL or running standalone binary entry points locally. + ```python + # Connect to worker daemon running on local port 12345: + handle = transport.remote(Worker, address="grpc://localhost:12345") + ``` + +3. **Remote-host Worker (Distributed network process / K8s pod)**: + - Deploys worker daemon processes across remote network hosts or Kubernetes pods. + - Workers dynamically register their network endpoints using Tunix peer discovery (`context.ipc.discovery`), allowing the orchestrator to resolve and connect without hardcoded IP addresses. + ```python + # Address resolved dynamically via peer discovery callback + handle = transport.remote(Worker, address=discovered_address) + ``` + +### Step 3: Submit Tasks to Workers via Actor Handles + +Regardless of whether the worker is running in the same process, on the same host, or on a remote host, the orchestrator interacts with worker handles using the identical asynchronous `asubmit` API: + +```python +async def run_workflow(handle): + # Submit method name and positional/keyword arguments + result = await handle.asubmit("ping", msg="hello") + print(result) # Output: "[worker-name] ack: hello" +``` + +--- + +## 3. Directory Structure + +``` +worker_transport/ +├── worker.py # Worker class definition +├── transport.py # Generic transport.local(cls, ...) and transport.remote(cls, address) API functions +├── remote_worker_server.py # Standalone remote worker daemon process entry point +├── worker_transport_test.py # Automated end-to-end integration test +└── README.md # User guide (this file) +``` diff --git a/tunix/experimental/distributed/examples/worker_transport/remote_worker_server.py b/tunix/experimental/distributed/examples/worker_transport/remote_worker_server.py new file mode 100644 index 000000000..49ff645f8 --- /dev/null +++ b/tunix/experimental/distributed/examples/worker_transport/remote_worker_server.py @@ -0,0 +1,66 @@ +# 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. + +"""Standalone remote worker process for worker_transport demo.""" + +import argparse +import asyncio +import pickle +from typing import Sequence + +from absl import app, flags +from tunix.experimental.distributed.examples.worker_transport.worker import Worker +from tunix.experimental.distributed.runtime.context import ProcessContext +from tunix.experimental.distributed.runtime.contexts.local_context import LocalProcessContext +from tunix.experimental.worker import remote_execution + +FLAGS = flags.FLAGS +flags.DEFINE_integer("port", 12345, "Port for remote worker gRPC server.") +flags.DEFINE_string("discovery_id", "remote_worker", "Discovery ID.") +flags.DEFINE_integer("discovery_port", 0, "Discovery port.") +flags.DEFINE_string("discovery_addrs", "", "Discovery server addresses.") + + +async def run_server_async(context: ProcessContext | None) -> None: + """Starts the gRPC worker server, registers with discovery after it starts, and waits for termination.""" + worker = Worker("remote") + server = remote_execution.GrpcRemoteExecutionServer(worker) + await server.start_serving_async(FLAGS.port) + + if context and context.ipc and context.ipc.discovery: + context.ipc.discovery.register( + metadata=pickle.dumps({"service_port": FLAGS.port}) + ) + + if server._server is not None: + await server._server.wait_for_termination() + + +def main(argv: Sequence[str], context: ProcessContext | None) -> None: + """Distributed process entry point for the remote worker.""" + del argv + if context is None and FLAGS.discovery_addrs: + args = argparse.Namespace( + discovery_id=FLAGS.discovery_id, + discovery_port=FLAGS.discovery_port or 0, + discovery_addrs=FLAGS.discovery_addrs, + ) + context = LocalProcessContext(args) + context.__enter__() + + asyncio.run(run_server_async(context)) + + +if __name__ == "__main__": + app.run(lambda argv: main(argv, context=None)) diff --git a/tunix/experimental/distributed/examples/worker_transport/transport.py b/tunix/experimental/distributed/examples/worker_transport/transport.py new file mode 100644 index 000000000..248acc6c5 --- /dev/null +++ b/tunix/experimental/distributed/examples/worker_transport/transport.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. + +"""Generic transport helper functions for worker actor creation. + +This module provides helper functions to initialize worker actor handles across +three deployment cases: + + 1. Same-process Worker (In-process): + local_handle = local(Worker, name="local") + + 2. Same-host Worker (Separate local OS process): + remote_handle = remote(Worker, address="grpc://localhost:12345") + + 3. Remote-host Worker (Distributed network process / K8s pod): + remote_handle = remote(Worker, address=discovered_address) +""" + +from typing import Any + +from tunix.experimental.worker import remote_execution + + +def local(cls: Any, *args: Any, **kwargs: Any) -> remote_execution.ActorHandle: + """Creates a local, in-process ActorHandle for target class `cls` instantiated with `*args, **kwargs`.""" + return remote_execution.remote(cls, transport="inprocess").remote( + *args, **kwargs + ) + + +def remote(cls: Any, address: str) -> remote_execution.ActorHandle: + """Creates a gRPC remote ActorHandle for target class `cls` at network address `address`.""" + return remote_execution.remote( + cls, transport="grpc", address=address + ).remote() diff --git a/tunix/experimental/distributed/examples/worker_transport/worker.py b/tunix/experimental/distributed/examples/worker_transport/worker.py new file mode 100644 index 000000000..d11154999 --- /dev/null +++ b/tunix/experimental/distributed/examples/worker_transport/worker.py @@ -0,0 +1,25 @@ +# 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. + +"""Common worker definition for worker_transport validation demo.""" + + +class Worker: + """Simple worker exposing a test method.""" + + def __init__(self, name: str): + self.name = name + + def ping(self, msg: str) -> str: + return f"[{self.name}] ack: {msg}"