diff --git a/pathwaysutils/experimental/gke/README.md b/pathwaysutils/experimental/gke/README.md new file mode 100644 index 0000000..45cc6b0 --- /dev/null +++ b/pathwaysutils/experimental/gke/README.md @@ -0,0 +1,223 @@ +# PathwaysJobSet Reference Implementation + +`PathwaysJobSet` is a client-side Python library designed to generate and +deploy Kubernetes `JobSet` resources (`kind: JobSet`) for running Pathways +workloads on Google Kubernetes Engine (GKE). + +It serves as a direct replacement for the custom `PathwaysJob` CRD (`kind: +PathwaysJob`), allowing you to deploy Pathways workloads using standard, +community-supported `JobSet` resources without needing a custom controller +installed on your GKE cluster. + +--- + +## Purpose & Overview + +This directory serves as the official **reference implementation** for +launching Pathways workloads via Kubernetes `JobSet`: + +- **Official Reference Standard:** Other Google teams and external customers + can base their JobSet specifications on this reference implementation and + expect their workloads to work reliably with Pathways. +- **Internal Integration Testing:** Our team uses this reference + implementation as the foundation for all of our automated integration + tests. +- **Minimal Required Specification:** This library defines the baseline, + minimal required Kubernetes specification to run Pathways on Cloud + workloads. + +### Production Workloads & Cluster Toolkit +While this reference implementation focuses on the minimal required +specification for Pathways, full-featured production deployments often +require additional Google Cloud Platform (GCP) features and infrastructure +integrations. + +- **Bridging the Gap with Cluster Toolkit:** Solutions such as + [Cluster Toolkit](https://cloud.google.com/cluster-toolkit/docs/overview) + bridge the gap between this minimal required implementation and + full-featured, enterprise-grade production workloads. +- **Guidance for Custom Workloads:** Advanced users and customers who need + custom JobSets or desire fine-grained control over their workload + definitions can refer to this `pathwaysutils` JobSet generator to see our + team's recommended best practices for Pathways JobSets, as well as reference + implementations for integrating with other GCP services (such as Cloud + Storage FUSE). + +--- + +## Key Features + +- **Client-Side Generation:** Generates standard `JobSet` YAML manifests + that can be inspected, version-controlled, or applied manually. +- **Standard Headless Execution:** Runs the Pathways Resource Manager (RM) + and Proxy as standalone containers in the head job. +- **GCSFuse Integration:** Easily mount Cloud Storage buckets to head + and/or worker pods using GCSFuse. +- **Colocated Python Support:** Inject a colocated Python sidecar container + and shared memory volume into worker pods for multi-agent or hybrid + workloads. +- **Elasticity Support:** Configure elastic slices for dynamic scaling. + +--- + +## Basic Usage + +```python +from pathwaysutils.experimental.gke import jobset + +# 1. Initialize the builder +pw_jobset = jobset.PathwaysJobSet( + name="my-pathways-workload", + namespace="default", + pathways_dir="gs://my-bucket/pathways-scratch", + tpu_type="v5e", + topology="4x8", + num_slices=1, +) + +# 2. Export to a standard JobSet YAML file +pw_jobset.export_yaml("jobset.yaml") + +# 3. Deploy directly to a GKE cluster (requires kubernetes configured) +pw_jobset.apply( + project_id="my-gcp-project", + region="us-central1", + cluster_id="my-gke-cluster", +) +``` + +--- + +## Examples + +### 1. Single-Slice v5e-32 (Non-elastic) +A standard single-slice workload on TPU v5e-32 (32 chips, `4x8` topology, 8 +VMs). + +```python +from pathwaysutils.experimental.gke import jobset + +js = jobset.PathwaysJobSet( + name="v5e-32-headless", + namespace="default", + pathways_dir="gs://my-bucket/scratch", + tpu_type="v5e", + topology="4x8", + num_slices=1, +) +js.export_yaml("v5e_32_headless.yaml") +``` + +### 2. Multislice v5p-4x4x4 (2 Slices) +A multislice workload using 2 slices of TPU v5p-4x4x4 (64 chips per slice). + +```python +from pathwaysutils.experimental.gke import jobset + +js = jobset.PathwaysJobSet( + name="v5p-multislice", + namespace="default", + pathways_dir="gs://my-bucket/scratch", + tpu_type="v5p", + topology="4x4x4", + num_slices=2, # 2 slices +) +js.export_yaml("v5p_multislice.yaml") +``` + +### 3. Multislice v6e with Elasticity +A multislice workload on TPU v6e-16 (topology `4x4`, 16 chips, 4 VMs per +slice) with 2 active slices initially, and elasticity configured for up to 4 +slices. + +```python +from pathwaysutils.experimental.gke import jobset + +js = jobset.PathwaysJobSet( + name="v6e-elastic", + namespace="default", + pathways_dir="gs://my-bucket/scratch", + tpu_type="v6e", + topology="4x4", + num_slices=2, # Initial/Active slices + elastic_slices=4, # Enable elasticity (informs proxy of max slices) +) +js.export_yaml("v6e_elastic.yaml") +``` + +### 4. Advanced Features: GCSFuse and Colocated Python +This example demonstrates chaining advanced features: + +- Mounting a GCS bucket to all containers (head and workers) using GCSFuse. +- Adding a colocated Python sidecar to the worker pods. + +```python +from pathwaysutils.experimental.gke import jobset + +js = ( + jobset.PathwaysJobSet( + name="advanced-workload", + namespace="default", + pathways_dir="gs://my-bucket/scratch", + tpu_type="v5e", + topology="4x8", + num_slices=1, + ) + .add_colocated_python() # Injects colocated python sidecar to workers + .add_gcsfuse( + containers="all", + mount_path="/data", + bucket="my-data-bucket", + read_only=True, + ) +) +js.export_yaml("advanced_workload.yaml") +``` + +--- + +## API Reference: `PathwaysJobSet` + +### Constructor `__init__` + +```python +def __init__( + self, + name: str, + namespace: str, + pathways_dir: str, + tpu_type: str, + topology: str, + num_slices: int, + max_restarts: int = 0, + max_slice_restarts: int = 0, + termination_grace_period_seconds: int | None = None, + pathways_version: str = "latest", + jobset_api_version: str = "v1alpha2", + elastic_slices: int = 0, + labels: Mapping[str, str] | None = None, + annotations: Mapping[str, str] | None = None, + shared_pathways_service: bool = False, +) +``` + +- `tpu_type`: Supported values include `"v5e"`, `"v5p"`, `"v6e"`, `"v4"`. +- `elastic_slices`: If `> 0`, enables elasticity by passing + `--num_elastic_slices` to the Pathways Proxy. +- `shared_pathways_service`: If `True`, configures the head job to run only + the Resource Manager (`pathways-rm`). + +### Methods + +- **`add_colocated_python()`**: Injects a colocated Python container and + shared memory volume (`/tmp/shared-memory`) into worker pods. +- **`add_gcsfuse(containers, mount_path, bucket, read_only=False)`**: Mounts + a GCS bucket using GCSFuse CSI driver. `containers` can be `"head"`, + `"worker"`, or `"all"`. +- **`to_dict()`**: Returns the compiled K8s JobSet resource as a Python + dictionary. +- **`export_yaml(filepath)`**: Serializes and writes the JobSet to a YAML + file. +- **`apply(project_id, region, cluster_id)`**: Deploys the JobSet to the + specified GKE cluster. Supports delete-and-recreate lifecycle if the JobSet + already exists. diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 6a69af1..6d321a7 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -11,17 +11,27 @@ # limitations under the License. """Pathways JobSet generator and builder (with Worker Job Config).""" +from __future__ import annotations + import hashlib import json import logging import math import time -from typing import Any, Mapping, Sequence -from kubernetes import client -from kubernetes import config as k8s_config +from typing import TYPE_CHECKING, Any, Mapping, Sequence import yaml -# GKE sidecar containers restartPolicy compatibility placeholder. +try: + import kubernetes +except ImportError as e: + raise ImportError( + "GKE utilities require `kubernetes`. " + "Please install pathwaysutils with GKE support:\n\n" + " pip install 'pathwaysutils[gke]'\n" + ) from e + +from kubernetes import client +from kubernetes import config as k8s_config _logger = logging.getLogger(__name__) @@ -72,7 +82,7 @@ def _deserialize_dict( - api_client: client.ApiClient, data_dict: Mapping[str, Any], klass: Any + api_client: Any, data_dict: Mapping[str, Any], klass: Any ) -> Any: class FakeResponse: @@ -93,8 +103,6 @@ def __init__( tpu_type: str, topology: str, num_slices: int, - user_pod_template: Mapping[str, Any] | None = None, - main_container_name: str = "main", max_restarts: int = 0, max_slice_restarts: int = 0, termination_grace_period_seconds: int | None = None, @@ -114,8 +122,6 @@ def __init__( tpu_type: TPU type (e.g., "v5e"). topology: TPU topology (e.g., "2x2"). num_slices: Number of slices. - user_pod_template: Optional user pod template for the head job. - main_container_name: Name of the main container in user_pod_template. max_restarts: Maximum number of restarts for the JobSet. max_slice_restarts: Maximum number of slice restarts. termination_grace_period_seconds: Optional termination grace period. @@ -124,12 +130,8 @@ def __init__( elastic_slices: Number of elastic slices. labels: Optional labels for the JobSet. annotations: Optional annotations for the JobSet. + shared_pathways_service: Whether to run only RM for Shared Pathways Service. """ - if shared_pathways_service and user_pod_template: - raise ValueError( - "Cannot enable shared_pathways_service when user_pod_template is" - " provided." - ) self._shared_pathways_service = shared_pathways_service self._name = name @@ -166,8 +168,6 @@ def __init__( num_slices=num_slices, instance_type=instance_type, image_tag=image_tag, - user_pod_template=user_pod_template, - main_container_name=main_container_name, elastic_slices=elastic_slices, shared_pathways_service=shared_pathways_service, ) @@ -185,7 +185,7 @@ def __init__( ) self._success_policy = None - if user_pod_template or shared_pathways_service: + if shared_pathways_service: self._success_policy = { "operator": "All", "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME], @@ -205,8 +205,6 @@ def _build_head_job_template( num_slices: int, instance_type: str, image_tag: str, - user_pod_template: Mapping[str, Any] | None, - main_container_name: str, elastic_slices: int, shared_pathways_service: bool, ) -> client.V1JobTemplateSpec: @@ -217,9 +215,8 @@ def _build_head_job_template( num_slices: Number of slices. instance_type: TPU instance type (e.g., "tpuv5:2x2"). image_tag: Version tag for Pathways images. - user_pod_template: Optional user pod template for the head job. - main_container_name: Name of the main container in user_pod_template. elastic_slices: Number of elastic slices. + shared_pathways_service: Whether to run only RM for Shared Pathways Service. Returns: The head job template. @@ -318,74 +315,20 @@ def _build_head_job_template( ), ) - api_client = client.ApiClient() + containers = [rm_container] + if not shared_pathways_service: + containers.append(proxy_container) - if user_pod_template: - user_template_obj = _deserialize_dict( - api_client, user_pod_template, client.V1PodTemplateSpec - ) - head_pod_spec = user_template_obj.spec - head_pod_spec.host_network = True - head_pod_spec.dns_policy = "ClusterFirstWithHostNet" - - rm_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - proxy_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - - init_containers = head_pod_spec.init_containers or [] - init_containers.extend([rm_container, proxy_container]) - head_pod_spec.init_containers = init_containers - - # Inject JAX env vars into main container. - jax_env = [ - client.V1EnvVar( - name="PATHWAYS_HEAD", - value_from=client.V1EnvVarSource( - field_ref=client.V1ObjectFieldSelector( - field_path=( - "metadata.labels['jobset.sigs.k8s.io/coordinator']" - ) - ) - ), - ), - client.V1EnvVar(name="JAX_PLATFORMS", value="proxy"), - client.V1EnvVar(name="XCLOUD_ENVIRONMENT", value="GCP"), - client.V1EnvVar( - name="JAX_BACKEND_TARGET", - value=f"grpc://$(PATHWAYS_HEAD):{PATHWAYS_PROXY_PORT}", - ), - ] - containers = head_pod_spec.containers or [] - for c in containers: - if c.name == main_container_name: - env = c.env or [] - env.extend(jax_env) - c.env = env - break - head_pod_spec.containers = containers - - annotations = user_pod_template.get("metadata", {}).get("annotations", {}) - labels = user_pod_template.get("metadata", {}).get("labels", {}) - else: - # Headless mode. - containers = [rm_container] - if not shared_pathways_service: - containers.append(proxy_container) - head_pod_spec = client.V1PodSpec( - host_network=True, - dns_policy="ClusterFirstWithHostNet", - containers=containers, - ) - annotations = {} - labels = {} - - if not head_pod_spec.restart_policy: - head_pod_spec.restart_policy = "Never" + head_pod_spec = client.V1PodSpec( + host_network=True, + dns_policy="ClusterFirstWithHostNet", + containers=containers, + restart_policy="Never", + ) - # Default annotations job_annotations = { "alpha.jobset.sigs.k8s.io/exclusive-topology": "kubernetes.io/hostname" } - job_annotations.update(annotations) head_job_template = client.V1JobTemplateSpec( metadata=client.V1ObjectMeta(annotations=job_annotations), @@ -396,7 +339,7 @@ def _build_head_job_template( parallelism=1, template=client.V1PodTemplateSpec( metadata=client.V1ObjectMeta( - annotations=job_annotations, labels=labels + annotations=job_annotations, labels={} ), spec=head_pod_spec, ), @@ -640,11 +583,10 @@ def add_colocated_python( client.V1VolumeMount(name=shm_volume_name, mount_path=shm_mount_path), ], ) - colocated_container.restart_policy = "Always" # pyrefly: ignore[missing-attribute] - init_containers = pod_spec.init_containers or [] - init_containers.append(colocated_container) - pod_spec.init_containers = init_containers + containers = pod_spec.containers or [] + containers.append(colocated_container) + pod_spec.containers = containers # Add volume mount to pathways-worker. for container in pod_spec.containers: diff --git a/pathwaysutils/proxy_backend.py b/pathwaysutils/proxy_backend.py index cf1f806..2507a8a 100644 --- a/pathwaysutils/proxy_backend.py +++ b/pathwaysutils/proxy_backend.py @@ -13,17 +13,27 @@ # limitations under the License. """Register the IFRT Proxy as a backend for JAX.""" +import os import jax from jax.extend import backend from jax.extend.backend import ifrt_proxy def register_backend_factory() -> None: + """Registers the IFRT Proxy backend factory with JAX.""" + + def make_client(): + options = ifrt_proxy.ClientConnectionOptions() + timeout_secs = os.environ.get("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS") + if timeout_secs: + options.connection_timeout_in_seconds = int(timeout_secs) + return ifrt_proxy.get_client( + jax.config.read("jax_backend_target"), + options, + ) + backend.register_backend_factory( "proxy", - lambda: ifrt_proxy.get_client( - jax.config.read("jax_backend_target"), - ifrt_proxy.ClientConnectionOptions(), - ), + make_client, priority=-1, ) diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index 6401bdd..082e9b4 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -2,12 +2,22 @@ import itertools import os from typing import Any +import unittest from unittest import mock from absl.testing import absltest from absl.testing import parameterized +import yaml + +try: + import kubernetes +except ImportError: + raise unittest.SkipTest( + "kubernetes is not installed (requires pathwaysutils[test])" + ) + from kubernetes import client +from kubernetes import config as k8s_config from pathwaysutils.experimental.gke import jobset -import yaml def normalize_k8s_spec(spec: Any) -> Any: @@ -167,98 +177,6 @@ def test_headless_head_job_containers(self): ) self.assertIn("--num_elastic_slices=2", proxy_container["args"]) - def test_non_headless_head_job_init_containers(self): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "ubuntu:latest", - }] - } - } - js = self._create_jobset( - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - config = js.to_dict() - helper = JobSetManifestHelper(config) - - pod_spec = helper.pod_specs["pathways-head"] - self.assertLen(pod_spec["initContainers"], 2) - self.assertIn("pathways-rm", helper.init_containers["pathways-head"]) - self.assertIn("pathways-proxy", helper.init_containers["pathways-head"]) - - rm_container = helper.init_containers["pathways-head"]["pathways-rm"] - proxy_container = helper.init_containers["pathways-head"]["pathways-proxy"] - self.assertEqual(rm_container["restartPolicy"], "Always") - self.assertEqual(proxy_container["restartPolicy"], "Always") - - def test_non_headless_head_job_jax_env(self): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "ubuntu:latest", - }] - } - } - js = self._create_jobset( - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - config = js.to_dict() - helper = JobSetManifestHelper(config) - - self.assertIn("jax-tpu", helper.containers["pathways-head"]) - main_container = helper.containers["pathways-head"]["jax-tpu"] - env_names = [e["name"] for e in main_container["env"]] - self.assertIn("PATHWAYS_HEAD", env_names) - self.assertIn("JAX_PLATFORMS", env_names) - self.assertIn("XCLOUD_ENVIRONMENT", env_names) - self.assertIn("JAX_BACKEND_TARGET", env_names) - - def test_non_headless_head_job_annotations(self): - user_pod_template = { - "metadata": {"annotations": {"example.com/annotation": "value"}}, - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "ubuntu:latest", - }] - } - } - js = self._create_jobset( - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - config = js.to_dict() - helper = JobSetManifestHelper(config) - - head_job = helper.jobs["pathways-head"] - self.assertEqual( - head_job["template"]["metadata"]["annotations"][ - "example.com/annotation" - ], - "value", - ) - self.assertEqual( - head_job["template"]["spec"]["template"]["metadata"]["annotations"][ - "example.com/annotation" - ], - "value", - ) - - def test_monkeypatch_restart_policy(self): - # Construct V1Container with restart_policy to test monkeypatch. - c = client.V1Container( - name="test", - restart_policy="Always" # pyrefly: ignore[unexpected-keyword] - ) # pytype: disable=wrong-keyword-args - self.assertEqual(getattr(c, "restart_policy"), "Always") - def test_worker_job_replicas(self): js = self._create_jobset(num_slices=2) @@ -358,7 +276,7 @@ def test_add_gcsfuse_read_only(self, read_only): ("worker", "pathways-worker", ["pathways-worker"]), ("explicit", ["pathways-worker", "pathways-rm"], ["pathways-worker", "pathways-rm"]), ) - def test_add_gcsfuse_container_filtering_headless( + def test_add_gcsfuse_container_filtering( self, containers_param, expected_containers ): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) @@ -383,49 +301,6 @@ def test_add_gcsfuse_container_filtering_headless( else: self.assertFalse(has_mount, f"Expected {c_name} in {job_name} NOT to have mount") - @parameterized.named_parameters( - ("all", "all", ["pathways-rm", "pathways-proxy", "pathways-worker", "jax-tpu"]), - ("explicit", ["pathways-worker", "jax-tpu"], ["pathways-worker", "jax-tpu"]), - ("explicit_rm", ["pathways-worker", "pathways-rm"], ["pathways-worker", "pathways-rm"]), - ) - def test_add_gcsfuse_container_filtering_non_headless( - self, containers_param, expected_containers - ): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "gcr.io/my-project/jax-tpu:latest", - }] - } - } - pw_jobset = self._create_jobset( - topology="2x2", - num_slices=1, - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - bucket_hash = int(hashlib.md5("my-bucket".encode()).hexdigest(), 16) % (10**8) - expected_vol_name = f"gcsfuse-{bucket_hash}" - - pw_jobset.add_gcsfuse( - containers=containers_param, - mount_path="/gcs/data", - bucket="my-bucket", - ) - helper = JobSetManifestHelper(pw_jobset.to_dict()) - - all_possible_containers = ["pathways-rm", "pathways-proxy", "pathways-worker", "jax-tpu"] - for c_name in all_possible_containers: - matches = helper.get_all_containers_by_name(c_name) - self.assertNotEmpty(matches) - for job_name, container in matches: - has_mount = any(m["mountPath"] == "/gcs/data" for m in container.get("volumeMounts", [])) - if c_name in expected_containers: - self.assertTrue(has_mount, f"Expected {c_name} in {job_name} to have mount") - else: - self.assertFalse(has_mount, f"Expected {c_name} in {job_name} NOT to have mount") - def test_add_gcsfuse_volumes_and_annotations(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) bucket_hash = int(hashlib.md5("my-bucket".encode()).hexdigest(), 16) % (10**8) @@ -544,9 +419,8 @@ def test_add_colocated_python_sidecar(self): pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom") helper = JobSetManifestHelper(pw_jobset.to_dict()) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) - sidecar = helper.init_containers["pathways-worker"]["colocated-python-sidecar"] - self.assertEqual(sidecar["restartPolicy"], "Always") + self.assertIn("colocated-python-sidecar", helper.containers["pathways-worker"]) + sidecar = helper.containers["pathways-worker"]["colocated-python-sidecar"] self.assertEqual(sidecar["image"], "gcr.io/my-project/colocated-python:custom") self.assertTrue( any( @@ -575,7 +449,7 @@ def test_add_colocated_python_preserves_init_containers(self): # Verify both exist self.assertIn("existing-init-container", helper.init_containers["pathways-worker"]) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) + self.assertIn("colocated-python-sidecar", helper.containers["pathways-worker"]) def test_add_colocated_python_volume_default(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) @@ -619,8 +493,8 @@ def test_add_colocated_python_custom_shm(self): ) helper = JobSetManifestHelper(pw_jobset.to_dict()) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) - sidecar = helper.init_containers["pathways-worker"]["colocated-python-sidecar"] + self.assertIn("colocated-python-sidecar", helper.containers["pathways-worker"]) + sidecar = helper.containers["pathways-worker"]["colocated-python-sidecar"] self.assertTrue( any( m["name"] == "shared-memory" and m["mountPath"] == "/tmp/custom-shm" @@ -655,35 +529,6 @@ def test_add_colocated_python_custom_shm(self): ) ) - def test_colocated_python_with_jax_command(self): - jax_command = "import jax; print(jax.devices());" - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "gcr.io/my-project/jax-tpu:latest", - "command": ["python3", "-c", jax_command], - }] - } - } - pw_jobset = self._create_jobset( - topology="2x2", - num_slices=1, - user_pod_template=user_pod_template, - main_container_name="jax-tpu", - ) - - pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom") - helper = JobSetManifestHelper(pw_jobset.to_dict()) - - self.assertIn("pathways-head", helper.jobs) - self.assertIn("jax-tpu", helper.containers["pathways-head"]) - jax_container = helper.containers["pathways-head"]["jax-tpu"] - self.assertEqual(jax_container["command"], ["python3", "-c", jax_command]) - - self.assertIn("pathways-worker", helper.jobs) - self.assertIn("colocated-python-sidecar", helper.init_containers["pathways-worker"]) - # Reference similar GKE / JobSet custom object unit test suites: # - Google3 GKE JobSet test suite: //depot/google3/cloud/ai/map/catmint/supervisor/client/python/orchestrators/gke_callbacks_test.py # - Upstream Kubernetes SIGs JobSet unit tests: https://github.com/kubernetes-sigs/jobset/blob/main/pkg/controllers/jobset_controller_test.go @@ -949,29 +794,6 @@ def test_shared_pathways_service(self): self.assertNotIn("pathways-proxy", helper.containers["pathways-head"]) self.assertLen(pod_spec["containers"], 1) - def test_shared_pathways_service_with_user_template_fails(self): - user_pod_template = { - "spec": { - "containers": [{ - "name": "jax-tpu", - "image": "gcr.io/my-project/jax-tpu:latest", - }] - } - } - with self.assertRaisesRegex( - ValueError, - "Cannot enable shared_pathways_service when user_pod_template is" - " provided.", - ): - jobset.PathwaysJobSet( - name="test-sps", - namespace="default", - pathways_dir="gs://bucket/scratch", - tpu_type="v5e", - topology="2x2", - num_slices=2, - shared_pathways_service=True, - user_pod_template=user_pod_template, - ) + if __name__ == "__main__": absltest.main() diff --git a/pathwaysutils/test/proxy_backend_test.py b/pathwaysutils/test/proxy_backend_test.py index fb8ad8c..13fce79 100644 --- a/pathwaysutils/test/proxy_backend_test.py +++ b/pathwaysutils/test/proxy_backend_test.py @@ -13,6 +13,7 @@ # limitations under the License. """Tests for the proxy backend module.""" +import os from unittest import mock from absl.testing import absltest @@ -54,6 +55,46 @@ def test_proxy_backend_registration(self): proxy_backend.register_backend_factory() self.assertIn("proxy", backend.backends()) + def test_proxy_backend_registration_with_timeout(self): + mock_get_client = self.enter_context( + mock.patch.object( + ifrt_proxy, + "get_client", + return_value=mock.MagicMock(), + ) + ) + self.enter_context( + mock.patch.dict( + os.environ, {"PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS": "42"} + ) + ) + proxy_backend.register_backend_factory() + self.assertIn("proxy", backend.backends()) + mock_get_client.assert_called_once() + args, _ = mock_get_client.call_args + self.assertEqual(args[0], "grpc://localhost:12345") + options = args[1] + self.assertEqual(options.connection_timeout_in_seconds, 42) + + def test_proxy_backend_registration_without_timeout(self): + mock_get_client = self.enter_context( + mock.patch.object( + ifrt_proxy, + "get_client", + return_value=mock.MagicMock(), + ) + ) + self.enter_context(mock.patch.dict(os.environ)) + os.environ.pop("PATHWAYS_PROXY_CONNECTION_TIMEOUT_SECS", None) + + proxy_backend.register_backend_factory() + self.assertIn("proxy", backend.backends()) + mock_get_client.assert_called_once() + args, _ = mock_get_client.call_args + self.assertEqual(args[0], "grpc://localhost:12345") + options = args[1] + self.assertNotEqual(options.connection_timeout_in_seconds, 42) + if __name__ == "__main__": absltest.main()