test: add worker transport validation test.

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()

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)
+```

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))

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()

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}"