From e4122fd78679931568565382d3bd967a9a28cdf6 Mon Sep 17 00:00:00 2001 From: Akanksha Gupta Date: Wed, 1 Jul 2026 10:58:41 -0700 Subject: [PATCH] Dynamically deploy tenant workers with custom sidecar image PiperOrigin-RevId: 941189177 --- .../deploy_pathways_service.py | 240 ++++++++---------- .../shared_pathways_service/gke_utils.py | 92 +++++++ .../shared_pathways_service/isc_pathways.py | 106 +++++++- .../shared_pathways_service/tpu_specs.py | 109 ++++++++ .../yamls/kueue-sps.yaml | 32 +++ .../yamls/shared-rm.yaml | 52 ++++ .../{pw-service.yaml => tenant-worker.yaml} | 90 +------ 7 files changed, 512 insertions(+), 209 deletions(-) create mode 100644 pathwaysutils/experimental/shared_pathways_service/tpu_specs.py create mode 100644 pathwaysutils/experimental/shared_pathways_service/yamls/kueue-sps.yaml create mode 100644 pathwaysutils/experimental/shared_pathways_service/yamls/shared-rm.yaml rename pathwaysutils/experimental/shared_pathways_service/yamls/{pw-service.yaml => tenant-worker.yaml} (58%) diff --git a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py index bfd6979..60c8e64 100644 --- a/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py +++ b/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py @@ -1,36 +1,29 @@ """Deploys Pathways service to a Kubernetes cluster using a JobSet template.""" from collections.abc import Callable, Sequence -import dataclasses import logging import math import os import string +import subprocess from typing import Any from absl import app from absl import flags -from kubernetes import client -from kubernetes import config -import yaml +from pathwaysutils.experimental.shared_pathways_service import tpu_specs _logger = logging.getLogger(__name__) # Flag definitions FLAGS = flags.FLAGS -_JOBSET_NAME = flags.DEFINE_string( - "jobset_name", "pathways-service", "Name of the JobSet" +_DEPLOYMENT_NAME = flags.DEFINE_string( + "deployment_name", None, "Deployment name for the Pathways service" ) _JAX_VERSION = flags.DEFINE_string( - "jax_version", "0.9.0", "JAX version (e.g., 0.9.0)" + "jax_version", "0.10.1", "JAX version (e.g., 0.10.1)" ) _SERVER_IMAGE = flags.DEFINE_string( "server_image", None, "Full path to the server Docker image" ) -_SIDECAR_IMAGE = flags.DEFINE_string( - "sidecar_image", - "us-docker.pkg.dev/cloud-tpu-v2-images/pathways-colocated-python/sidecar:20260423-python_3.12-jax_0.10.0", - "Full path to the sidecar Docker image", -) _TPU_TYPE = flags.DEFINE_enum( "tpu_type", "v6e", ["v5e", "v5p", "v6e", "tpu7x"], "TPU type" ) @@ -43,14 +36,39 @@ _GCS_BUCKET = flags.DEFINE_string( "gcs_bucket", "gs://pathways-test-bucket", - "GCS bucket name for scratch space", + "GCS bucket name for RM checkpoint location", +) +_RM_TEMPLATE_FILE = flags.DEFINE_string( + "rm_template_file", + os.path.join( + os.path.dirname(__file__), + "yamls/shared-rm.yaml", + ), + "Path to the Shared RM YAML template file", +) +_KUEUE_TEMPLATE_FILE = flags.DEFINE_string( + "kueue_template_file", + os.path.join( + os.path.dirname(__file__), + "yamls/kueue-sps.yaml", + ), + "Path to the Kueue SPS YAML template file", +) +_NUM_PREDEPLOYED_WORKERS = flags.DEFINE_integer( + "num_predeployed_workers", 0, "Number of worker JobSets to deploy upfront." +) +_SIDECAR_IMAGE = flags.DEFINE_string( + "sidecar_image", + "us-docker.pkg.dev/cloud-tpu-v2-images/pathways-colocated-python/sidecar:20260423-python_3.12-jax_0.10.0", + "Full path to the sidecar Docker image for workers.", ) -_TEMPLATE_FILE = flags.DEFINE_string( - "template_file", +_WORKER_TEMPLATE_FILE = flags.DEFINE_string( + "worker_template_file", os.path.join( - os.path.dirname(__file__), "yamls/pw-service.yaml", + os.path.dirname(__file__), + "yamls/tenant-worker.yaml", ), - "Path to the JobSet YAML template file", + "Path to the worker JobSet YAML template file", ) _DRY_RUN = flags.DEFINE_boolean( "dry_run", @@ -60,15 +78,6 @@ _SIDECAR_SHM_DIR = "/tmp/sidecar_dir" -@dataclasses.dataclass(frozen=True) -class TPUConfig: - """Holds configuration details for a specific TPU type.""" - machine_type: str - chips_per_vm: int - accelerator_label: str - instance_prefix: str - - def _validate_topology(topology): """Validates the topology flag format.""" try: @@ -95,63 +104,9 @@ def _validate_topology(topology): ) -def get_tpu_config(tpu_type: str) -> TPUConfig: - """Returns a TPUConfig object containing TPU configuration details.""" - tpu_configs = { - "v5e": TPUConfig( - machine_type="ct5lp-hightpu-4t", - chips_per_vm=4, - accelerator_label="tpu-v5-lite-podslice", - instance_prefix="tpuv5e", - ), - "v5p": TPUConfig( - machine_type="ct5p-hightpu-4t", - chips_per_vm=4, - accelerator_label="tpu-v5p-slice", - instance_prefix="tpuv5", - ), - "v6e": TPUConfig( - machine_type="ct6e-standard-4t", - chips_per_vm=4, - accelerator_label="tpu-v6e-slice", - instance_prefix="tpuv6e", - ), - "tpu7x": TPUConfig( - machine_type="tpu7x-standard-4t", - chips_per_vm=4, - accelerator_label="tpu7x", - instance_prefix="tpu7x", - ), - } - if tpu_type not in tpu_configs: - raise ValueError( - f"Unsupported TPU type: {tpu_type}. Supported types are:" - f" {list(tpu_configs.keys())}" - ) - return tpu_configs[tpu_type] - - -def calculate_vms_per_slice(topology: str, chips_per_vm: int) -> int: - """Calculates the number of VMs per slice based on the topology.""" - try: - dims = [int(d) for d in topology.split("x")] - total_chips = math.prod(dims) - if total_chips % chips_per_vm != 0: - raise ValueError( - f"Total chips ({total_chips}) in topology {topology} is not divisible" - f" by chips_per_vm ({chips_per_vm})" - ) - return total_chips // chips_per_vm - except ValueError as e: - raise ValueError( - f"Invalid topology format: {topology}. Expected format like 'AxB' or" - f" 'AxBxC'. {e}" - ) from e - - def load_and_substitute_template( template_path: str, context: dict[str, Any] -) -> dict[str, Any]: +) -> str: """Loads and substitutes the string.Template from the given path.""" try: with open(template_path, "r") as f: @@ -164,70 +119,87 @@ def load_and_substitute_template( _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) + return template.substitute(context) -def deploy_jobset(jobset_yaml: dict[str, Any]) -> None: - """Deploys the JobSet to the current Kubernetes cluster.""" +def deploy_yaml_str(yaml_str: str) -> None: + """Deploys the given YAML string to the GKE cluster using kubectl.""" try: - config.load_kube_config() - api = client.CustomObjectsApi() - api.create_namespaced_custom_object( - group="jobset.x-k8s.io", - version="v1alpha2", - namespace=jobset_yaml["metadata"]["namespace"], - body=jobset_yaml, - plural="jobsets", - ) - _logger.info( - "JobSet '%s' created successfully.", jobset_yaml["metadata"]["name"] + subprocess.run( + ["kubectl", "apply", "-f", "-"], + input=yaml_str, + check=True, + text=True, + capture_output=True, ) - except client.rest.ApiException: - _logger.exception("Error creating JobSet") - except config.ConfigException: - _logger.exception("Error loading Kubernetes configuration") + _logger.info("Successfully applied YAML.") + except subprocess.CalledProcessError as e: + _logger.exception("Error applying YAML: %s", e.stderr) + raise def run_deployment( - tpu_type, - topology, - num_slices, - jobset_name, - gcs_bucket, - server_image, - sidecar_image, - template_file, - dry_run, - deploy_func: Callable[[dict[str, Any]], None] = deploy_jobset, + deployment_name: str, + tpu_type: str, + topology: str, + num_slices: int, + gcs_bucket: str, + server_image: str, + rm_template_file: str, + kueue_template_file: str, + worker_template_file: str, + num_predeployed_workers: int, + sidecar_image: str, + dry_run: bool, + deploy_func: Callable[[str], None] = deploy_yaml_str, ) -> 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) + dims = [int(d) for d in topology.split("x")] + chips_per_slice = math.prod(dims) + + tpu_params = tpu_specs.get_tpu_params(tpu_type, topology) context = { - "JOBSET_NAME": jobset_name, + "DEPLOYMENT_NAME": deployment_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, + "TPU_TYPE": tpu_type, + "NOMINAL_QUOTA": str(chips_per_slice * num_slices), + **tpu_params, } - jobset_config = load_and_substitute_template(template_file, context) + rm_config_str = load_and_substitute_template(rm_template_file, context) + kueue_config_str = load_and_substitute_template(kueue_template_file, context) - _logger.info("--- Generated JobSet YAML ---") - _logger.info("\n%s", yaml.dump(jobset_config)) + _logger.info("--- Generated RM YAML ---") + _logger.info("\n%s", rm_config_str) + _logger.info("--- Generated Kueue YAML ---") + _logger.info("\n%s", kueue_config_str) if not dry_run: - _logger.info("Deploying JobSet...") - deploy_func(jobset_config) + _logger.info("Deploying RM...") + deploy_func(rm_config_str) + _logger.info("Deploying Kueue Configs...") + deploy_func(kueue_config_str) + + if num_predeployed_workers > 0: + for i in range(1, num_predeployed_workers + 1): + worker_context = context.copy() + worker_context.update({ + "WORKER_NAME": f"{deployment_name}-w{i}", + "QUEUE_NAME": f"{deployment_name}-shared-tpu-local-q", + "SIDECAR_IMAGE": sidecar_image, + "SIDECAR_SHM_DIR": _SIDECAR_SHM_DIR, + "PATHWAYS_RM_SERVICE_ADDRESS": f"{deployment_name}-rm-svc:29001", + }) + worker_config_str = load_and_substitute_template( + worker_template_file, worker_context + ) + _logger.info("--- Generated Worker %d YAML ---", i) + _logger.info("\n%s", worker_config_str) + _logger.info("Deploying Worker %d...", i) + deploy_func(worker_config_str) else: _logger.info("Dry run mode, not deploying.") @@ -248,22 +220,32 @@ def main(argv: Sequence[str]) -> None: else: server_image = f"us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-{_JAX_VERSION.value}" + if _DEPLOYMENT_NAME.value: + # Take the first 8 characters to limit the length. GKE Pod name has a max + # length limit of 63 characters. + deployment_name = _DEPLOYMENT_NAME.value[:8] + else: + deployment_name = f"pw-{_TPU_TYPE.value}" + run_deployment( + deployment_name=deployment_name, tpu_type=_TPU_TYPE.value, topology=_TOPOLOGY.value, num_slices=_NUM_SLICES.value, - jobset_name=_JOBSET_NAME.value, gcs_bucket=_GCS_BUCKET.value, server_image=server_image, + rm_template_file=_RM_TEMPLATE_FILE.value, + kueue_template_file=_KUEUE_TEMPLATE_FILE.value, + worker_template_file=_WORKER_TEMPLATE_FILE.value, + num_predeployed_workers=_NUM_PREDEPLOYED_WORKERS.value, 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 + "Error: Template file not found at %s", _RM_TEMPLATE_FILE.value ) diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index ac17acf..1fd0766 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py @@ -531,3 +531,95 @@ def get_worker_sidecar_image( return None +def get_rm_server_image( + pathways_service: str, namespace: str = "default" +) -> str | None: + """Gets the server image used by the Pathways Resource Manager Service.""" + pathways_head_hostname = pathways_service.split(":")[0] + _validate_k8s_name(namespace) + _validate_k8s_name(pathways_head_hostname) + + # 1. Get the Service to find its selector + get_service_cmd = [ + "kubectl", + "get", + "service", + pathways_head_hostname, + "-n", + namespace, + "-o", + "json", + ] + try: + result = subprocess.run( + get_service_cmd, + check=True, + capture_output=True, + text=True, + ) + service_data = json.loads(result.stdout) + selector = service_data.get("spec", {}).get("selector") + if not selector: + _logger.warning( + "No selector found for service %s in namespace %s", + pathways_head_hostname, + namespace, + ) + return None + except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e: + _logger.exception( + "Failed to get selector for service %s: %r", + pathways_head_hostname, + e, + ) + return None + + # 2. Format the label selector + label_selector = ",".join(f"{k}={v}" for k, v in selector.items()) + + # 3. Get Pods matching the selector + get_pods_cmd = [ + "kubectl", + "get", + "pods", + "-n", + namespace, + "-l", + label_selector, + "-o", + "json", + ] + try: + result = subprocess.run( + get_pods_cmd, + check=True, + capture_output=True, + text=True, + ) + pods_data = json.loads(result.stdout) + items = pods_data.get("items", []) + if not items: + _logger.warning( + "No pods found matching selector %s in namespace %s", + label_selector, + namespace, + ) + return None + + # Get the image from the first container of the first matching pod + pod_spec = items[0].get("spec", {}) + containers = pod_spec.get("containers", []) + if containers: + image = containers[0].get("image") + if image: + return image + + except (subprocess.CalledProcessError, json.JSONDecodeError, KeyError) as e: + _logger.exception( + "Failed to get pods or parse image for selector %s: %r", + label_selector, + e, + ) + return None + + return None diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index 339b1f3..f44651b 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py @@ -18,12 +18,16 @@ import pathwaysutils from pathwaysutils.experimental.shared_pathways_service import gke_utils from pathwaysutils.experimental.shared_pathways_service import metrics_collector +from pathwaysutils.experimental.shared_pathways_service import tpu_specs from pathwaysutils.experimental.shared_pathways_service import validators PROXY_FILEPATH = os.path.join( os.path.dirname(__file__), "yamls/pw-proxy.yaml" ) +WORKER_FILEPATH = os.path.join( + os.path.dirname(__file__), "yamls/tenant-worker.yaml" +) # TODO(b/459935429): Hardcoding the port and using hostNetwork: true in the # proxy YAML limits us to one proxy server pod per node. Consider alternative # networking configurations to allow multiple proxies per node if needed. @@ -36,6 +40,9 @@ DEFAULT_PROXY_IMAGE = ( "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest" ) +DEFAULT_SERVER_IMAGE = ( + "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest" +) _logger = logging.getLogger(__name__) @@ -90,6 +97,38 @@ def from_list(cls, options: Iterable[str] | None) -> "ProxyOptions": ) +def _deploy_pathways_workers( + *, + worker_name: str, + queue_name: str, + server_image: str, + sidecar_image: str, + gcs_scratch_location: str, + pathways_service: str, + tpu_params: dict[str, str], +) -> None: + """Deploys the Pathways worker JobSet to the GKE cluster.""" + try: + with open(WORKER_FILEPATH, "r") as f: + yaml_template = f.read() + except OSError as err: + raise ValueError("Could not read file: " + WORKER_FILEPATH) from err + + template = string.Template(yaml_template) + substituted_yaml = template.substitute( + WORKER_NAME=worker_name, + QUEUE_NAME=queue_name, + SERVER_IMAGE=server_image, + SIDECAR_IMAGE=sidecar_image, + GCS_SCRATCH_LOCATION=gcs_scratch_location, + SIDECAR_SHM_DIR="/tmp/sidecar_dir", + PATHWAYS_RM_SERVICE_ADDRESS=pathways_service, + **tpu_params, + ) + _logger.info("Deploying Tenant Workers: %s", worker_name) + gke_utils.deploy_gke_yaml(substituted_yaml) + + def _deploy_pathways_proxy_server( *, pathways_service: str, @@ -244,10 +283,14 @@ class _ISCPathways: proxy_job_name: The name to use for the deployed proxy. proxy_pod_name: The name of the proxy pod, assigned during deployment. proxy_server_image: The image to use for the proxy server. + server_image: The server image to use for the workers. proxy_options: Configuration options for the Pathways proxy. metrics_collector: The metrics collector instance if enabled. start_time: The start time of the TPU assignment. total_chips: The total number of TPU chips expected across all instances. + sidecar_image: The custom sidecar image for the tenant worker. + queue_name: The name of the Kueue queue to use for the workers. + worker_name: The name of the worker JobSet. """ def __init__( @@ -263,6 +306,9 @@ def __init__( proxy_server_image: str, proxy_options: ProxyOptions | None = None, collect_service_metrics: bool = False, + sidecar_image: str | None = None, + queue_name: str = "shared-tpu-local-queue", + server_image: str = DEFAULT_SERVER_IMAGE, ): """Initializes the TPU manager.""" self.cluster = cluster @@ -276,7 +322,13 @@ def __init__( self._port_forward_process = None self._proxy_port = None self.proxy_server_image = proxy_server_image + self.server_image = server_image self.proxy_options = proxy_options or ProxyOptions() + self.sidecar_image = sidecar_image + self.queue_name = queue_name + self.worker_name = ( + f"isc-worker-{os.environ.get('USER', 'user')}-{self._proxy_job_name.split('-')[-1]}" + ) self._old_jax_platforms = None if collect_service_metrics: raw_collector = metrics_collector.MetricsCollector( @@ -301,7 +353,11 @@ def __repr__(self): f"pathways_service='{self.pathways_service}', " f"expected_tpu_instances={self.expected_tpu_instances}, " f"_proxy_job_name='{self._proxy_job_name}', " - f"proxy_options={self.proxy_options})" + f"proxy_server_image='{self.proxy_server_image}', " + f"server_image='{self.server_image}', " + f"proxy_options={self.proxy_options}, " + f"sidecar_image='{self.sidecar_image}', " + f"queue_name='{self.queue_name}')" ) def _get_total_chips(self) -> int: @@ -334,6 +390,20 @@ def __enter__(self): try: self.start_time = time.time() + if self.sidecar_image: + tpu_type_str = next(iter(self.expected_tpu_instances.keys())) + tpu_gen, topology = tpu_specs.parse_tpu_type_string(tpu_type_str) + tpu_params = tpu_specs.get_tpu_params(tpu_gen, topology) + _deploy_pathways_workers( + worker_name=self.worker_name, + queue_name=self.queue_name, + server_image=self.server_image, + sidecar_image=self.sidecar_image, + gcs_scratch_location=self.bucket, + pathways_service=self.pathways_service, + tpu_params=tpu_params, + ) + _deploy_pathways_proxy_server( pathways_service=self.pathways_service, proxy_job_name=self._proxy_job_name, @@ -409,6 +479,9 @@ def _cleanup(self) -> None: # Delete the proxy GKE job. _logger.info("Deleting Pathways proxy...") gke_utils.delete_gke_resource("job", self._proxy_job_name) + if self.sidecar_image: + _logger.info("Deleting Tenant Workers...") + gke_utils.delete_gke_resource("jobset", self.worker_name) _logger.info("Pathways proxy GKE job deletion complete.") # Restore JAX variables. @@ -437,6 +510,9 @@ def connect( proxy_server_image: str = DEFAULT_PROXY_IMAGE, proxy_options: Sequence[str] | None = None, collect_service_metrics: bool = False, + sidecar_image: str | None = None, + queue_name: str = "shared-tpu-local-queue", + server_image: str | None = None, ) -> Iterator["_ISCPathways"]: """Connects to a Pathways server if the cluster exists. If not, creates it. @@ -456,6 +532,10 @@ def connect( provided, no extra options will be used. collect_service_metrics: Whether to collect usage metrics for Shared Pathways Service. + sidecar_image: The custom sidecar image for the tenant worker. + queue_name: The Kueue LocalQueue name. + server_image: The server image to use for workers. If not provided, a + default will be used. Yields: The Pathways manager. @@ -469,11 +549,26 @@ def connect( cluster_name=cluster, project_id=project, location=region ) - proxy_options_obj = ProxyOptions.from_list(proxy_options) - if proxy_options_obj.sidecar: - sidecar_image = gke_utils.get_worker_sidecar_image( + if not server_image: + inferred_image = gke_utils.get_rm_server_image( pathways_service=pathways_service ) + if inferred_image: + _logger.info("Inferred RM server image: %s", inferred_image) + server_image = inferred_image + else: + _logger.warning( + "Could not infer RM server image, falling back to default: %s", + DEFAULT_SERVER_IMAGE, + ) + server_image = DEFAULT_SERVER_IMAGE + + proxy_options_obj = ProxyOptions.from_list(proxy_options) + if proxy_options_obj.sidecar: + if not sidecar_image: + sidecar_image = gke_utils.get_worker_sidecar_image( + pathways_service=pathways_service + ) if sidecar_image: validators.validate_sidecar_image_versions(sidecar_image) _logger.info("Validation complete.") @@ -496,6 +591,9 @@ def connect( proxy_server_image=proxy_server_image, proxy_options=proxy_options_obj, collect_service_metrics=collect_service_metrics, + sidecar_image=sidecar_image, + queue_name=queue_name, + server_image=server_image, ) as t: if t.proxy_pod_name: num_slices = sum(t.expected_tpu_instances.values()) diff --git a/pathwaysutils/experimental/shared_pathways_service/tpu_specs.py b/pathwaysutils/experimental/shared_pathways_service/tpu_specs.py new file mode 100644 index 0000000..0f0a6ca --- /dev/null +++ b/pathwaysutils/experimental/shared_pathways_service/tpu_specs.py @@ -0,0 +1,109 @@ +"""TPU specifications and helper functions for Shared Pathways Service.""" + +import dataclasses +import logging +import math + +_logger = logging.getLogger(__name__) + + +@dataclasses.dataclass(frozen=True) +class TPUConfig: + """Holds configuration details for a specific TPU type.""" + + machine_type: str + chips_per_vm: int + accelerator_label: str + instance_prefix: str + + +def get_tpu_config(tpu_type: str) -> TPUConfig: + """Returns a TPUConfig object containing TPU configuration details.""" + tpu_configs = { + "v5e": TPUConfig( + machine_type="ct5lp-hightpu-4t", + chips_per_vm=4, + accelerator_label="tpu-v5-lite-podslice", + instance_prefix="tpuv5e", + ), + "v5p": TPUConfig( + machine_type="ct5p-hightpu-4t", + chips_per_vm=4, + accelerator_label="tpu-v5p-slice", + instance_prefix="tpuv5", + ), + "v6e": TPUConfig( + machine_type="ct6e-standard-4t", + chips_per_vm=4, + accelerator_label="tpu-v6e-slice", + instance_prefix="tpuv6e", + ), + "tpu7x": TPUConfig( + machine_type="tpu7x-standard-4t", + chips_per_vm=4, + accelerator_label="tpu7x", + instance_prefix="tpu7x", + ), + } + if tpu_type not in tpu_configs: + raise ValueError( + f"Unsupported TPU type: {tpu_type}. Supported types are:" + f" {list(tpu_configs.keys())}" + ) + return tpu_configs[tpu_type] + + +def calculate_vms_per_slice(topology: str, chips_per_vm: int) -> int: + """Calculates the number of VMs per slice based on the topology.""" + try: + dims = [int(d) for d in topology.split("x")] + total_chips = math.prod(dims) + if total_chips % chips_per_vm != 0: + raise ValueError( + f"Total chips ({total_chips}) in topology {topology} is not divisible" + f" by chips_per_vm ({chips_per_vm})" + ) + return total_chips // chips_per_vm + except ValueError as e: + raise ValueError( + f"Invalid topology format: {topology}. Expected format like 'AxB' or" + f" 'AxBxC'. {e}" + ) from e + + +def parse_tpu_type_string(tpu_type_str: str) -> tuple[str, str]: + """Parses a tpu_type string (e.g., 'tpuv5e:4x8') into (tpu_type, topology).""" + if ":" not in tpu_type_str: + raise ValueError( + f"Invalid tpu_type string: {tpu_type_str}. Expected format" + " 'type:topology'" + ) + prefix, topology = tpu_type_str.split(":", 1) + + # Map prefix back to tpu_type key + prefix_map = { + "tpuv5e": "v5e", + "tpuv5": "v5p", + "tpuv6e": "v6e", + "tpu7x": "tpu7x", + } + if prefix not in prefix_map: + raise ValueError( + f"Unsupported TPU prefix: {prefix}. Supported prefixes are:" + f" {list(prefix_map.keys())}" + ) + + return prefix_map[prefix], topology + + +def get_tpu_params(tpu_type: str, topology: str) -> dict[str, str]: + """Returns a dictionary of TPU parameters for GKE templates.""" + tpu_config = get_tpu_config(tpu_type) + vms_per_slice = calculate_vms_per_slice(topology, tpu_config.chips_per_vm) + return { + "ACCELERATOR_LABEL": tpu_config.accelerator_label, + "TOPOLOGY": topology, + "VMS_PER_SLICE": str(vms_per_slice), + "CHIPS_PER_VM": str(tpu_config.chips_per_vm), + "INSTANCE_TYPE": f"{tpu_config.instance_prefix}:{topology}", + } diff --git a/pathwaysutils/experimental/shared_pathways_service/yamls/kueue-sps.yaml b/pathwaysutils/experimental/shared_pathways_service/yamls/kueue-sps.yaml new file mode 100644 index 0000000..97d07be --- /dev/null +++ b/pathwaysutils/experimental/shared_pathways_service/yamls/kueue-sps.yaml @@ -0,0 +1,32 @@ +apiVersion: kueue.x-k8s.io/v1beta1 +kind: ResourceFlavor +metadata: + name: ${DEPLOYMENT_NAME}-tpu-${TPU_TYPE} +spec: + nodeLabels: + cloud.google.com/gke-tpu-accelerator: ${ACCELERATOR_LABEL} + cloud.google.com/gke-tpu-topology: ${TOPOLOGY} +--- +apiVersion: kueue.x-k8s.io/v1beta1 +kind: ClusterQueue +metadata: + name: ${DEPLOYMENT_NAME}-shared-tpu-q +spec: + namespaceSelector: {} # Match all namespaces + cohort: tpu-pool + resourceGroups: + - coveredResources: ["google.com/tpu"] + flavors: + - name: tpu-${TPU_TYPE} + resources: + - name: "google.com/tpu" + nominalQuota: "${NOMINAL_QUOTA}" # 2 slices of 32 chips = 64 chips + +--- +apiVersion: kueue.x-k8s.io/v1beta1 +kind: LocalQueue +metadata: + name: ${DEPLOYMENT_NAME}-shared-tpu-local-q + namespace: default +spec: + clusterQueue: ${DEPLOYMENT_NAME}-shared-tpu-q diff --git a/pathwaysutils/experimental/shared_pathways_service/yamls/shared-rm.yaml b/pathwaysutils/experimental/shared_pathways_service/yamls/shared-rm.yaml new file mode 100644 index 0000000..2fab284 --- /dev/null +++ b/pathwaysutils/experimental/shared_pathways_service/yamls/shared-rm.yaml @@ -0,0 +1,52 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: ${DEPLOYMENT_NAME}-rm + namespace: default +spec: + replicas: 1 + selector: + matchLabels: + app: ${DEPLOYMENT_NAME}-rm + template: + metadata: + labels: + app: ${DEPLOYMENT_NAME}-rm + spec: + containers: + - name: ${DEPLOYMENT_NAME}-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} + ports: + - containerPort: 29001 + protocol: TCP + - containerPort: 29002 + protocol: TCP + resources: + limits: + cpu: "8" + memory: 32G + restartPolicy: Always +--- +apiVersion: v1 +kind: Service +metadata: + name: ${DEPLOYMENT_NAME}-rm-svc + namespace: default +spec: + selector: + app: ${DEPLOYMENT_NAME}-rm + ports: + - name: rm-port + port: 29001 + targetPort: 29001 + - name: rm-port-2 + port: 29002 + targetPort: 29002 + type: ClusterIP diff --git a/pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml b/pathwaysutils/experimental/shared_pathways_service/yamls/tenant-worker.yaml similarity index 58% rename from pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml rename to pathwaysutils/experimental/shared_pathways_service/yamls/tenant-worker.yaml index a02750e..07cca01 100644 --- a/pathwaysutils/experimental/shared_pathways_service/yamls/pw-service.yaml +++ b/pathwaysutils/experimental/shared_pathways_service/yamls/tenant-worker.yaml @@ -1,11 +1,12 @@ apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: - name: ${JOBSET_NAME} + name: ${WORKER_NAME} namespace: default + labels: + kueue.x-k8s.io/queue-name: ${QUEUE_NAME} spec: - coordinator: - replicatedJob: pathways-head + suspend: true failurePolicy: maxRestarts: 1 restartStrategy: Recreate @@ -13,61 +14,8 @@ spec: 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} + replicas: 1 # We run 1 slice per JobSet template: spec: backoffLimit: 1000000 @@ -85,7 +33,7 @@ spec: imagePullPolicy: Always args: - --server_port=29005 - - --resource_manager_address=$$(PATHWAYS_HEAD):29001 + - --resource_manager_address=${PATHWAYS_RM_SERVICE_ADDRESS} - --gcs_scratch_location=${GCS_SCRATCH_LOCATION} - --cloud_pathways_sidecar_shm_directory=${SIDECAR_SHM_DIR} env: @@ -98,9 +46,7 @@ spec: - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER value: "false" - name: MEGASCALE_NUM_SLICES - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + value: "1" # Each JobSet is 1 slice - name: JOBSET_NAME valueFrom: fieldRef: @@ -110,17 +56,11 @@ spec: 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'] + value: "0" # Only 1 slice in this JobSet - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + value: pathways-rm-service - name: MEGASCALE_COORDINATOR_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + value: pathways-rm-service ports: - containerPort: 29005 protocol: TCP @@ -144,27 +84,25 @@ spec: imagePullPolicy: Always env: - name: GRPC_SERVER_ADDRESS - value: '''0.0.0.0:50051''' + 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 + value: '0' - name: GLOG_v - value: '5' # Extreme verbosity for all C++ modules + value: '5' - name: TF_CPP_MIN_LOG_LEVEL value: '0' - name: TF_CPP_MIN_VLOG_LEVEL - value: '5' # TF/XLA verbose logging + value: '5' - 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