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/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