diff --git a/pathwaysutils/experimental/gke/jobset.py b/pathwaysutils/experimental/gke/jobset.py index 28bee85..0626497 100644 --- a/pathwaysutils/experimental/gke/jobset.py +++ b/pathwaysutils/experimental/gke/jobset.py @@ -576,11 +576,11 @@ def _add_volume_to_pod_spec( def add_colocated_python( self, - image: str, + image: str | None = None, shm_mount_path: str = "/tmp/shared-memory", shm_size_limit: str | None = None, ) -> "PathwaysJobSet": - """Adds colocated python sidecar to the worker pods.""" + """Adds colocated python sidecar and/or shared-memory volume to worker pods.""" pod_spec = self._worker_job_template.spec.template.spec # Add shared memory volume if not exists. @@ -598,30 +598,33 @@ def add_colocated_python( ) pod_spec.volumes = volumes - # Add colocated python container. - colocated_container = client.V1Container( - name="colocated-python-sidecar", - image=image, - image_pull_policy="Always", - env=[ - client.V1EnvVar(name="GRPC_SERVER_ADDRESS", value="0.0.0.0:50051"), - client.V1EnvVar( - name="CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY", - value=shm_mount_path, - ), - ], - 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 - ), - ], - ) + if image: + # Add colocated python container. + colocated_container = client.V1Container( + name="colocated-python-sidecar", + image=image, + image_pull_policy="Always", + env=[ + client.V1EnvVar( + name="GRPC_SERVER_ADDRESS", value="0.0.0.0:50051" + ), + client.V1EnvVar( + name="CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY", + value=shm_mount_path, + ), + ], + 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 + ), + ], + ) - containers = pod_spec.containers or [] - containers.append(colocated_container) - pod_spec.containers = containers + containers = pod_spec.containers or [] + containers.append(colocated_container) + pod_spec.containers = containers # Add volume mount to pathways-worker. for container in pod_spec.containers: diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index 144c999..6b2df8e 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py @@ -1,5 +1,6 @@ """GKE utils for deploying and managing the Pathways proxy.""" +from collections.abc import Iterable import datetime import functools import logging @@ -109,15 +110,20 @@ def deploy_gke_yaml(yaml: str, action: str = "apply") -> None: def delete_gke_resource( - resource_type: str, name: str, namespace: str = "default" + resource_type: str, + name: str, + namespace: str = "default", + *, + wait: bool = True, ) -> None: """Deletes the given resource from the GKE cluster. Args: - resource_type: The type of resource to delete (e.g. "deployment", - "service", "job"). + resource_type: The type of resource to delete (e.g. "deployment", "service", + "job"). name: The name of the resource. namespace: The namespace of the resource. + wait: Whether to wait for the resource deletion to complete. """ _validate_k8s_name(resource_type) _validate_k8s_name(name) @@ -132,9 +138,13 @@ def delete_gke_resource( "-n", namespace, "--ignore-not-found", + ] + if not wait: + command.append("--wait=false") + command.extend([ "--", name, - ] + ]) try: result = subprocess.run( command, @@ -150,7 +160,6 @@ def delete_gke_resource( raise - def get_pod_from_job(job_name: str) -> str: """Returns the pod name for the given job. @@ -672,7 +681,6 @@ def stream_pod_logs(pod_name: str) -> subprocess.Popen[str]: raise - def wait_for_deployment( name: str, namespace: str = "default", timeout: int = 300 ) -> None: @@ -850,4 +858,204 @@ def get_compatible_proxy_server_image(server_image: str) -> str: return new_repo +def _read_worker_pod( + core_api: client.CoreV1Api, + pod_name: str, + namespace: str = "default", +) -> client.V1Pod: + """Reads a worker Pod by exact Pod name or by Indexed Job hostname.""" + try: + return core_api.read_namespaced_pod(name=pod_name, namespace=namespace) + except client.exceptions.ApiException as e: + if e.status != 404 or "-" not in pod_name: + raise + job_name, completion_index = pod_name.rsplit("-", 1) + if not completion_index.isdigit(): + raise + label_selector = ( + f"job-name={job_name}," + f"batch.kubernetes.io/job-completion-index={completion_index}" + ) + pod_list = core_api.list_namespaced_pod( + namespace=namespace, label_selector=label_selector + ) + items = getattr(pod_list, "items", None) or [] + active_pods = [ + p + for p in items + if getattr(getattr(p, "metadata", None), "deletion_timestamp", None) + is None + and getattr(getattr(p, "status", None), "phase", None) + in ("Pending", "Running") + ] + candidates = active_pods or items + if not candidates: + raise + return candidates[0] + + +def inject_ephemeral_sidecar( + *, + pod_name: str, + container_name: str = "ephemeral-python-sidecar", + image: str, + port: int = 50052, + shm_mount_path: str = "/tmp/sidecar_dir", + namespace: str = "default", +) -> None: + """Injects an ephemeral sidecar container into the given worker Pod.""" + _validate_k8s_name(pod_name) + _validate_k8s_name(container_name) + _validate_k8s_name(namespace) + + core_api = _get_k8s_core_api() + pod = _read_worker_pod(core_api, pod_name, namespace=namespace) + actual_pod_name = ( + pod.metadata.name + if getattr(pod, "metadata", None) + and isinstance(getattr(pod.metadata, "name", None), str) + else pod_name + ) + existing_ephemeral = pod.spec.ephemeral_containers or [] + if any(c.name == container_name for c in existing_ephemeral): + _logger.info( + "Ephemeral container '%s' already exists on pod '%s'. Skipping" + " injection.", + container_name, + actual_pod_name, + ) + return + + ephemeral_container = client.V1EphemeralContainer( + name=container_name, + image=image, + image_pull_policy="Always", + target_container_name="pathways-worker", + args=[ + f"--port={port}", + "--logtostderr", + "--stderrthreshold=0", + "--v=1", + ], + env=[ + client.V1EnvVar(name="GRPC_SERVER_ADDRESS", value=f"0.0.0.0:{port}"), + client.V1EnvVar( + name="CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY", + value=shm_mount_path, + ), + 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", + ), + ], + volume_mounts=[ + client.V1VolumeMount(name="shared-tmp", mount_path="/tmp"), + client.V1VolumeMount(name="shared-memory", mount_path=shm_mount_path), + ], + ) + patch_body = { + "spec": { + "ephemeralContainers": existing_ephemeral + [ephemeral_container] + } + } + _logger.info( + "Injecting ephemeral sidecar '%s' (image=%s, port=%d) into pod '%s'...", + container_name, + image, + port, + actual_pod_name, + ) + core_api.patch_namespaced_pod_ephemeralcontainers( + name=actual_pod_name, + namespace=namespace, + body=patch_body, + ) + + +def wait_for_ephemeral_container( + pod_name: str, + container_name: str = "ephemeral-python-sidecar", + namespace: str = "default", + timeout: int = 300, +) -> None: + """Waits for an ephemeral container in a Pod to enter the Running state.""" + _validate_k8s_name(pod_name) + _validate_k8s_name(container_name) + _validate_k8s_name(namespace) + core_api = _get_k8s_core_api() + start_time = time.time() + while time.time() - start_time < timeout: + pod = _read_worker_pod(core_api, pod_name, namespace=namespace) + actual_pod_name = ( + pod.metadata.name + if getattr(pod, "metadata", None) + and isinstance(getattr(pod.metadata, "name", None), str) + else pod_name + ) + statuses = pod.status.ephemeral_container_statuses or [] + for status in statuses: + if status.name == container_name: + if status.state and status.state.running is not None: + _logger.info( + "Ephemeral container '%s' is running in pod '%s'.", + container_name, + actual_pod_name, + ) + return + if status.state and status.state.terminated is not None: + reason = status.state.terminated.reason + exit_code = status.state.terminated.exit_code + raise RuntimeError( + f"Ephemeral container '{container_name}' in pod" + f" '{actual_pod_name}' terminated unexpectedly (reason={reason}," + f" exit_code={exit_code})." + ) + time.sleep(2) + raise RuntimeError( + f"Timeout waiting for ephemeral container '{container_name}' in pod" + f" '{pod_name}' to become running." + ) + + +def delete_worker_pods( + pod_names: Iterable[str], + namespace: str = "default", +) -> None: + """Deletes the specified worker Pods so the JobSet controller recreates them.""" + _validate_k8s_name(namespace) + worker_jobs: set[str] = set() + for pod_name in pod_names: + try: + resolved_name = pod_name + try: + core_api = _get_k8s_core_api() + pod = _read_worker_pod(core_api, pod_name, namespace=namespace) + if getattr(pod, "metadata", None): + if isinstance(getattr(pod.metadata, "name", None), str): + resolved_name = pod.metadata.name + labels = getattr(pod.metadata, "labels", None) + if isinstance(labels, dict): + job_name = labels.get("job-name") or labels.get( + "batch.kubernetes.io/job-name" + ) + if isinstance(job_name, str) and job_name: + worker_jobs.add(job_name) + except Exception: # pylint: disable=broad-exception-caught + resolved_name = pod_name + delete_gke_resource("pod", resolved_name, namespace=namespace, wait=False) + except Exception as e: # pylint: disable=broad-exception-caught + _logger.exception("Failed to delete worker pod '%s': %r", pod_name, e) + + for job_name in sorted(worker_jobs): + try: + delete_gke_resource("job", job_name, namespace=namespace, wait=False) + except Exception as e: # pylint: disable=broad-exception-caught + _logger.exception("Failed to delete worker job '%s': %r", job_name, e) diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index 6c8411d..65dea59 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py @@ -8,6 +8,7 @@ import logging import os import random +import re import signal import string import subprocess @@ -44,6 +45,10 @@ DEFAULT_PROXY_IMAGE = ( "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest" ) +EPHEMERAL_SIDECAR_CONTAINER_NAME = "ephemeral-python-sidecar" +DEFAULT_SIDECAR_PORT = 50051 +EPHEMERAL_SIDECAR_PORT = 50052 +DEFAULT_SIDECAR_CONNECT_TIMEOUT = "300s" _logger = logging.getLogger(__name__) @@ -57,16 +62,26 @@ class ProxyOptions: proxy server. xla_flags: A list of XLA flags to pass to the proxy server. sidecar: Whether to use the worker sidecar or not. + sidecar_image: Optional custom sidecar image to inject as an ephemeral + container into placed worker pods. + sidecar_port: Port used by the external sidecar gRPC server. + sidecar_connect_timeout: Timeout for connecting to the external sidecar. """ use_insecure_credentials: bool = False xla_flags: list[str] = dataclasses.field(default_factory=list) sidecar: bool = False + sidecar_image: str | None = None + sidecar_port: int = DEFAULT_SIDECAR_PORT + sidecar_connect_timeout: str = DEFAULT_SIDECAR_CONNECT_TIMEOUT @classmethod def from_list(cls, options: Iterable[str] | None) -> "ProxyOptions": """Creates a ProxyOptions object from a list of 'key:value' strings.""" use_insecure = False use_sidecar = False + sidecar_image = None + sidecar_port = DEFAULT_SIDECAR_PORT + sidecar_connect_timeout = DEFAULT_SIDECAR_CONNECT_TIMEOUT xla_flags = [] for option in options or []: if ":" in option: @@ -76,6 +91,12 @@ def from_list(cls, options: Iterable[str] | None) -> "ProxyOptions": use_insecure = value.strip().lower() == "true" elif key_strip == "sidecar": use_sidecar = value.strip().lower() == "true" + elif key_strip == "sidecar_image": + val_strip = value.strip() + if val_strip: + sidecar_image = val_strip + use_sidecar = True + sidecar_port = EPHEMERAL_SIDECAR_PORT elif key_strip == "xla_flags": val_strip = value.strip() if ( @@ -95,6 +116,9 @@ def from_list(cls, options: Iterable[str] | None) -> "ProxyOptions": use_insecure_credentials=use_insecure, xla_flags=xla_flags, sidecar=use_sidecar, + sidecar_image=sidecar_image, + sidecar_port=sidecar_port, + sidecar_connect_timeout=sidecar_connect_timeout, ) @@ -151,8 +175,17 @@ def _deploy_pathways_proxy_server( ) proxy_args_str = "\n" + proxy_args_str - if proxy_options.sidecar: + if proxy_options.sidecar or proxy_options.sidecar_image: proxy_args_str += "\n - --sidecar_name=external" + if proxy_options.sidecar_image: + proxy_args_str += ( + "\n -" + f" --cloud_pathways_external_sidecar_port={proxy_options.sidecar_port}" + "\n -" + " --cloud_pathways_external_sidecar_connect_timeout=" + f"{proxy_options.sidecar_connect_timeout}" + "\n - --vmodule=location_map=2" + ) escaped_proxy_server_image = ( proxy_server_image.replace("\\", "\\\\") @@ -179,12 +212,66 @@ def _deploy_pathways_proxy_server( _logger.info("Successfully deployed Pathways proxy.") +def _extract_pod_names_from_log_line(line: str) -> list[str]: + """Extracts worker Pod names from a pw-proxy log line.""" + pods = [] + # Format 1: Transition slice ... unplaced -> placed on worker pods: [p1, p2] + match = re.search(r"on worker pods:\s*\[([^\]]+)\]", line, re.IGNORECASE) + if match: + for item in match.group(1).split(","): + pod = item.strip().split(".")[0].split(":")[0] + if pod: + pods.append(pod) + # Format 2: VLOG(2) Sidecar init request resolved_addresses: "pod.jobset:port" + for addr in re.findall(r'resolved_addresses:\s*"([^"]+)"', line): + pod = addr.strip().split(".")[0].split(":")[0] + if pod: + pods.append(pod) + return pods + + +def _inject_ephemeral_sidecars( + pod_names: Iterable[str], + sidecar_image: str, + sidecar_port: int, +) -> None: + """Injects and waits for the ephemeral sidecar container on placed pods.""" + pods = sorted(set(pod_names)) + if not pods: + _logger.warning( + "No placed worker pods found in proxy logs; skipping ephemeral sidecar" + " injection." + ) + return + _logger.info( + "Injecting ephemeral sidecar image '%s' into %d placed worker pod(s): %s", + sidecar_image, + len(pods), + pods, + ) + for pod_name in pods: + gke_utils.inject_ephemeral_sidecar( + pod_name=pod_name, + container_name=EPHEMERAL_SIDECAR_CONTAINER_NAME, + image=sidecar_image, + port=sidecar_port, + ) + for pod_name in pods: + gke_utils.wait_for_ephemeral_container( + pod_name=pod_name, + container_name=EPHEMERAL_SIDECAR_CONTAINER_NAME, + ) + + def _wait_for_placement( log_process: subprocess.Popen[str], num_slices: int, metrics_collector_inst: Any = None, start_time: float | None = None, total_chips: int = 0, + sidecar_image: str | None = None, + sidecar_port: int = EPHEMERAL_SIDECAR_PORT, + placed_pods_out: set[str] | None = None, ) -> None: """Waits for the placement to be complete by checking proxy logs.""" _logger.info("Streaming proxy logs until the placement is complete...") @@ -195,7 +282,9 @@ def _wait_for_placement( "FAILED_PRECONDITION", ] end_phrase = "unplaced -> placed" + reconfig_phrase = "component reconfiguration after placement changes" placement_count = 0 + placed_pods: set[str] = set() if not log_process.stdout: _logger.error("Log streaming process stdout is empty. Terminating.") @@ -206,11 +295,28 @@ def _wait_for_placement( f"STDERR: {stderr}" ) + def _complete_placement() -> None: + if placed_pods_out is not None: + placed_pods_out.update(placed_pods) + if sidecar_image: + _inject_ephemeral_sidecars(placed_pods, sidecar_image, sidecar_port) + if metrics_collector_inst is not None: + metrics_collector_inst.record_active_user(True) + metrics_collector_inst.record_capacity_in_use(total_chips) + if start_time: + duration = time.time() - start_time + metrics_collector_inst.record_assignment_time(duration) + metrics_collector_inst.record_successful_request() + + waiting_for_vlog_pods = False for line in log_process.stdout: line_lower = line.lower() if any(keyword.lower() in line_lower for keyword in keywords): _logger.info("Proxy log: %s", line.strip()) + for pod in _extract_pod_names_from_log_line(line): + placed_pods.add(pod) + if end_phrase.lower() in line_lower: placement_count += 1 if placement_count < num_slices: @@ -221,13 +327,18 @@ def _wait_for_placement( ) else: _logger.info("TPU placement for %d slice(s) complete!", num_slices) - metrics_collector_inst.record_active_user(True) - metrics_collector_inst.record_capacity_in_use(total_chips) - if start_time: - duration = time.time() - start_time - metrics_collector_inst.record_assignment_time(duration) - metrics_collector_inst.record_successful_request() - break + if sidecar_image and not placed_pods: + waiting_for_vlog_pods = True + else: + _complete_placement() + return + + if waiting_for_vlog_pods and reconfig_phrase in line_lower: + _complete_placement() + return + + if waiting_for_vlog_pods: + _complete_placement() def _restore_env_var(key: str, original_value: str | None) -> None: @@ -306,6 +417,7 @@ def __init__( self._old_jax_platforms_config = None self._old_jax_backend_target_config = None self.total_chips = self._get_total_chips() + self._placed_worker_pods: set[str] = set() self._cleaned_up = False self._cleanup_lock = threading.Lock() self._original_signal_handlers: dict[signal.Signals, Any] = {} @@ -433,6 +545,25 @@ def __enter__(self): os.environ[_JAX_PLATFORMS_KEY.upper()] = _JAX_PLATFORM_PROXY os.environ[_JAX_BACKEND_TARGET_KEY.upper()] = jax_backend_target + if self.proxy_options.sidecar_image and self.proxy_pod_name: + num_slices = sum(self.expected_tpu_instances.values()) + self._log_process = gke_utils.stream_pod_logs(self.proxy_pod_name) + placement_thread = threading.Thread( + target=_wait_for_placement, + args=( + self._log_process, + num_slices, + self.metrics_collector, + self.start_time, + self.total_chips, + self.proxy_options.sidecar_image, + self.proxy_options.sidecar_port, + self._placed_worker_pods, + ), + daemon=True, + ) + placement_thread.start() + pathwaysutils.initialize() _logger.info( "Interactive supercomputing proxy client ready for cluster '%s'.", @@ -506,6 +637,17 @@ def _cleanup(self) -> None: "Failed to delete Pathways proxy GKE job: %r", e ) + # Delete placed worker pods that had ephemeral sidecars injected so the + # JobSet controller recreates clean worker pods. + if self._placed_worker_pods: + _logger.info( + "Deleting %d worker pod(s) with ephemeral sidecar containers: %s", + len(self._placed_worker_pods), + sorted(self._placed_worker_pods), + ) + gke_utils.delete_worker_pods(sorted(self._placed_worker_pods)) + self._placed_worker_pods.clear() + # Restore JAX variables. _logger.info("Restoring JAX env and config variables...") _restore_env_var(_JAX_PLATFORMS_KEY.upper(), self._old_jax_platforms) @@ -583,6 +725,7 @@ def connect( proxy_job_name: str | None = None, proxy_server_image: str | None = None, proxy_options: Sequence[str] | None = None, + sidecar_image: str | None = None, collect_service_metrics: bool = False, ) -> Iterator["_ISCPathways"]: """Connects to a Pathways server if the cluster exists. If not, creates it. @@ -599,10 +742,12 @@ def connect( random name will be generated. proxy_server_image: (Deprecated) The proxy server image to use. If not provided, it will be auto-detected from the Pathways service. If the given - proxy image is incompatible with the Pathways service, it will be - replaced with the compatible proxy image. + proxy image is incompatible with the Pathways service, it will be replaced + with the compatible proxy image. proxy_options: Configuration options for the Pathways proxy. If not provided, no extra options will be used. + sidecar_image: Optional custom colocated Python sidecar image to inject as + an ephemeral container into assigned worker pods. collect_service_metrics: Whether to collect usage metrics for Shared Pathways Service. @@ -625,7 +770,7 @@ def connect( cluster=cluster, project=project, location=region ) - server_image, sidecar_image = gke_utils.get_pathways_service_images( + server_image, service_sidecar_image = gke_utils.get_pathways_service_images( pathways_service ) compatible_proxy_image = gke_utils.get_compatible_proxy_server_image( @@ -645,8 +790,16 @@ def connect( proxy_server_image = compatible_proxy_image proxy_options_obj = ProxyOptions.from_list(proxy_options) - if proxy_options_obj.sidecar and sidecar_image: - validators.validate_sidecar_image_versions(sidecar_image) + if sidecar_image: + proxy_options_obj.sidecar = True + proxy_options_obj.sidecar_image = sidecar_image + proxy_options_obj.sidecar_port = EPHEMERAL_SIDECAR_PORT + + effective_sidecar_image = ( + proxy_options_obj.sidecar_image or service_sidecar_image + ) + if proxy_options_obj.sidecar and effective_sidecar_image: + validators.validate_sidecar_image_versions(effective_sidecar_image) _logger.info("Validation complete.") if not proxy_job_name: @@ -670,20 +823,21 @@ def connect( collect_service_metrics=collect_service_metrics, ) as t: if t.proxy_pod_name: - num_slices = sum(t.expected_tpu_instances.values()) - t._log_process = gke_utils.stream_pod_logs(t.proxy_pod_name) - placement_thread = threading.Thread( - target=_wait_for_placement, - args=( - t._log_process, - num_slices, - t.metrics_collector, - t.start_time, - t.total_chips, - ), - daemon=True, - ) - placement_thread.start() + if not proxy_options_obj.sidecar_image: + num_slices = sum(t.expected_tpu_instances.values()) + t._log_process = gke_utils.stream_pod_logs(t.proxy_pod_name) + placement_thread = threading.Thread( + target=_wait_for_placement, + args=( + t._log_process, + num_slices, + t.metrics_collector, + t.start_time, + t.total_chips, + ), + daemon=True, + ) + placement_thread.start() else: _logger.warning( "proxy_pod_name not set on _ISCPathways instance, skipping background" diff --git a/pathwaysutils/experimental/shared_pathways_service/run_connect_example.py b/pathwaysutils/experimental/shared_pathways_service/run_connect_example.py index 0ba7509..79a9923 100644 --- a/pathwaysutils/experimental/shared_pathways_service/run_connect_example.py +++ b/pathwaysutils/experimental/shared_pathways_service/run_connect_example.py @@ -46,6 +46,12 @@ ' "key:value". For example: --proxy_options=use_insecure_credentials:true' ' or --proxy_options=xla_flags:"--xla_flag1 --xla_flag2"', ) +flags.DEFINE_string( + "sidecar_image", + None, + "Optional custom colocated Python sidecar container image to inject as an" + " ephemeral container into assigned worker pods.", +) flags.DEFINE_bool( "collect_service_metrics", @@ -76,6 +82,7 @@ def main(argv: Sequence[str]) -> None: proxy_job_name=FLAGS.proxy_job_name, proxy_server_image=FLAGS.proxy_server_image, proxy_options=FLAGS.proxy_options, + sidecar_image=FLAGS.sidecar_image, collect_service_metrics=FLAGS.collect_service_metrics, ): # your-workload diff --git a/pathwaysutils/experimental/shared_pathways_service/run_workload.py b/pathwaysutils/experimental/shared_pathways_service/run_workload.py index 3ba940a..0b1d258 100644 --- a/pathwaysutils/experimental/shared_pathways_service/run_workload.py +++ b/pathwaysutils/experimental/shared_pathways_service/run_workload.py @@ -67,6 +67,12 @@ ' "key:value". For example: --proxy_options=use_insecure_credentials:true' ' or --proxy_options=xla_flags:"--xla_flag1 --xla_flag2"', ) +_SIDECAR_IMAGE = flags.DEFINE_string( + "sidecar_image", + None, + "Optional custom colocated Python sidecar container image to inject as an" + " ephemeral container into assigned worker pods.", +) _COMMAND = flags.DEFINE_string( "command", None, "The command to run on TPUs.", required=True ) @@ -92,6 +98,7 @@ def run_command( command: str, proxy_server_image: str | None = None, proxy_options: Sequence[str] | None = None, + sidecar_image: str | None = None, collect_service_metrics: bool = False, connect_fn: Callable[..., ContextManager[Any]] = isc_pathways.connect, ) -> None: @@ -109,6 +116,7 @@ def run_command( proxy_server_image: (Deprecated) The proxy server image to use. If not provided, it will be auto-detected from the Pathways service. proxy_options: Configuration options for the Pathways proxy. + sidecar_image: Optional custom colocated Python sidecar container image. collect_service_metrics: Whether to collect usage metrics for Shared Pathways Service. Defaults to False. connect_fn: The function to use for establishing the connection context, @@ -126,19 +134,20 @@ def run_command( stacklevel=2, ) logging.info("Connecting to Shared Pathways Service...") - with connect_fn( + connect_kwargs: dict[str, Any] = dict( cluster=cluster, project=project, region=region, gcs_bucket=gcs_bucket, pathways_service=pathways_service, expected_tpu_instances={tpu_type: tpu_count}, - proxy_server_image=( - proxy_server_image if proxy_server_image else None - ), + proxy_server_image=(proxy_server_image if proxy_server_image else None), proxy_options=proxy_options, collect_service_metrics=collect_service_metrics, - ): + ) + if sidecar_image: + connect_kwargs["sidecar_image"] = sidecar_image + with connect_fn(**connect_kwargs): logging.info("Connection established. Running command: %r", command) command_args = shlex.split(command) proc = subprocess.Popen(command_args, env=os.environ.copy()) @@ -179,6 +188,7 @@ def main(argv: Sequence[str]) -> None: command=_COMMAND.value, proxy_server_image=_PROXY_SERVER_IMAGE.value, proxy_options=_PROXY_OPTIONS.value, + sidecar_image=_SIDECAR_IMAGE.value, collect_service_metrics=_COLLECT_SERVICE_METRICS.value, ) diff --git a/pathwaysutils/test/experimental/gke/jobset_test.py b/pathwaysutils/test/experimental/gke/jobset_test.py index 41fd382..51b4de3 100644 --- a/pathwaysutils/test/experimental/gke/jobset_test.py +++ b/pathwaysutils/test/experimental/gke/jobset_test.py @@ -106,7 +106,6 @@ def get_all_containers_by_name( return matches - class PathwaysJobSetTest(parameterized.TestCase): def _create_jobset( @@ -328,7 +327,7 @@ def test_add_gcsfuse_volumes_and_annotations(self): def test_add_gcsfuse_handles_none_metadata(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Force metadata to be None to simulate imported templates or raw specs without metadata pw_jobset._head_job_template.metadata = None pw_jobset._head_job_template.spec.template.metadata = None @@ -342,7 +341,7 @@ def test_add_gcsfuse_handles_none_metadata(self): bucket="my-bucket", ) helper = JobSetManifestHelper(pw_jobset.to_dict()) - + self.assertEqual(helper.job_metadatas["pathways-head"].get("annotations", {}).get("gke-gcsfuse/volumes"), "true") self.assertEqual(helper.pod_metadatas["pathways-head"].get("annotations", {}).get("gke-gcsfuse/volumes"), "true") self.assertEqual(helper.job_metadatas["pathways-worker"].get("annotations", {}).get("gke-gcsfuse/volumes"), "true") @@ -350,7 +349,7 @@ def test_add_gcsfuse_handles_none_metadata(self): def test_add_gcsfuse_preserves_existing_metadata(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Pre-populate metadata, annotations, and labels pw_jobset._head_job_template.metadata = client.V1ObjectMeta( labels={"existing-job-label": "value"}, @@ -403,14 +402,14 @@ def test_add_gcsfuse_preserves_existing_volumes(self): def test_add_colocated_python_handles_none_volumes(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Force volumes to be None pw_jobset._worker_job_template.spec.template.spec.volumes = None # Should not crash and should correctly add volume pw_jobset.add_colocated_python(image="gcr.io/my-project/colocated-python:custom") helper = JobSetManifestHelper(pw_jobset.to_dict()) - + self.assertIn("shared-memory", helper.volumes["pathways-worker"]) def test_add_colocated_python_sidecar(self): @@ -436,9 +435,28 @@ def test_add_colocated_python_sidecar(self): ) ) + def test_add_colocated_python_without_image_provisions_shm_only(self): + pw_jobset = self._create_jobset(topology="2x2", num_slices=1) + + pw_jobset.add_colocated_python(image="") + helper = JobSetManifestHelper(pw_jobset.to_dict()) + + self.assertNotIn( + "colocated-python-sidecar", helper.containers["pathways-worker"] + ) + self.assertIn("shared-memory", helper.volumes["pathways-worker"]) + worker = helper.containers["pathways-worker"]["pathways-worker"] + self.assertTrue( + any( + m["name"] == "shared-memory" + and m["mountPath"] == "/tmp/shared-memory" + for m in worker["volumeMounts"] + ) + ) + def test_add_colocated_python_preserves_init_containers(self): pw_jobset = self._create_jobset(topology="2x2", num_slices=1) - + # Pre-populate init container on worker pod worker_spec = pw_jobset._worker_job_template.spec.template.spec existing_init = client.V1Container(name="existing-init-container", image="ubuntu:latest") diff --git a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py index dbc9340..73911bd 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py @@ -1593,7 +1593,160 @@ def test_terminate_process_custom_timeout(self): mock_proc.terminate.assert_called_once() mock_proc.wait.assert_called_once_with(timeout=10) + def test_inject_ephemeral_sidecar_success(self): + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + mock_pod = mock.MagicMock() + mock_pod.spec.ephemeral_containers = None + mock_core_api.read_namespaced_pod.return_value = mock_pod + + gke_utils.inject_ephemeral_sidecar( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + image="custom-sidecar:latest", + port=50052, + shm_mount_path="/tmp/sidecar_dir", + namespace="default", + ) + + mock_core_api.read_namespaced_pod.assert_called_once_with( + name="worker-pod-0-0", namespace="default" + ) + mock_core_api.patch_namespaced_pod_ephemeralcontainers.assert_called_once() + _, kwargs = mock_core_api.patch_namespaced_pod_ephemeralcontainers.call_args + self.assertEqual(kwargs["name"], "worker-pod-0-0") + self.assertEqual(kwargs["namespace"], "default") + ephemeral_list = kwargs["body"]["spec"]["ephemeralContainers"] + self.assertLen(ephemeral_list, 1) + self.assertEqual(ephemeral_list[0].name, "ephemeral-python-sidecar") + self.assertEqual(ephemeral_list[0].image, "custom-sidecar:latest") + self.assertEqual(ephemeral_list[0].target_container_name, "pathways-worker") + self.assertEqual( + ephemeral_list[0].args, + ["--port=50052", "--logtostderr", "--stderrthreshold=0", "--v=1"], + ) + + def test_inject_ephemeral_sidecar_resolves_indexed_job_pod(self): + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + mock_core_api.read_namespaced_pod.side_effect = ( + client.exceptions.ApiException(status=404, reason="Not Found") + ) + mock_pod = mock.MagicMock() + mock_pod.metadata.name = "worker-pod-0-0-abcde" + mock_pod.metadata.deletion_timestamp = None + mock_pod.status.phase = "Running" + mock_pod.spec.ephemeral_containers = None + mock_pod_list = mock.MagicMock() + mock_pod_list.items = [mock_pod] + mock_core_api.list_namespaced_pod.return_value = mock_pod_list + + gke_utils.inject_ephemeral_sidecar( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + image="custom-sidecar:latest", + port=50052, + namespace="default", + ) + + mock_core_api.list_namespaced_pod.assert_called_once_with( + namespace="default", + label_selector=( + "job-name=worker-pod-0,batch.kubernetes.io/job-completion-index=0" + ), + ) + _, kwargs = mock_core_api.patch_namespaced_pod_ephemeralcontainers.call_args + self.assertEqual(kwargs["name"], "worker-pod-0-0-abcde") + + def test_inject_ephemeral_sidecar_already_exists_skips(self): + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + existing = client.V1EphemeralContainer( + name="ephemeral-python-sidecar", image="custom-sidecar:latest" + ) + mock_pod = mock.MagicMock() + mock_pod.spec.ephemeral_containers = [existing] + mock_core_api.read_namespaced_pod.return_value = mock_pod + + gke_utils.inject_ephemeral_sidecar( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + image="custom-sidecar:latest", + ) + + mock_core_api.patch_namespaced_pod_ephemeralcontainers.assert_not_called() + + def test_wait_for_ephemeral_container_running_success(self): + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + status = mock.MagicMock() + status.name = "ephemeral-python-sidecar" + status.state.running = mock.MagicMock() + status.state.terminated = None + mock_pod = mock.MagicMock() + mock_pod.status.ephemeral_container_statuses = [status] + mock_core_api.read_namespaced_pod.return_value = mock_pod + + gke_utils.wait_for_ephemeral_container( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + timeout=10, + ) + mock_core_api.read_namespaced_pod.assert_called_once_with( + name="worker-pod-0-0", namespace="default" + ) + + def test_wait_for_ephemeral_container_terminated_raises(self): + mock_core_api = mock.MagicMock(spec=client.CoreV1Api) + self.enter_context( + mock.patch.object( + gke_utils, "_get_k8s_core_api", return_value=mock_core_api + ) + ) + status = mock.MagicMock() + status.name = "ephemeral-python-sidecar" + status.state.running = None + status.state.terminated.reason = "Error" + status.state.terminated.exit_code = 1 + mock_pod = mock.MagicMock() + mock_pod.status.ephemeral_container_statuses = [status] + mock_core_api.read_namespaced_pod.return_value = mock_pod + + with self.assertRaisesRegex(RuntimeError, "terminated unexpectedly"): + gke_utils.wait_for_ephemeral_container( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + timeout=10, + ) + + def test_delete_worker_pods_calls_delete_gke_resource(self): + mock_delete = self.enter_context( + mock.patch.object(gke_utils, "delete_gke_resource", autospec=True) + ) + gke_utils.delete_worker_pods( + ["worker-pod-0-0", "worker-pod-0-1"], namespace="default" + ) + mock_delete.assert_has_calls([ + mock.call("pod", "worker-pod-0-0", namespace="default", wait=False), + mock.call("pod", "worker-pod-0-1", namespace="default", wait=False), + ]) + if __name__ == "__main__": absltest.main() - diff --git a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py index 9eb7620..697d8e6 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py @@ -143,7 +143,6 @@ def test_deploy_pathways_proxy_server_image_escaping(self): substituted_yaml, ) - def test_deploy_pathways_proxy_server_with_insecure_credentials_success(self): mock_deploy_gke_yaml = self.enter_context( mock.patch.object( @@ -1418,6 +1417,178 @@ def test_connect_proxy_server_image_deprecation_warning(self): ): pass + def test_proxy_options_from_list_with_sidecar_image(self): + options = isc_pathways.ProxyOptions.from_list( + ["sidecar_image:us-docker.pkg.dev/my-proj/sidecar:custom"] + ) + self.assertTrue(options.sidecar) + self.assertEqual( + options.sidecar_image, "us-docker.pkg.dev/my-proj/sidecar:custom" + ) + self.assertEqual(options.sidecar_port, 50052) + + def test_deploy_pathways_proxy_server_with_custom_sidecar_image(self): + mock_deploy_gke_yaml = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "deploy_gke_yaml", autospec=True + ) + ) + isc_pathways._deploy_pathways_proxy_server( + pathways_service="test-service:8080", + proxy_job_name="test-proxy", + expected_instances={"tpuv6e:2x2": 1}, + gcs_scratch_location="test-bucket", + proxy_server_image="test-image:latest", + proxy_options=isc_pathways.ProxyOptions( + sidecar=True, + sidecar_image="custom-sidecar:latest", + sidecar_port=50052, + sidecar_connect_timeout="300s", + ), + ) + mock_deploy_gke_yaml.assert_called_once() + substituted_yaml = mock_deploy_gke_yaml.call_args[0][0] + self.assertIn("- --sidecar_name=external", substituted_yaml) + self.assertIn( + "- --cloud_pathways_external_sidecar_port=50052", substituted_yaml + ) + self.assertIn( + "- --cloud_pathways_external_sidecar_connect_timeout=300s", + substituted_yaml, + ) + self.assertIn("- --vmodule=location_map=2", substituted_yaml) + + def test_wait_for_placement_injects_ephemeral_sidecar_into_placed_pods(self): + mock_inject = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "inject_ephemeral_sidecar", autospec=True + ) + ) + mock_wait_ephemeral = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "wait_for_ephemeral_container", + autospec=True, + ) + ) + mock_process = mock.create_autospec(subprocess.Popen, instance=True) + mock_process.stdout = io.StringIO( + "I1001 12:00:00.000000 1 location_map.cc:1042] Transition slice" + " ClientSliceId(1) unplaced -> placed on worker pods:" + " [worker-pod-1-0, worker-pod-1-1]\n" + ) + placed_pods: set[str] = set() + + isc_pathways._wait_for_placement( + mock_process, + num_slices=1, + metrics_collector_inst=mock.Mock(), + sidecar_image="custom-sidecar:latest", + sidecar_port=50052, + placed_pods_out=placed_pods, + ) + + self.assertEqual(placed_pods, {"worker-pod-1-0", "worker-pod-1-1"}) + mock_inject.assert_has_calls([ + mock.call( + pod_name="worker-pod-1-0", + container_name="ephemeral-python-sidecar", + image="custom-sidecar:latest", + port=50052, + ), + mock.call( + pod_name="worker-pod-1-1", + container_name="ephemeral-python-sidecar", + image="custom-sidecar:latest", + port=50052, + ), + ]) + mock_wait_ephemeral.assert_has_calls([ + mock.call( + pod_name="worker-pod-1-0", + container_name="ephemeral-python-sidecar", + ), + mock.call( + pod_name="worker-pod-1-1", + container_name="ephemeral-python-sidecar", + ), + ]) + + def test_wait_for_placement_extracts_pods_from_vlog_sidecar_init(self): + mock_inject = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "inject_ephemeral_sidecar", autospec=True + ) + ) + mock_wait_ephemeral = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "wait_for_ephemeral_container", + autospec=True, + ) + ) + mock_process = mock.create_autospec(subprocess.Popen, instance=True) + mock_process.stdout = io.StringIO( + "I1001 12:00:00.000000 1 location_map.cc:1042] Transition slice" + " ClientSliceId(1) unplaced -> placed\n" + "I1001 12:00:00.000100 1 location_map.cc:1178] Sidecar init request:" + ' resolved_addresses: "worker-pod-0-0.my-jobset:29005"\n' + "I1001 12:00:00.000200 1 location_map.cc:933] Component" + " reconfiguration after placement changes took 1ms\n" + ) + placed_pods: set[str] = set() + + isc_pathways._wait_for_placement( + mock_process, + num_slices=1, + metrics_collector_inst=mock.Mock(), + sidecar_image="custom-sidecar:latest", + sidecar_port=50052, + placed_pods_out=placed_pods, + ) + + self.assertEqual(placed_pods, {"worker-pod-0-0"}) + mock_inject.assert_called_once_with( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + image="custom-sidecar:latest", + port=50052, + ) + mock_wait_ephemeral.assert_called_once_with( + pod_name="worker-pod-0-0", + container_name="ephemeral-python-sidecar", + ) + + def test_cleanup_deletes_placed_worker_pods(self): + manager = isc_pathways._ISCPathways( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_job_name="test-proxy", + proxy_server_image="test-image", + ) + manager._placed_worker_pods = {"worker-pod-1-1", "worker-pod-1-0"} + self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "delete_gke_resource", autospec=True + ) + ) + mock_delete_workers = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "delete_worker_pods", autospec=True + ) + ) + + manager._cleanup() + + mock_delete_workers.assert_called_once_with( + ["worker-pod-1-0", "worker-pod-1-1"] + ) + self.assertEmpty(manager._placed_worker_pods) + class KubeConfigCredentialsTest(parameterized.TestCase): """Tests for the kube config context helpers.""" diff --git a/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py b/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py index 7323000..aa8bd46 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py @@ -192,6 +192,7 @@ def test_main_calls_run_command(self): command="echo hello", proxy_server_image="", proxy_options=[], + sidecar_image=None, collect_service_metrics=False, ) @@ -222,9 +223,41 @@ def test_main_calls_run_command_with_metrics_enabled(self): command="echo hello", proxy_server_image="", proxy_options=[], + sidecar_image=None, collect_service_metrics=True, ) + @flagsaver.flagsaver( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + tpu_type="tpuv6e:4x8", + tpu_count=1, + command="echo hello", + sidecar_image="custom-sidecar:latest", + ) + def test_main_calls_run_command_with_sidecar_image(self): + with mock.patch.object( + run_workload, "run_command", autospec=True + ) as mock_run_command: + run_workload.main(["unused_argv"]) + mock_run_command.assert_called_once_with( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + tpu_type="tpuv6e:4x8", + tpu_count=1, + command="echo hello", + proxy_server_image="", + proxy_options=[], + sidecar_image="custom-sidecar:latest", + collect_service_metrics=False, + ) + if __name__ == "__main__": # Provide dummy values for required flags to satisfy absl verification.