From de948daf82c3b22747385a2421a40a65163a6700 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Fri, 17 Jul 2026 09:17:09 -0700 Subject: [PATCH] Replace Shared Pathways Service YAML templates with PathwaysJobSet. PiperOrigin-RevId: 949606101 --- .../deploy_pathways_service.py | 101 ++++----- .../yamls/pw-service.yaml | 195 ------------------ 2 files changed, 51 insertions(+), 245 deletions(-) delete mode 100644 pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml diff --git a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py index bfd6979..72903ee 100644 --- a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py +++ b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py @@ -5,12 +5,12 @@ import logging import math import os -import string from typing import Any from absl import app from absl import flags from kubernetes import client from kubernetes import config +from pathwaysutils.experimental.gke import jobset import yaml _logger = logging.getLogger(__name__) @@ -45,13 +45,6 @@ "gs://pathways-test-bucket", "GCS bucket name for scratch space", ) -_TEMPLATE_FILE = flags.DEFINE_string( - "template_file", - os.path.join( - os.path.dirname(__file__), "yamls/pw-service.yaml", - ), - "Path to the JobSet YAML template file", -) _DRY_RUN = flags.DEFINE_boolean( "dry_run", False, @@ -149,25 +142,6 @@ def calculate_vms_per_slice(topology: str, chips_per_vm: int) -> int: ) from e -def load_and_substitute_template( - template_path: str, context: dict[str, Any] -) -> dict[str, Any]: - """Loads and substitutes the string.Template from the given path.""" - try: - with open(template_path, "r") as f: - template_str = f.read() - except OSError as err: - raise ValueError( - f"Could not read template file: {template_path}: {err}" - ) from err - - _logger.info("Template file: %s", template_path) - _logger.info("Context: %s", context) - template = string.Template(template_str) - _logger.info("Template: %s", template) - substituted_yaml = template.substitute(context) - return yaml.safe_load(substituted_yaml) - def deploy_jobset(jobset_yaml: dict[str, Any]) -> None: """Deploys the JobSet to the current Kubernetes cluster.""" @@ -198,29 +172,61 @@ def run_deployment( gcs_bucket, server_image, sidecar_image, - template_file, dry_run, deploy_func: Callable[[dict[str, Any]], None] = deploy_jobset, ) -> None: """Executes the deployment logic.""" - tpu_config = get_tpu_config(tpu_type) - vms_per_slice = calculate_vms_per_slice(topology, tpu_config.chips_per_vm) - - context = { - "JOBSET_NAME": jobset_name, - "SERVER_IMAGE": server_image, - "SIDECAR_IMAGE": sidecar_image, - "SIDECAR_SHM_DIR": _SIDECAR_SHM_DIR, - "GCS_SCRATCH_LOCATION": gcs_bucket, - "NUM_SLICES": num_slices, - "INSTANCE_TYPE": f"{tpu_config.instance_prefix}:{topology}", - "VMS_PER_SLICE": vms_per_slice, - "CHIPS_PER_VM": tpu_config.chips_per_vm, - "ACCELERATOR_LABEL": tpu_config.accelerator_label, - "TOPOLOGY": topology, - } + # Use PathwaysJobSet builder instead of YAML template. + pw_jobset = jobset.PathwaysJobSet( + name=jobset_name, + namespace="default", + pathways_dir=gcs_bucket, + tpu_type=tpu_type, + topology=topology, + num_slices=num_slices, + shared_pathways_service=True, + ) - jobset_config = load_and_substitute_template(template_file, context) + # If custom server_image is provided, mutate the templates to use it. + if server_image: + # Mutate head job. + for container in pw_jobset.head_job_template.spec.template.spec.containers: + if container.name == "pathways-rm": + container.image = server_image + # Mutate worker job. + for container in pw_jobset.worker_job_template.spec.template.spec.containers: + if container.name == "pathways-worker": + container.image = server_image + + # Add colocated python sidecar. + pw_jobset.add_colocated_python(image=sidecar_image, shm_mount_path=_SIDECAR_SHM_DIR) + + # Mutate the sidecar configuration to match what HEAD expects. + worker_spec = pw_jobset.worker_job_template.spec.template.spec + + # 1. Add extra logging env vars to sidecar. + for container in worker_spec.init_containers: + if container.name == "colocated-python-sidecar": + container.env.extend([ + client.V1EnvVar(name="PYTHONUNBUFFERED", value="1"), + client.V1EnvVar(name="LOGLEVEL", value="DEBUG"), + client.V1EnvVar(name="GLOG_minloglevel", value="0"), + client.V1EnvVar(name="GLOG_v", value="5"), + client.V1EnvVar(name="TF_CPP_MIN_LOG_LEVEL", value="0"), + client.V1EnvVar(name="TF_CPP_MIN_VLOG_LEVEL", value="5"), + client.V1EnvVar(name="TPU_MIN_LOG_LEVEL", value="0"), + client.V1EnvVar(name="GLOG_vmodule", value="jax_array_handlers=5,type_handlers=5,tensorstore_utils=5"), + ]) + + # 2. Add arg to pathways-worker container (in addition to env var set by builder). + for container in worker_spec.containers: + if container.name == "pathways-worker": + args = container.args or [] + if not any(a.startswith("--cloud_pathways_sidecar_shm_directory=") for a in args): + args.append(f"--cloud_pathways_sidecar_shm_directory={_SIDECAR_SHM_DIR}") + container.args = args + + jobset_config = pw_jobset.to_dict() _logger.info("--- Generated JobSet YAML ---") _logger.info("\n%s", yaml.dump(jobset_config)) @@ -256,15 +262,10 @@ def main(argv: Sequence[str]) -> None: gcs_bucket=_GCS_BUCKET.value, server_image=server_image, sidecar_image=_SIDECAR_IMAGE.value, - template_file=_TEMPLATE_FILE.value, dry_run=_DRY_RUN.value, ) except ValueError as e: _logger.exception("Error: %s", e) - except FileNotFoundError: - _logger.exception( - "Error: Template file not found at %s", _TEMPLATE_FILE.value - ) if __name__ == "__main__": diff --git a/pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml b/pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml deleted file mode 100644 index a02750e..0000000 --- a/pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml +++ /dev/null @@ -1,195 +0,0 @@ -apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: ${JOBSET_NAME} - namespace: default -spec: - coordinator: - replicatedJob: pathways-head - failurePolicy: - maxRestarts: 1 - restartStrategy: Recreate - network: - enableDNSHostnames: true - publishNotReadyAddresses: true - replicatedJobs: - - name: pathways-head - replicas: 1 - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname - spec: - backoffLimit: 3 - completionMode: Indexed - completions: 1 - parallelism: 1 - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname - spec: - containers: - - name: pathways-rm - image: ${SERVER_IMAGE} - imagePullPolicy: Always - args: - - --server_port=29001 - - --gcs_scratch_location=${GCS_SCRATCH_LOCATION} - - --node_type=resource_manager - - --instance_count=${NUM_SLICES} - - --instance_type=${INSTANCE_TYPE} - env: - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: HOST_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: TPU_SKIP_MDS_QUERY - value: "true" - ports: - - containerPort: 29001 - protocol: TCP - - containerPort: 29002 - protocol: TCP - resources: - limits: - cpu: "8" - memory: 32G - dnsPolicy: ClusterFirstWithHostNet - hostNetwork: true - restartPolicy: OnFailure - - name: worker - replicas: ${NUM_SLICES} - template: - spec: - backoffLimit: 1000000 - completionMode: Indexed - completions: ${VMS_PER_SLICE} - parallelism: ${VMS_PER_SLICE} - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool - spec: - containers: - - name: pathways-worker - image: ${SERVER_IMAGE} - imagePullPolicy: Always - args: - - --server_port=29005 - - --resource_manager_address=$$(PATHWAYS_HEAD):29001 - - --gcs_scratch_location=${GCS_SCRATCH_LOCATION} - - --cloud_pathways_sidecar_shm_directory=${SIDECAR_SHM_DIR} - env: - - name: TPU_MIN_LOG_LEVEL - value: "0" - - name: TF_CPP_MIN_LOG_LEVEL - value: "0" - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER - value: "false" - - name: MEGASCALE_NUM_SLICES - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: MEGASCALE_SLICE_ID - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: MEGASCALE_COORDINATOR_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - ports: - - containerPort: 29005 - protocol: TCP - - containerPort: 29006 - protocol: TCP - - containerPort: 8471 - protocol: TCP - - containerPort: 8080 - protocol: TCP - resources: - limits: - google.com/tpu: "${CHIPS_PER_VM}" - volumeMounts: - - name: shared-tmp - mountPath: /tmp - - name: sidecar-shared-memory - mountPath: ${SIDECAR_SHM_DIR} - initContainers: - - name: colocated-python-sidecar - image: ${SIDECAR_IMAGE} - imagePullPolicy: Always - env: - - name: GRPC_SERVER_ADDRESS - value: '''0.0.0.0:50051''' - - name: CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY - value: ${SIDECAR_SHM_DIR} - - name: PYTHONUNBUFFERED - value: '1' - # --- High Verbosity Logging Variables --- - - name: LOGLEVEL - value: 'DEBUG' - - name: GLOG_minloglevel - value: '0' # 0 = INFO level base - - name: GLOG_v - value: '5' # Extreme verbosity for all C++ modules - - name: TF_CPP_MIN_LOG_LEVEL - value: '0' - - name: TF_CPP_MIN_VLOG_LEVEL - value: '5' # TF/XLA verbose logging - - name: TPU_MIN_LOG_LEVEL - value: '0' - - name: GLOG_vmodule - value: 'jax_array_handlers=5,type_handlers=5,tensorstore_utils=5' - # ---------------------------------------- - ports: - - containerPort: 50051 - protocol: TCP - resources: {} - restartPolicy: Always - volumeMounts: - - name: shared-tmp - mountPath: /tmp - - name: sidecar-shared-memory - mountPath: ${SIDECAR_SHM_DIR} - dnsPolicy: ClusterFirstWithHostNet - hostNetwork: true - nodeSelector: - cloud.google.com/gke-tpu-accelerator: ${ACCELERATOR_LABEL} - cloud.google.com/gke-tpu-topology: ${TOPOLOGY} - restartPolicy: OnFailure - volumes: - - name: shared-tmp - hostPath: - path: /tmp - type: DirectoryOrCreate - - name: sidecar-shared-memory - emptyDir: - medium: Memory - startupPolicy: - startupPolicyOrder: InOrder - successPolicy: - operator: All