test: add unit tests to distributed modules

.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

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

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

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

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