pathwaysutils 0.1.7__tar.gz → 0.1.8__tar.gz

{pathwaysutils-0.1.7 → pathwaysutils-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.7
+Version: 0.1.8
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>

{pathwaysutils-0.1.7 → pathwaysutils-0.1.8}/pathwaysutils/__init__.py

@@ -21,4 +21,4 @@ is_pathways_backend_used: Callable[[], bool] = _initialize.is_pathways_backend_u
 del _initialize
 
 # When changing this, also update the CHANGELOG.md.
-__version__ = "v0.1.7"
+__version__ = "v0.1.8"

pathwaysutils-0.1.8/pathwaysutils/experimental/concatenate_by_mesh_axis.py

@@ -0,0 +1,152 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Pathwaysutils JAX concatenate_by_mesh_axis."""
+
+from collections.abc import Sequence
+import itertools
+from typing import Any
+import jax
+import numpy as np
+from pathwaysutils import jax as pw_jax
+
+
+def concatenate_by_mesh_axis(
+    array_trees: Sequence[Any],
+    mesh_axis: str,
+) -> Any:
+  """Concatenates meshes by an axis. Returns arrays on the concatenated mesh.
+
+  Note: This API always donates the given arrays.
+
+  Args:
+    array_trees: Sequence of PyTrees of JAX arrays with `NamedSharding`. All
+      PyTrees in the sequence have the same structure. All arrays in each PyTree
+      are sharded/replicated on the same mesh. The input arrays are always
+      donated.
+    mesh_axis: Mesh axis to concatenate.
+
+  Returns:
+    A PyTree with the same structure as `array_trees[i]`. It has arrays with
+    their shards concatenated to match a concatenated mesh.
+  """
+  if not array_trees:
+    return array_trees
+
+  def _get_named_sharding(array: jax.Array) -> jax.sharding.NamedSharding:
+    if not isinstance(array, jax.Array):
+      raise ValueError(f"Elements must be jax.Array. Got {type(array)}")
+    sharding = array.sharding
+    if not isinstance(sharding, jax.sharding.NamedSharding):
+      raise ValueError(f"Expected NamedSharding. Got {type(sharding)}")
+    return sharding
+
+  flats_and_defs = [jax.tree_util.tree_flatten(at) for at in array_trees]
+  flattened_arrays = [fd[0] for fd in flats_and_defs]
+  treedefs = [fd[1] for fd in flats_and_defs]
+  input_treedef = treedefs[0]
+  for td in treedefs[1:]:
+    if td != input_treedef:
+      raise ValueError(
+          "All array trees must have the same treedef. Got"
+          f" {td} vs. {input_treedef}"
+      )
+
+  # Convert to have the output array structure in the outer list, and each entry
+  # be a list of arrays from each shard for the concatenated output array.
+  input_flat_arrays = list(zip(*flattened_arrays))
+
+  if not flattened_arrays[0]:
+    return array_trees[0]
+
+  # Extract the shared mesh from each PyTree (from an arbitrary array in each).
+  meshes_to_concatenate = [
+      _get_named_sharding(array).mesh for array in input_flat_arrays[0]
+  ]
+
+  # Validate that the meshes are compatible.
+  reference_mesh = meshes_to_concatenate[0]
+  if mesh_axis not in reference_mesh.axis_names:
+    raise ValueError(
+        f"mesh_axis '{mesh_axis}' not found in mesh axis names:"
+        f" {reference_mesh.axis_names}"
+    )
+  mesh_axis_idx = reference_mesh.axis_names.index(mesh_axis)
+  for mesh in meshes_to_concatenate:
+    if mesh.axis_names != reference_mesh.axis_names:
+      raise ValueError(
+          "Meshes must have the same axis names. Got"
+          f" {mesh} vs. {reference_mesh}."
+      )
+    if (
+        mesh.axis_sizes[:mesh_axis_idx]
+        != reference_mesh.axis_sizes[:mesh_axis_idx]
+        or mesh.axis_sizes[mesh_axis_idx + 1 :]
+        != reference_mesh.axis_sizes[mesh_axis_idx + 1 :]
+    ):
+      raise ValueError(
+          "Arrays must have the same mesh axis sizes for all axes except"
+          f" {mesh_axis}. Got {mesh} vs. {reference_mesh}."
+      )
+
+  # Construct list of the mesh axis section boundaries.
+  devices = [mesh.devices for mesh in meshes_to_concatenate]
+  mesh_axis_sections = list(
+      itertools.accumulate(
+          mesh.axis_sizes[mesh_axis_idx] for mesh in meshes_to_concatenate
+      )
+  )
+
+  concatenated_mesh = jax.sharding.Mesh(
+      np.concatenate(devices, mesh_axis_idx),
+      axis_names=reference_mesh.axis_names,
+      axis_types=reference_mesh.axis_types,
+  )
+
+  def _get_output_sharding(
+      arrays: Sequence[jax.Array],
+  ) -> jax.sharding.NamedSharding:
+    reference_sharding = _get_named_sharding(arrays[0])
+    reference_spec = reference_sharding.spec
+    return jax.sharding.NamedSharding(
+        concatenated_mesh,
+        reference_spec,
+        memory_kind=reference_sharding.memory_kind,
+    )
+
+  def _sharded_dim_idx_for_sharding(
+      sharding: jax.sharding.NamedSharding,
+  ) -> int:
+    sharded_dim_idx = -1
+    for dim_idx, dim_spec in enumerate(sharding.spec):
+      flat_dim_spec, _ = jax.tree_util.tree_flatten(dim_spec)
+      if mesh_axis in flat_dim_spec:
+        sharded_dim_idx = dim_idx
+        break
+    return sharded_dim_idx
+
+  out_shardings = [_get_output_sharding(arrays) for arrays in input_flat_arrays]
+  sharded_dim_idxs = [
+      _sharded_dim_idx_for_sharding(sharding) for sharding in out_shardings
+  ]
+
+  flat_output_arrays = pw_jax.concatenate_by_mesh_axis(
+      arrays=input_flat_arrays,
+      sharded_dim_idxs=sharded_dim_idxs,
+      mesh_axis_sizes=concatenated_mesh.axis_sizes,
+      mesh_axis_idx=mesh_axis_idx,
+      mesh_axis_sections=mesh_axis_sections,
+      out_shardings=out_shardings,
+      donate=True,
+  )
+
+  return jax.tree_util.tree_unflatten(input_treedef, flat_output_arrays)

pathwaysutils-0.1.8/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py

@@ -0,0 +1,261 @@
+"""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
+from typing import Any
+from absl import app
+from absl import flags
+from kubernetes import client
+from kubernetes import config
+import yaml
+
+_logger = logging.getLogger(__name__)
+
+# Flag definitions
+FLAGS = flags.FLAGS
+_JOBSET_NAME = flags.DEFINE_string(
+    "jobset_name", "pathways-service", "Name of the JobSet"
+)
+_JAX_VERSION = flags.DEFINE_string(
+    "jax_version", "0.9.0", "JAX version (e.g., 0.9.0)"
+)
+_SERVER_IMAGE = flags.DEFINE_string(
+    "server_image", None, "Full path to the server Docker image"
+)
+_TPU_TYPE = flags.DEFINE_enum(
+    "tpu_type", "v6e", ["v5e", "v5p", "v6e", "tpu7x"], "TPU type"
+)
+_TOPOLOGY = flags.DEFINE_string(
+    "topology", "2x2", "TPU topology (e.g., 4x8, 2x2x2)"
+)
+_NUM_SLICES = flags.DEFINE_integer(
+    "num_slices", 2, "Number of TPU slices"
+)
+_GCS_BUCKET = flags.DEFINE_string(
+    "gcs_bucket",
+    "gs://pathways-test-bucket",
+    "GCS bucket name for scratch space",
+)
+_TEMPLATE_FILE = flags.DEFINE_string(
+    "template_file",
+    os.path.join(
+        os.path.dirname(__file__), "yamls/pw-service.yaml",
+    ),
+    "Path to the JobSet YAML template file",
+)
+_DRY_RUN = flags.DEFINE_boolean(
+    "dry_run",
+    False,
+    "If true, only print the generated YAML without deploying.",
+)
+
+
+@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:
+    dims = topology.split("x")
+    if not (2 <= len(dims) <= 3):
+      return False
+    for dim in dims:
+      if not dim.isdigit():
+        return False
+      if int(dim) <= 0:
+        return False
+    return True
+  except ValueError:
+    return False
+
+
+flags.register_validator(
+    "topology",
+    _validate_topology,
+    message=(
+        "--topology must be in the format like 'AxB' or 'AxBxC', where A, B, C"
+        " are positive integers."
+    ),
+)
+
+
+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="tpuv5p",
+      ),
+      "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="tpu-v7-slice",
+          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]:
+  """Loads and substitutes the string.Template from the given path."""
+  try:
+    with open(template_path, "r") as f:
+      template_str = f.read()
+  except OSError as err:
+    raise ValueError(
+        f"Could not read template file: {template_path}: {err}"
+    ) from err
+
+  _logger.info("Template file: %s", template_path)
+  _logger.info("Context: %s", context)
+  template = string.Template(template_str)
+  _logger.info("Template: %s", template)
+  substituted_yaml = template.substitute(context)
+  return yaml.safe_load(substituted_yaml)
+
+
+def deploy_jobset(jobset_yaml: dict[str, Any]) -> None:
+  """Deploys the JobSet to the current Kubernetes cluster."""
+  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"]
+    )
+  except client.rest.ApiException:
+    _logger.exception("Error creating JobSet")
+  except config.ConfigException:
+    _logger.exception("Error loading Kubernetes configuration")
+
+
+def run_deployment(
+    tpu_type,
+    topology,
+    num_slices,
+    jobset_name,
+    gcs_bucket,
+    server_image,
+    template_file,
+    dry_run,
+    deploy_func: Callable[[dict[str, Any]], None] = deploy_jobset,
+) -> None:
+  """Executes the deployment logic."""
+  tpu_config = get_tpu_config(tpu_type)
+  vms_per_slice = calculate_vms_per_slice(topology, tpu_config.chips_per_vm)
+
+  context = {
+      "JOBSET_NAME": jobset_name,
+      "SERVER_IMAGE": server_image,
+      "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,
+  }
+
+  jobset_config = load_and_substitute_template(template_file, context)
+
+  _logger.info("--- Generated JobSet YAML ---")
+  _logger.info("\n%s", yaml.dump(jobset_config))
+
+  if not dry_run:
+    _logger.info("Deploying JobSet...")
+    deploy_func(jobset_config)
+  else:
+    _logger.info("Dry run mode, not deploying.")
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError("Too many command-line arguments.")
+
+  try:
+    if (
+        flags.FLAGS["jax_version"].present
+        and flags.FLAGS["server_image"].present
+    ):
+      raise ValueError("Cannot provide both --jax_version and --server_image")
+
+    if _SERVER_IMAGE.value:
+      server_image = _SERVER_IMAGE.value
+    else:
+      server_image = f"us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-{_JAX_VERSION.value}"
+
+    run_deployment(
+        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,
+        template_file=_TEMPLATE_FILE.value,
+        dry_run=_DRY_RUN.value,
+    )
+  except ValueError as e:
+    _logger.exception("Error: %s", e)
+  except FileNotFoundError:
+    _logger.exception(
+        "Error: Template file not found at %s", _TEMPLATE_FILE.value
+    )
+
+
+if __name__ == "__main__":
+  app.run(main)

{pathwaysutils-0.1.7 → pathwaysutils-0.1.8}/pathwaysutils/experimental/shared_pathways_service/gke_utils.py

@@ -298,6 +298,32 @@ def enable_port_forwarding(
   return (port_available, port_forward_process)
 
 
+def stream_pod_logs(pod_name: str) -> subprocess.Popen[str]:
+  """Streams logs from the given pod.
+
+  Args:
+    pod_name: The name of the pod.
+
+  Returns:
+    The process for streaming the logs.
+
+  Raises:
+    Exception: If the log streaming fails.
+  """
+  command = ["kubectl", "logs", "-f", pod_name]
+  try:
+    return subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,  # Line buffered
+    )
+  except Exception as _:
+    _logger.exception("Error streaming logs for pod %s", pod_name)
+    raise
+
+
 def delete_gke_job(job_name: str) -> None:
   """Deletes the given job from the GKE cluster.