diff --git a/.github/workflows/cpu-tests.yml b/.github/workflows/cpu-tests.yml
index 6dd826193..eec904026 100644
--- a/.github/workflows/cpu-tests.yml
+++ b/.github/workflows/cpu-tests.yml
@@ -79,6 +79,12 @@ jobs:
         run: |
           python -m pytest tests/models/flash_attention_cache_test.py -v --tb=short
 
+      - name: Compile gRPC Protobuf definitions
+        run: |
+          python -m pip install grpcio-tools
+          find tunix/experimental/distributed -name "*.proto" -exec \
+            python -m grpc_tools.protoc -I. --python_out=. --grpc_python_out=. {} +
+
       - name: Run experimental tests
         run: |
           python -m pytest tests/experimental/ -v --tb=short
diff --git a/tests/experimental/distributed/deployment/yaml_generator_test.py b/tests/experimental/distributed/deployment/yaml_generator_test.py
new file mode 100644
index 000000000..00c3ac780
--- /dev/null
+++ b/tests/experimental/distributed/deployment/yaml_generator_test.py
@@ -0,0 +1,135 @@
+# 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.
+
+"""Unit tests for Kubernetes deployment YAML manifest generator."""
+
+import io
+import os
+import sys
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from tunix.experimental.distributed.deployment import yaml_generator
+
+
+def _get_template_path(filename: str) -> str:
+  # Locate template file relative to yaml_generator module directory
+  path = os.path.join(
+      os.path.dirname(yaml_generator.__file__),
+      "yamls",
+      filename,
+  )
+  return path
+
+
+class YamlGeneratorTest(parameterized.TestCase):
+
+  def test_generate_cpu_yaml(self):
+    template_file = _get_template_path("jobset.cpu.yaml")
+    argv = [
+        "yaml_generator.py",
+        template_file,
+        "--jobset_name=test-cpu-job",
+        "--cpu_machine=n2-standard-64",
+        "--worker_container_port=9999",
+    ]
+    with mock.patch.object(sys, "argv", argv):
+      with mock.patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+        yaml_generator.main()
+        rendered = mock_stdout.getvalue()
+        self.assertIn("test-cpu-job", rendered)
+        self.assertIn("9999", rendered)
+
+  @parameterized.named_parameters(
+      (
+          "tpu7x_small",
+          "tpu7x:4x4x4",
+          "test-tpu7x",
+          "tpu7x",
+      ),
+      (
+          "tpu7x_large",
+          "tpu7x:4x4x8",
+          "test-tpu7x-large",
+          "tpu7x",
+      ),
+      (
+          "tpuv5",
+          "tpuv5:2x2x2",
+          "test-tpuv5",
+          "tpu-v5p-slice",
+      ),
+      (
+          "tpuv5e",
+          "tpuv5e:2x4",
+          "test-tpuv5e",
+          "tpu-v5-lite-podslice",
+      ),
+      (
+          "tpuv6e",
+          "tpuv6e:2x4",
+          "test-tpuv6e",
+          "tpu-v6e-slice",
+      ),
+      (
+          "tpuv6ea",
+          "tpuv6ea:2x4",
+          "test-tpuv6ea",
+          "tpu-v6ea-slice",
+      ),
+  )
+  def test_generate_tpu_slice(
+      self, tpu_slice, jobset_name, expected_accelerator
+  ):
+    template_file = _get_template_path("jobset.pathways.yaml")
+    argv = [
+        "yaml_generator.py",
+        template_file,
+        f"--tpu_slice={tpu_slice}",
+        f"--jobset_name={jobset_name}",
+    ]
+    with mock.patch.object(sys, "argv", argv):
+      with mock.patch("sys.stdout", new_callable=io.StringIO) as mock_stdout:
+        yaml_generator.main()
+        rendered = mock_stdout.getvalue()
+        self.assertIn(expected_accelerator, rendered)
+        self.assertIn(jobset_name, rendered)
+
+  @parameterized.named_parameters(
+      (
+          "unsupported_tpu_type",
+          "unknown_tpu:4x4",
+          ValueError,
+      ),
+      (
+          "invalid_num_chips",
+          "tpu7x:1x2",
+          AssertionError,
+      ),
+  )
+  def test_invalid_slice_raises(self, tpu_slice, expected_exception):
+    template_file = _get_template_path("jobset.pathways.yaml")
+    argv = [
+        "yaml_generator.py",
+        template_file,
+        f"--tpu_slice={tpu_slice}",
+    ]
+    with mock.patch.object(sys, "argv", argv):
+      with self.assertRaises(expected_exception):
+        yaml_generator.main()
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/distributed/runtime/context_test.py b/tests/experimental/distributed/runtime/context_test.py
new file mode 100644
index 000000000..07e497b4b
--- /dev/null
+++ b/tests/experimental/distributed/runtime/context_test.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.
+
+"""Unit tests for base distributed runtime context interfaces."""
+
+from absl.testing import absltest
+from tunix.experimental.distributed.runtime import context
+
+
+class ContextTest(absltest.TestCase):
+
+  def test_jax_context_defaults(self):
+    jax_ctx = context.JaxContext()
+    self.assertIsNone(jax_ctx.initialize())
+
+  def test_discovery_context_defaults(self):
+    disc_ctx = context.DiscoveryContext()
+    self.assertIsNone(disc_ctx.on_register(lambda host, port, meta: None))
+    self.assertIsNone(disc_ctx.register(b"metadata"))
+
+  def test_ipc_context_defaults(self):
+    ipc_ctx = context.IpcContext()
+    self.assertIsInstance(ipc_ctx.discovery, context.DiscoveryContext)
+
+  def test_process_context_defaults(self):
+    proc_ctx = context.ProcessContext()
+    with proc_ctx as entered:
+      self.assertIs(entered, proc_ctx)
+    self.assertIsInstance(proc_ctx.jax, context.JaxContext)
+    self.assertIsInstance(proc_ctx.ipc, context.IpcContext)
+    self.assertIsInstance(proc_ctx.ipc.discovery, context.DiscoveryContext)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/distributed/runtime/contexts/k8s_context_test.py b/tests/experimental/distributed/runtime/contexts/k8s_context_test.py
new file mode 100644
index 000000000..b677b2678
--- /dev/null
+++ b/tests/experimental/distributed/runtime/contexts/k8s_context_test.py
@@ -0,0 +1,111 @@
+# 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.
+
+"""Unit tests for Kubernetes distributed runtime contexts."""
+
+import argparse
+import os
+import sys
+from unittest import mock
+
+from absl.testing import absltest
+import portpicker
+from tunix.experimental.distributed.runtime import context
+from tunix.experimental.distributed.runtime.contexts import k8s_context
+
+
+class K8sContextTest(absltest.TestCase):
+
+  def test_resolve_discovery_address(self):
+    resolved = k8s_context.resolve_discovery_address("door:12345")
+    self.assertEqual(resolved, "door-proc-0-0.door:12345")
+
+  def test_resolve_self_hostname_success(self):
+    envs = {
+        "JOBSET_NAME": "myjobset",
+        "REPLICATED_JOB_NAME": "worker",
+        "JOB_INDEX": "0",
+        "POD_INDEX": "2",
+    }
+    with mock.patch.dict(os.environ, envs):
+      fqdn = k8s_context.resolve_self_hostname()
+      self.assertEqual(fqdn, "myjobset-worker-0-2.myjobset")
+
+  def test_resolve_self_hostname_missing_env(self):
+    envs = {
+        "JOBSET_NAME": "myjobset",
+    }
+    with mock.patch.dict(os.environ, envs, clear=True):
+      with self.assertRaises(ValueError):
+        k8s_context.resolve_self_hostname()
+
+  def test_k8s_jax_context_pathways(self):
+    envs = {
+        "JAX_PLATFORMS": "proxy",
+        "JAX_BACKEND_TARGET": "0.0.0.0:8000",
+    }
+    mock_pw = mock.MagicMock()
+    with mock.patch.dict(os.environ, envs):
+      with mock.patch.dict(sys.modules, {"pathwaysutils": mock_pw}):
+        k8s_context.K8sJaxContext().initialize()
+        mock_pw.initialize.assert_called_once()
+
+  def test_k8s_jax_context_mcjax(self):
+    envs = {}
+    mock_jax = mock.MagicMock()
+    with mock.patch.dict(os.environ, envs, clear=True):
+      with mock.patch.dict(sys.modules, {"jax": mock_jax}):
+        k8s_context.K8sJaxContext().initialize()
+        mock_jax.distributed.initialize.assert_called_once()
+
+  def test_k8s_discovery_context_register(self):
+    envs = {
+        "JOBSET_NAME": "myjobset",
+        "REPLICATED_JOB_NAME": "worker",
+        "JOB_INDEX": "0",
+        "POD_INDEX": "1",
+    }
+    port = portpicker.pick_unused_port()
+    args = argparse.Namespace(
+        discovery_port=port,
+        discovery_addrs="door:8888",
+    )
+
+    with mock.patch.dict(os.environ, envs):
+      with k8s_context.K8sDiscoveryContext(args) as disc_ctx:
+        with mock.patch(
+            "tunix.experimental.distributed.runtime.contexts.k8s_context.discovery.register"
+        ) as mock_reg:
+          disc_ctx.register(b"pod-meta")
+          mock_reg.assert_called_once_with(
+              "door-proc-0-0.door:8888",
+              "myjobset-worker-0-1.myjobset",
+              port,
+              b"pod-meta",
+          )
+
+  def test_k8s_process_context(self):
+    args = argparse.Namespace(
+        discovery_port=portpicker.pick_unused_port(),
+        discovery_addrs="door:8888",
+    )
+    proc_ctx = k8s_context.K8sProcessContext(args)
+    with proc_ctx as entered:
+      self.assertIs(entered, proc_ctx)
+      self.assertIsInstance(proc_ctx.jax, k8s_context.K8sJaxContext)
+      self.assertIsInstance(proc_ctx.ipc, k8s_context.K8sIpcContext)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/distributed/runtime/contexts/local_context_test.py b/tests/experimental/distributed/runtime/contexts/local_context_test.py
new file mode 100644
index 000000000..acabb775d
--- /dev/null
+++ b/tests/experimental/distributed/runtime/contexts/local_context_test.py
@@ -0,0 +1,77 @@
+# 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.
+
+"""Unit tests for local distributed runtime contexts."""
+
+import argparse
+from unittest import mock
+
+from absl.testing import absltest
+import portpicker
+from tunix.experimental.distributed.runtime import context
+from tunix.experimental.distributed.runtime.contexts import local_context
+
+
+class LocalContextTest(absltest.TestCase):
+
+  def test_resolve_discovery_address(self):
+    resolved = local_context.resolve_discovery_address("worker-id:12345")
+    self.assertEqual(resolved, "localhost:12345")
+
+  @mock.patch(
+      "tunix.experimental.distributed.runtime.discovery.discovery.grpc.server"
+  )
+  def test_local_discovery_context_lifecycle_and_registration(
+      self, mock_grpc_server
+  ):
+    port = portpicker.pick_unused_port()
+    args = argparse.Namespace(
+        discovery_port=port,
+        discovery_addrs="leader:9999",
+    )
+
+    with local_context.LocalDiscoveryContext(args) as disc_ctx:
+      cb = mock.MagicMock()
+      disc_ctx.on_register(cb)
+      self.assertTrue(disc_ctx._server.is_started())
+
+      with mock.patch(
+          "tunix.experimental.distributed.runtime.contexts.local_context.discovery.register"
+      ) as mock_reg:
+        disc_ctx.register(b"my-metadata")
+        mock_reg.assert_called_once_with(
+            "localhost:9999", "localhost", port, b"my-metadata"
+        )
+
+    self.assertFalse(disc_ctx._server.is_started())
+
+  def test_local_process_context(self):
+    args = argparse.Namespace(
+        discovery_port=portpicker.pick_unused_port(),
+        discovery_addrs="leader:9999",
+    )
+    proc_ctx = local_context.LocalProcessContext(args)
+
+    with proc_ctx as entered:
+      self.assertIs(entered, proc_ctx)
+      self.assertIsInstance(proc_ctx.jax, context.JaxContext)
+      self.assertIsInstance(proc_ctx.ipc, local_context.LocalIpcContext)
+      self.assertIsInstance(
+          proc_ctx.ipc.discovery, local_context.LocalDiscoveryContext
+      )
+      self.assertTrue(proc_ctx.ipc.discovery._server.is_started() is False)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/distributed/runtime/discovery/discovery_test.py b/tests/experimental/distributed/runtime/discovery/discovery_test.py
new file mode 100644
index 000000000..81efc9106
--- /dev/null
+++ b/tests/experimental/distributed/runtime/discovery/discovery_test.py
@@ -0,0 +1,107 @@
+# 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.
+
+"""Unit tests for gRPC peer discovery server and registration helper."""
+
+import threading
+from unittest import mock
+
+from absl.testing import absltest
+import grpc
+import portpicker
+from tunix.experimental.distributed.runtime.discovery import discovery
+
+
+class DiscoveryTest(absltest.TestCase):
+
+  def test_start_with_zero_port_raises(self):
+    server = discovery.DiscoveryServer()
+    with self.assertRaises(ValueError):
+      server.start(0, lambda h, p, m: None)
+
+  @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)
+    try:
+      with self.assertRaises(RuntimeError):
+        server.start(port, lambda h, p, m: None)
+    finally:
+      server.stop()
+
+  @mock.patch.object(grpc, "server")
+  def test_server_register_rpc(self, mock_grpc_server):
+    server = discovery.DiscoveryServer()
+    port = 8888
+    received = {}
+
+    def callback(hostname, p, metadata):
+      received["hostname"] = hostname
+      received["port"] = p
+      received["metadata"] = metadata
+
+    server.start(port, callback)
+    try:
+      self.assertTrue(server.is_started())
+      mock_grpc_server.return_value.add_insecure_port.assert_called_once_with(
+          f"[::]:{port}"
+      )
+      mock_grpc_server.return_value.start.assert_called_once()
+    finally:
+      server.stop()
+
+  def test_register_empty_address_raises(self):
+    with self.assertRaises(ValueError):
+      discovery.register("", "node-0", 1234, b"meta")
+
+  @mock.patch.object(grpc, "insecure_channel")
+  @mock.patch("time.sleep")
+  def test_register_retry_on_unavailable(self, mock_sleep, mock_channel):
+    mock_stub_cls = mock.MagicMock()
+    mock_stub = mock_stub_cls.return_value
+
+    unavailable_error = grpc.RpcError()
+    unavailable_error.code = lambda: grpc.StatusCode.UNAVAILABLE
+    unavailable_error.details = lambda: "unavailable"
+
+    mock_stub.Register.side_effect = [unavailable_error, None]
+
+    with mock.patch(
+        "tunix.experimental.distributed.runtime.discovery.discovery_service_pb2_grpc.DiscoveryServiceStub",
+        return_value=mock_stub,
+    ):
+      discovery.register("localhost:9999", "node-0", 1234, b"meta")
+
+    self.assertEqual(mock_stub.Register.call_count, 2)
+    mock_sleep.assert_called_once()
+
+  @mock.patch.object(grpc, "insecure_channel")
+  def test_register_non_retryable_error_raises(self, mock_channel):
+    mock_stub = mock.MagicMock()
+    invalid_error = grpc.RpcError()
+    invalid_error.code = lambda: grpc.StatusCode.INVALID_ARGUMENT
+    invalid_error.details = lambda: "invalid arg"
+    mock_stub.Register.side_effect = invalid_error
+
+    with mock.patch(
+        "tunix.experimental.distributed.runtime.discovery.discovery_service_pb2_grpc.DiscoveryServiceStub",
+        return_value=mock_stub,
+    ):
+      with self.assertRaises(RuntimeError):
+        discovery.register("localhost:9999", "node-0", 1234, b"meta")
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/distributed/runtime/executors/executor_test.py b/tests/experimental/distributed/runtime/executors/executor_test.py
new file mode 100644
index 000000000..56b0cc2ff
--- /dev/null
+++ b/tests/experimental/distributed/runtime/executors/executor_test.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.
+
+"""Unit tests for local and Kubernetes process executors."""
+
+import argparse
+from unittest import mock
+
+from absl.testing import absltest
+import portpicker
+from tunix.experimental.distributed.runtime.contexts import k8s_context
+from tunix.experimental.distributed.runtime.contexts import local_context
+from tunix.experimental.distributed.runtime.executors import k8s_executor
+from tunix.experimental.distributed.runtime.executors import local_executor
+
+
+class ExecutorTest(absltest.TestCase):
+
+  def test_local_executor_run(self):
+    executor = local_executor.LocalExecutor()
+    args = argparse.Namespace(
+        discovery_port=portpicker.pick_unused_port(),
+        discovery_addrs="door:1234",
+    )
+
+    received = {}
+
+    def main_fn(argv, ctx):
+      received["argv"] = argv
+      received["ctx"] = ctx
+
+    executor.run(main_fn, ["--foo=bar"], args)
+    self.assertEqual(received["argv"], ["--foo=bar"])
+    self.assertIsInstance(received["ctx"], local_context.LocalProcessContext)
+
+  def test_k8s_executor_run(self):
+    executor = k8s_executor.K8sExecutor()
+    args = argparse.Namespace(
+        discovery_port=portpicker.pick_unused_port(),
+        discovery_addrs="door:1234",
+    )
+
+    received = {}
+
+    def main_fn(argv, ctx):
+      received["argv"] = argv
+      received["ctx"] = ctx
+
+    executor.run(main_fn, ["--bar=baz"], args)
+    self.assertEqual(received["argv"], ["--bar=baz"])
+    self.assertIsInstance(received["ctx"], k8s_context.K8sProcessContext)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/distributed/runtime/main_test.py b/tests/experimental/distributed/runtime/main_test.py
new file mode 100644
index 000000000..f0370a65c
--- /dev/null
+++ b/tests/experimental/distributed/runtime/main_test.py
@@ -0,0 +1,131 @@
+# 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.
+
+"""Unit tests for distributed runtime main entry point."""
+
+from unittest import mock
+from absl.testing import absltest
+from tunix.experimental.distributed.runtime import main as distributed_main
+
+
+def dummy_process_main(argv, context):
+  pass
+
+
+class MainTest(absltest.TestCase):
+
+  def test_import_symbol_valid(self):
+    fn = distributed_main.import_symbol(
+        "tunix.experimental.distributed.runtime.main.split_process_argv"
+    )
+    self.assertIs(fn, distributed_main.split_process_argv)
+
+  def test_import_symbol_invalid_no_dot(self):
+    with self.assertRaises(ValueError):
+      distributed_main.import_symbol("nodotsymbol")
+
+  def test_import_symbol_module_not_found(self):
+    with self.assertRaises(ModuleNotFoundError):
+      distributed_main.import_symbol(
+          "tunix.experimental.distributed.nonexistent.module.symbol"
+      )
+
+  def test_import_symbol_attribute_error(self):
+    with self.assertRaises(AttributeError):
+      distributed_main.import_symbol(
+          "tunix.experimental.distributed.runtime.main.NonexistentSymbol"
+      )
+
+  def test_split_process_argv_no_process_flag(self):
+    argv = ["--foo=bar", "--baz=123"]
+    slices = distributed_main.split_process_argv(argv)
+    self.assertEqual(slices, [["--foo=bar", "--baz=123"]])
+
+  def test_split_process_argv_multiple_processes(self):
+    argv = [
+        "--process",
+        "--process_main=foo.main",
+        "--port=123",
+        "--process",
+        "--process_main=bar.main",
+        "--port=456",
+    ]
+    slices = distributed_main.split_process_argv(argv)
+    self.assertEqual(
+        slices,
+        [
+            ["--process_main=foo.main", "--port=123"],
+            ["--process_main=bar.main", "--port=456"],
+        ],
+    )
+
+  def test_prepare_process_success(self):
+    argv = [
+        f"--process_main={__name__}.dummy_process_main",
+        "--discovery_id=worker-0",
+        "--discovery_port=8080",
+        "--discovery_addrs=leader:8080",
+        "--custom_flag=hello",
+    ]
+    prepared = distributed_main.prepare_process(argv)
+    self.assertEqual(prepared.main_fn.__name__, dummy_process_main.__name__)
+    self.assertEqual(prepared.context_args.discovery_id, "worker-0")
+    self.assertEqual(prepared.context_args.discovery_port, 8080)
+    self.assertEqual(prepared.context_args.discovery_addrs, "leader:8080")
+    self.assertEqual(prepared.argv, ["--custom_flag=hello"])
+
+  def test_prepare_process_invalid_main(self):
+    argv = [f"--process_main={__name__}.nonexistent_fn"]
+    with self.assertRaises(ValueError):
+      distributed_main.prepare_process(argv)
+
+  @mock.patch.object(distributed_main, "import_symbol")
+  def test_main_single_process(self, mock_import_symbol):
+    mock_executor = mock.MagicMock()
+    mock_import_symbol.side_effect = lambda fqn: (
+        mock_executor if "Executor" in fqn else dummy_process_main
+    )
+
+    argv = [
+        "prog",
+        "--process_executor=tunix.experimental.distributed.runtime.executor.LocalExecutor",
+        f"--process_main={__name__}.dummy_process_main",
+        "--custom_arg=abc",
+    ]
+    distributed_main.main(argv)
+    mock_executor.return_value.run.assert_called_once()
+
+  @mock.patch.object(distributed_main, "import_symbol")
+  def test_main_multi_process(self, mock_import_symbol):
+    mock_executor = mock.MagicMock()
+    mock_import_symbol.side_effect = lambda fqn: (
+        mock_executor if "Executor" in fqn else dummy_process_main
+    )
+
+    argv = [
+        "prog",
+        "--process_executor=tunix.experimental.distributed.runtime.executor.LocalExecutor",
+        "--process",
+        f"--process_main={__name__}.dummy_process_main",
+        "--p=1",
+        "--process",
+        f"--process_main={__name__}.dummy_process_main",
+        "--p=2",
+    ]
+    distributed_main.main(argv)
+    self.assertEqual(mock_executor.return_value.run.call_count, 2)
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tunix/experimental/distributed/examples/vllm_rollout/orchestrator.py b/tunix/experimental/distributed/examples/vllm_rollout/orchestrator.py
index c05a72117..1fc8229e7 100644
--- a/tunix/experimental/distributed/examples/vllm_rollout/orchestrator.py
+++ b/tunix/experimental/distributed/examples/vllm_rollout/orchestrator.py
@@ -25,8 +25,8 @@
 import time
 from typing import Any, Sequence
 
+from tunix.experimental.common import datatypes as data_types
 from tunix.experimental.distributed.runtime.context import ProcessContext
-from tunix.experimental.rollout import data_types
 from tunix.experimental.worker import remote_execution
 
 
diff --git a/tunix/experimental/distributed/runtime/contexts/k8s_context.py b/tunix/experimental/distributed/runtime/contexts/k8s_context.py
index 823641594..35c961e4c 100644
--- a/tunix/experimental/distributed/runtime/contexts/k8s_context.py
+++ b/tunix/experimental/distributed/runtime/contexts/k8s_context.py
@@ -77,7 +77,7 @@ class K8sJaxContext(context.JaxContext):
 
   def initialize(self) -> None:
     """Initializes Pathways or standard JAX distributed runtime based on environment."""
-    if "proxy" in os.environ.get("JAX_PLATFORMS") and os.environ.get(
+    if "proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get(
         "JAX_BACKEND_TARGET"
     ):
       logging.info("initializing Pathways runtime")
diff --git a/tunix/experimental/distributed/runtime/discovery/discovery.py b/tunix/experimental/distributed/runtime/discovery/discovery.py
index 80dca0ae9..243fc15f4 100644
--- a/tunix/experimental/distributed/runtime/discovery/discovery.py
+++ b/tunix/experimental/distributed/runtime/discovery/discovery.py
@@ -79,6 +79,7 @@ def stop(self, timeout: float | None = None) -> None:
     if self._server:
       self._server.stop(timeout)
       self._server.wait_for_termination(timeout)
+      self._server = None
 
 
 def register(