Skip to content

Expose tenant worker configurations in workload runner CLI #275

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Expose tenant worker configurations in workload runner CLI #275

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
240 changes: 111 additions & 129 deletions pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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"
)
Expand All @@ -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",
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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.")

Expand All @@ -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
)


Expand Down
Loading
Loading