From 02ab62f074bcd4a79b57869be9bf41eca710174b Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Tue, 28 Jul 2026 08:23:42 -0700 Subject: [PATCH] Make pathways proxy image and rm/worker image configurable PiperOrigin-RevId: 955271890 --- pathwaysutils/experimental/gke/jobset.py | 67 +++++++++++++++++-- .../test/experimental/gke/jobset_test.py | 51 ++++++++++++++ 2 files changed, 113 insertions(+), 5 deletions(-) diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index cedee48..9cd5470 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -81,6 +81,16 @@ } +def _format_image(image: str, default_tag: str) -> str: + """Appends default_tag if image does not already specify a tag or digest.""" + if "@" in image: + return image + last_slash = image.rfind("/") + if ":" in image[last_slash + 1:]: + return image + return f"{image}:{default_tag}" + + def _deserialize_dict( api_client: Any, data_dict: Mapping[str, Any], klass: Any ) -> Any: @@ -112,6 +122,8 @@ def __init__( labels: Mapping[str, str] | None = None, annotations: Mapping[str, str] | None = None, shared_pathways_service: bool = False, + pathways_rm_and_worker_image: str = DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE, + pathways_proxy_image: str = DEFAULT_PATHWAYS_PROXY_IMAGE, ): """Initializes the instance. @@ -131,8 +143,13 @@ def __init__( labels: Optional labels for the JobSet. annotations: Optional annotations for the JobSet. shared_pathways_service: Whether to run only RM for Shared Pathways Service. + pathways_rm_and_worker_image: Base Docker image for Resource Manager and + Worker containers. + pathways_proxy_image: Base Docker image for Proxy container. """ self._shared_pathways_service = shared_pathways_service + self._pathways_rm_and_worker_image = pathways_rm_and_worker_image + self._pathways_proxy_image = pathways_proxy_image self._name = name self._namespace = namespace @@ -170,6 +187,8 @@ def __init__( image_tag=image_tag, elastic_slices=elastic_slices, shared_pathways_service=shared_pathways_service, + pathways_rm_and_worker_image=pathways_rm_and_worker_image, + pathways_proxy_image=pathways_proxy_image, ) # Build worker template. @@ -182,6 +201,7 @@ def __init__( image_tag=image_tag, max_slice_restarts=max_slice_restarts, termination_grace_period_seconds=termination_grace_period_seconds, + pathways_rm_and_worker_image=pathways_rm_and_worker_image, ) self._success_policy = None @@ -199,6 +219,14 @@ def head_job_template(self) -> client.V1JobTemplateSpec: def worker_job_template(self) -> client.V1JobTemplateSpec: return self._worker_job_template + @property + def pathways_rm_and_worker_image(self) -> str: + return self._pathways_rm_and_worker_image + + @property + def pathways_proxy_image(self) -> str: + return self._pathways_proxy_image + def _build_head_job_template( self, pathways_dir: str, @@ -207,6 +235,8 @@ def _build_head_job_template( image_tag: str, elastic_slices: int, shared_pathways_service: bool, + pathways_rm_and_worker_image: str, + pathways_proxy_image: str, ) -> client.V1JobTemplateSpec: """Builds the head job template for the JobSet. @@ -217,12 +247,14 @@ def _build_head_job_template( image_tag: Version tag for Pathways images. elastic_slices: Number of elastic slices. shared_pathways_service: Whether to run only RM for Shared Pathways Service. + pathways_rm_and_worker_image: Base Docker image for Resource Manager. + pathways_proxy_image: Base Docker image for Proxy container. Returns: The head job template. """ - rm_image = f"{DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE}:{image_tag}" - proxy_image = f"{DEFAULT_PATHWAYS_PROXY_IMAGE}:{image_tag}" + rm_image = _format_image(pathways_rm_and_worker_image, image_tag) + proxy_image = _format_image(pathways_proxy_image, image_tag) rm_args = [ f"--server_port={PATHWAYS_RM_PORT}", @@ -355,8 +387,10 @@ def _build_worker_job_template( image_tag: str, max_slice_restarts: int, termination_grace_period_seconds: int | None, + pathways_rm_and_worker_image: str, ) -> client.V1JobTemplateSpec: - worker_image = f"{DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE}:{image_tag}" + """Builds the worker job template for the JobSet.""" + worker_image = _format_image(pathways_rm_and_worker_image, image_tag) args = [ f"--resource_manager_address=$(PATHWAYS_HEAD):{PATHWAYS_RM_PORT}", @@ -574,7 +608,9 @@ def add_colocated_python( ports=[client.V1ContainerPort(container_port=50051)], volume_mounts=[ client.V1VolumeMount(name="shared-tmp", mount_path="/tmp"), - client.V1VolumeMount(name=shm_volume_name, mount_path=shm_mount_path), + client.V1VolumeMount( + name=shm_volume_name, mount_path=shm_mount_path + ), ], ) @@ -630,7 +666,9 @@ def add_gcsfuse( for job_template in (self._head_job_template, self._worker_job_template): pod_spec = job_template.spec.template.spec - all_containers = (pod_spec.containers or []) + (pod_spec.init_containers or []) + all_containers = (pod_spec.containers or []) + ( + pod_spec.init_containers or [] + ) matching = self._filter_matching_containers(containers, all_containers) if not matching: @@ -760,6 +798,25 @@ def import_yaml(cls, filepath: str) -> "PathwaysJobSet": instance._head_job_template = head_job_template instance._worker_job_template = worker_job_template + rm_image = DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE + proxy_image = DEFAULT_PATHWAYS_PROXY_IMAGE + if ( + head_job_template.spec + and head_job_template.spec.template + and head_job_template.spec.template.spec + ): + pod_spec = head_job_template.spec.template.spec + all_containers = (pod_spec.containers or []) + ( + pod_spec.init_containers or [] + ) + for c in all_containers: + if c.name == "pathways-rm" and c.image: + rm_image = c.image + elif c.name == "pathways-proxy" and c.image: + proxy_image = c.image + instance._pathways_rm_and_worker_image = rm_image + instance._pathways_proxy_image = proxy_image + instance._success_policy = config["spec"].get("successPolicy") return instance diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index cc2fea6..41fd382 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -654,6 +654,57 @@ def test_export_import_roundtrip(self): normalize_k8s_spec(imported_jobset.to_dict()), ) + def test_configurable_images(self): + pw_jobset = jobset.PathwaysJobSet( + name="test-workload", + namespace="default", + pathways_dir="gs://bucket/scratch", + tpu_type="v5e", + topology="2x2", + num_slices=1, + pathways_version="v2.0", + pathways_rm_and_worker_image="gcr.io/custom/server", + pathways_proxy_image="gcr.io/custom/proxy:tagged", + ) + + self.assertEqual( + pw_jobset.pathways_rm_and_worker_image, "gcr.io/custom/server" + ) + self.assertEqual( + pw_jobset.pathways_proxy_image, "gcr.io/custom/proxy:tagged" + ) + + # Verify head job containers. + head_containers = { + c.name: c.image + for c in pw_jobset.head_job_template.spec.template.spec.containers + } + self.assertEqual(head_containers["pathways-rm"], "gcr.io/custom/server:v2.0") + self.assertEqual( + head_containers["pathways-proxy"], "gcr.io/custom/proxy:tagged" + ) + + # Verify worker job container. + worker_image = ( + pw_jobset.worker_job_template.spec.template.spec.containers[0].image + ) + self.assertEqual(worker_image, "gcr.io/custom/server:v2.0") + + # Verify roundtrip import. + temp_filepath = os.path.join( + self.create_tempdir().full_path, "jobset_custom.yaml" + ) + pw_jobset.export_yaml(temp_filepath) + imported_jobset = jobset.PathwaysJobSet.import_yaml(temp_filepath) + + self.assertEqual( + imported_jobset.pathways_rm_and_worker_image, + "gcr.io/custom/server:v2.0", + ) + self.assertEqual( + imported_jobset.pathways_proxy_image, "gcr.io/custom/proxy:tagged" + ) + def test_import_validation_failures(self): temp_dir = self.create_tempdir().full_path