pathwaysutils 0.1.8__tar.gz → 0.1.9__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.8
+Version: 0.1.9
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
@@ -10,7 +10,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 License-File: LICENSE
 Requires-Dist: absl-py
 Requires-Dist: fastapi
-Requires-Dist: jax>=0.8.0
+Requires-Dist: jax>=0.8.3
 Requires-Dist: orbax-checkpoint
 Requires-Dist: uvicorn
 Requires-Dist: requests

{pathwaysutils-0.1.8 → pathwaysutils-0.1.9}/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.8"
+__version__ = "v0.1.9"

{pathwaysutils-0.1.8 → pathwaysutils-0.1.9}/pathwaysutils/elastic/elastic.py

@@ -18,10 +18,13 @@ retries a function in case of `jax.errors.JaxRuntimeError` caused by slice down
 events. It also provides a utility for waiting for slices to become active.
 """
 
+from abc import ABC, abstractmethod
 import collections
-from collections.abc import Mapping, Sequence, Set
+from collections.abc import Callable, Mapping, Sequence, Set
+import itertools
 import logging
 import time
+from typing import Any
 
 import jax
 import numpy as np
@@ -41,44 +44,87 @@ _ELASTIC_DOWN_ADDITIONAL_ERROR_TYPES = frozenset([
 ])
 
 
-def _plus_one(x: jax.Array) -> jax.Array:
-  """Adds one to each element in the array.
+class SliceHealthChecker(ABC):
+  """Base class for slice health checkers.
 
-  Used to test if a slice is active.
+  Implementations must provide `dispatch` and `validate` methods.
+  `dispatch` starts the health check operations, and `validate` waits for them
+  and returns the active slice indices.
+  """
 
-  Args:
-    x: The array to add one to.
+  def __init__(self, slice_to_devices: Mapping[int, Sequence[jax.Device]]):
+    self.slice_to_devices = slice_to_devices
 
-  Returns:
-      The array with one added to each element.
-  """
-  return x + 1
+  @abstractmethod
+  def dispatch(self) -> None:
+    """Dispatches the JAX operations in the background."""
 
+  @abstractmethod
+  def validate(self) -> Set[int]:
+    """Blocks on results and returns the set of active slice indices.
 
-def _simple_execution(devices: Sequence[jax.Device]) -> jax.Array:
-  """Simple execution to test if a slice is active.
+    Returns:
+      Set of active slice indices.
+    """
 
-  This function is used to test if a slice is active. It executes a simple
-  computation on the devices and returns the result. If any of the devices are
-  not active, the returned array will fail with a JaxRuntimeError used.
 
-  Simply executing this function is not enough to determine if the slice is
-  active. We also need to check the value of the returned array.
+class DefaultSliceHealthChecker(SliceHealthChecker):
+  """Default implementation that checks each slice individually.
 
-  Args:
-    devices: The devices to execute on.
+  It executes a simple computation on each slice independently to verify
+  that the devices are active and responsive.
 
-  Returns:
-    The result of the execution.
+  The computation involves creating a zero array of the same size as the
+  number of devices in the slice, adding (_SIMPLE_EXECUTION_TEST_VALUE - 1)
+  to it, and then running a pmap to add 1 to each element. This verifies
+  that the devices are active and can perform computations.
   """
-  if not devices:
-    raise ValueError("No devices")
 
-  test_input = np.zeros(len(devices), dtype=float) + (
-      _SIMPLE_EXECUTION_TEST_VALUE - 1
-  )
+  def _plus_one(self, x: jax.Array) -> jax.Array:
+    return x + 1
+
+  def _simple_execution(self, devices: Sequence[jax.Device]) -> jax.Array:
+    if not devices:
+      raise ValueError("No devices")
+
+    test_input = np.zeros(len(devices), dtype=float) + (
+        _SIMPLE_EXECUTION_TEST_VALUE - 1
+    )
+
+    return jax.pmap(self._plus_one, devices=devices)(test_input)
+
+  def dispatch(self) -> None:
+    self.results = {
+        slice_index: self._simple_execution(devices)
+        for slice_index, devices in self.slice_to_devices.items()
+    }
+
+  def validate(self) -> Set[int]:
+    active_slice_indices = set()
+    for slice_index, x in self.results.items():
+      expected = (
+          np.zeros(len(self.slice_to_devices[slice_index]), dtype=float)
+          + _SIMPLE_EXECUTION_TEST_VALUE
+      )
+      try:
+        jax.block_until_ready(x)
+        if np.allclose(x, expected):
+          active_slice_indices.add(slice_index)
+        else:
+          msg = (
+              f"Error with _simple_execution for slice_index={slice_index}. "
+              f"Expected: {expected!r}, Actual: {x!r}"
+          )
+          _logger.error(msg)
+          raise ValueError(msg)
+      except jax.errors.JaxRuntimeError as error:
+        _logger.debug(
+            "Caught JaxRuntimeError for slice_index=%s: %s", slice_index, error
+        )
+        if not is_error_due_to_slice_down(error):
+          raise
+    return active_slice_indices
 
-  return jax.pmap(_plus_one, devices=devices)(test_input)
 
 
 def get_slice_to_devices(
@@ -94,6 +140,7 @@ def get_slice_to_devices(
 @timing.timeit
 def get_active_slice_indices(
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
+    checker: SliceHealthChecker | None = None,
 ) -> Set[int]:
   """Returns the set of active slices indices.
 
@@ -101,6 +148,8 @@ def get_active_slice_indices(
     slice_to_devices: A mapping from slice index to devices. If None,
       `get_slice_to_devices(jax.devices())` is used to gather all available
       devices and group them by slice.
+    checker: A SliceHealthChecker instance. If None, DefaultSliceHealthChecker
+      is used.
 
   Returns:
     A set of integers representing the indices of the active slices.
@@ -109,55 +158,16 @@ def get_active_slice_indices(
     _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(tuple(jax.devices()))
 
+  if checker is None:
+    checker = DefaultSliceHealthChecker(slice_to_devices)
+
   _logger.debug(
       "Getting active slice indices for slices: %s",
       sorted(list(slice_to_devices.keys())),
   )
 
-  active_slice_indices = set()
-
-  results = {
-      slice_index: _simple_execution(devices)
-      for slice_index, devices in slice_to_devices.items()
-  }
-
-  for slice_index, x in results.items():
-    _logger.debug("Checking slice_index=%s", slice_index)
-    expected = (
-        np.zeros(len(slice_to_devices[slice_index]), dtype=float)
-        + _SIMPLE_EXECUTION_TEST_VALUE
-    )
-    try:
-      with timing.Timer(f"Checking {slice_index=}"):
-        _logger.debug("Blocking until ready for slice_index=%s", slice_index)
-        jax.block_until_ready(x)
-        _logger.debug("Execution finished for slice_index=%s", slice_index)
-        if np.allclose(x, expected):
-          active_slice_indices.add(slice_index)
-          _logger.debug("slice_index=%s active", slice_index)
-        else:
-          _logger.error(
-              "Error with _simple_execution for slice_index=%s. "
-              "This should never happen. Expected: %r, Actual: %r",
-              slice_index,
-              expected,
-              x,
-          )
-          raise ValueError(
-              f"Error with _simple_execution for slice_index={slice_index}."
-          )
-    except jax.errors.JaxRuntimeError as error:
-      _logger.debug(
-          "Caught JaxRuntimeError for slice_index=%s: %s", slice_index, error
-      )
-      if not is_error_due_to_slice_down(error):
-        _logger.info("Re-raising error for slice_index=%s", slice_index)
-        raise
-      _logger.debug("slice_index=%s bad", slice_index)
-
-  _logger.debug("active_slice_indices=%s", active_slice_indices)
-
-  return active_slice_indices
+  checker.dispatch()
+  return checker.validate()
 
 
 def wait_for_slices(
@@ -165,6 +175,7 @@ def wait_for_slices(
     poll_interval: float | int = 10,
     timeout: float | int | None = None,
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
+    checker: SliceHealthChecker | None = None,
 ) -> Set[int]:
   """Waits until after at least `slice_count` slices become active.
 
@@ -177,6 +188,7 @@ def wait_for_slices(
       timeout.
     slice_to_devices: A mapping from slice index to devices. If None,
       `get_slice_to_devices(jax.devices())` is used.
+    checker: A SliceHealthChecker instance to use for checking health.
 
   Returns:
     The active slice indices
@@ -201,7 +213,7 @@ def wait_for_slices(
     check_start_time = time.time()
 
     _logger.debug("Checking active slices...")
-    active_slice_indices = get_active_slice_indices(slice_to_devices)
+    active_slice_indices = get_active_slice_indices(slice_to_devices, checker)
     if len(active_slice_indices) >= slice_count:
       _logger.info(
           "Sufficient slices active: %s >= %s. Active indices: %s",

pathwaysutils-0.1.9/pathwaysutils/experimental/gke/__init__.py

@@ -0,0 +1,11 @@
+# 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-0.1.9/pathwaysutils/experimental/gke/jobset.py

@@ -0,0 +1,414 @@
+# 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.
+"""Pathways JobSet generator and builder (Head Job Config)."""
+
+import json
+import logging
+from typing import Any, Mapping
+from kubernetes import client
+
+# GKE sidecar containers restartPolicy compatibility placeholder.
+
+_logger = logging.getLogger(__name__)
+
+# Core constants.
+PATHWAYS_HEAD_JOB_NAME = "pathways-head"
+PATHWAYS_WORKER_JOB_NAME = "pathways-worker"
+
+DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE = (
+    "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server"
+)
+DEFAULT_PATHWAYS_PROXY_IMAGE = (
+    "us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server"
+)
+
+PATHWAYS_PROXY_PORT = 29000
+PATHWAYS_RM_PORT = 29001
+
+MACHINE_TYPE_TO_TPU_VERSION_MAP = {
+    "tpu7x-standard-4t": "tpu7x",
+    "tpu7x": "tpu7x",
+    "ct6e-standard-4t": "tpuv6e",
+    "v6e": "tpuv6e",
+    "ct6e-standard-8t": "tpuv6e1t",
+    "ct5p-hightpu-4t": "tpuv5",
+    "v5p": "tpuv5",
+    "ct5lp-hightpu-4t": "tpuv5e",
+    "v5e": "tpuv5e",
+    "ct5lp-hightpu-8t": "tpuv5e1t",
+    "ct4p-hightpu-4t": "tpuv4",
+    "v4": "tpuv4",
+}
+
+MACHINE_TYPE_TO_GKE_ACCELERATOR_TYPE_MAP = {
+    "tpu7x-standard-4t": "tpu7x",
+    "tpu7x": "tpu7x",
+    "ct6e-standard-4t": "tpu-v6e-slice",
+    "v6e": "tpu-v6e-slice",
+    "ct6e-standard-8t": "tpu-v6e-slice",
+    "ct5p-hightpu-4t": "tpu-v5p-slice",
+    "v5p": "tpu-v5p-slice",
+    "ct5lp-hightpu-4t": "tpu-v5-lite-podslice",
+    "v5e": "tpu-v5-lite-podslice",
+    "ct5lp-hightpu-8t": "tpu-v5-lite-podslice",
+    "ct4p-hightpu-4t": "tpu-v4-podslice",
+    "v4": "tpu-v4-podslice",
+}
+
+
+def _deserialize_dict(
+    api_client: client.ApiClient, data_dict: Mapping[str, Any], klass: Any
+) -> Any:
+  class FakeResponse:
+
+    def __init__(self, data):
+      self.data = data
+
+  return api_client.deserialize(FakeResponse(json.dumps(data_dict)), klass)
+
+
+class PathwaysJobSet:
+  """Generates JobSet configuration for Pathways (with Head Job Config)."""
+
+  def __init__(
+      self,
+      name: str,
+      namespace: str,
+      pathways_dir: str,
+      tpu_type: str,
+      topology: str,
+      num_slices: int,
+      user_pod_template: Mapping[str, Any] | None = None,
+      main_container_name: str = "main",
+      max_restarts: int = 0,
+      pathways_version: str = "latest",
+      jobset_api_version: str = "v1alpha2",
+      elastic_slices: int = 0,
+      labels: Mapping[str, str] | None = None,
+      annotations: Mapping[str, str] | None = None,
+  ):
+    """Initializes the instance.
+
+    Args:
+      name: Name of the JobSet.
+      namespace: Namespace of the JobSet.
+      pathways_dir: GCS path for Pathways scratch space.
+      tpu_type: TPU type (e.g., "v5e").
+      topology: TPU topology (e.g., "2x2").
+      num_slices: Number of slices.
+      user_pod_template: Optional user pod template for the head job.
+      main_container_name: Name of the main container in user_pod_template.
+      max_restarts: Maximum number of restarts for the JobSet.
+      pathways_version: Version tag for Pathways images.
+      jobset_api_version: API version of JobSet.
+      elastic_slices: Number of elastic slices.
+      labels: Optional labels for the JobSet.
+      annotations: Optional annotations for the JobSet.
+    """
+    self._name = name
+    self._namespace = namespace
+    self._jobset_api_version = jobset_api_version
+    self._max_restarts = max_restarts
+    self._worker_replicas = num_slices
+    self._labels = dict(labels) if labels else {}
+    self._annotations = dict(annotations) if annotations else {}
+
+    tpu_version = MACHINE_TYPE_TO_TPU_VERSION_MAP.get(tpu_type.lower())
+    if not tpu_version:
+      raise ValueError(f"Unsupported TPU type: {tpu_type}")
+
+    instance_type = f"{tpu_version}:{topology}"
+    image_tag = pathways_version
+
+    # Build head template.
+    self._head_job_template = self._build_head_job_template(
+        pathways_dir=pathways_dir,
+        num_slices=num_slices,
+        instance_type=instance_type,
+        image_tag=image_tag,
+        user_pod_template=user_pod_template,
+        main_container_name=main_container_name,
+        elastic_slices=elastic_slices,
+    )
+
+    # Build minimal worker template (placeholder)
+    self._worker_job_template = self._build_minimal_job_template("worker")
+
+    self._success_policy = None
+    if user_pod_template:
+      self._success_policy = {
+          "operator": "All",
+          "targetReplicatedJobs": [PATHWAYS_HEAD_JOB_NAME],
+      }
+
+  def _build_minimal_job_template(self, role: str) -> client.V1JobTemplateSpec:
+    """Builds a minimal job template for a given role."""
+    pod_spec = client.V1PodSpec(
+        containers=[
+            client.V1Container(name=f"placeholder-{role}", image="ubuntu")
+        ]
+    )
+    job_spec = client.V1JobSpec(
+        template=client.V1PodTemplateSpec(
+            metadata=client.V1ObjectMeta(labels={"role": role}), spec=pod_spec
+        )
+    )
+    return client.V1JobTemplateSpec(spec=job_spec)
+
+  def _build_head_job_template(
+      self,
+      pathways_dir: str,
+      num_slices: int,
+      instance_type: str,
+      image_tag: str,
+      user_pod_template: Mapping[str, Any] | None,
+      main_container_name: str,
+      elastic_slices: int,
+  ) -> client.V1JobTemplateSpec:
+    """Builds the head job template for the JobSet.
+
+    Args:
+      pathways_dir: GCS path for Pathways scratch space.
+      num_slices: Number of slices.
+      instance_type: TPU instance type (e.g., "tpuv5:2x2").
+      image_tag: Version tag for Pathways images.
+      user_pod_template: Optional user pod template for the head job.
+      main_container_name: Name of the main container in user_pod_template.
+      elastic_slices: Number of elastic slices.
+
+    Returns:
+      The head job template.
+    """
+    rm_image = f"{DEFAULT_PATHWAYS_RM_AND_WORKER_IMAGE}:{image_tag}"
+    proxy_image = f"{DEFAULT_PATHWAYS_PROXY_IMAGE}:{image_tag}"
+
+    rm_args = [
+        f"--server_port={PATHWAYS_RM_PORT}",
+        f"--gcs_scratch_location={pathways_dir}",
+        "--node_type=resource_manager",
+        f"--instance_count={num_slices}",
+        f"--instance_type={instance_type}",
+    ]
+    rm_env = [
+        client.V1EnvVar(
+            name="REPLICATED_JOB_NAME",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path="metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name']"
+                )
+            ),
+        ),
+        client.V1EnvVar(
+            name="JOBSET_NAME",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path=(
+                        "metadata.annotations['jobset.sigs.k8s.io/jobset-name']"
+                    )
+                )
+            ),
+        ),
+        client.V1EnvVar(
+            name="HOST_ADDRESS",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path=(
+                        "metadata.labels['jobset.sigs.k8s.io/coordinator']"
+                    )
+                )
+            ),
+        ),
+        client.V1EnvVar(name="TPU_SKIP_MDS_QUERY", value="true"),
+    ]
+    rm_container = client.V1Container(
+        name="pathways-rm",
+        image=rm_image,
+        image_pull_policy="Always",
+        args=rm_args,
+        env=rm_env,
+        ports=[
+            client.V1ContainerPort(
+                container_port=PATHWAYS_RM_PORT, protocol="TCP"
+            ),
+            client.V1ContainerPort(container_port=29002, protocol="TCP"),
+        ],
+        resources=client.V1ResourceRequirements(
+            limits={"cpu": "8", "memory": "32G"}
+        ),
+    )
+
+    proxy_args = [
+        f"--server_port={PATHWAYS_PROXY_PORT}",
+        f"--resource_manager_address=$(PATHWAYS_HEAD):{PATHWAYS_RM_PORT}",
+        f"--gcs_scratch_location={pathways_dir}",
+    ]
+    if elastic_slices > 0:
+      proxy_args.append(f"--num_elastic_slices={elastic_slices}")
+
+    proxy_env = [
+        client.V1EnvVar(
+            name="PATHWAYS_HEAD",
+            value_from=client.V1EnvVarSource(
+                field_ref=client.V1ObjectFieldSelector(
+                    field_path=(
+                        "metadata.labels['jobset.sigs.k8s.io/coordinator']"
+                    )
+                )
+            ),
+        )
+    ]
+    proxy_container = client.V1Container(
+        name="pathways-proxy",
+        image=proxy_image,
+        image_pull_policy="Always",
+        args=proxy_args,
+        env=proxy_env,
+        ports=[
+            client.V1ContainerPort(
+                container_port=PATHWAYS_PROXY_PORT, protocol="TCP"
+            )
+        ],
+        resources=client.V1ResourceRequirements(
+            limits={"cpu": "16", "memory": "100G"}
+        ),
+    )
+
+    api_client = client.ApiClient()
+
+    if user_pod_template:
+      user_template_obj = _deserialize_dict(
+          api_client, user_pod_template, client.V1PodTemplateSpec
+      )
+      head_pod_spec = user_template_obj.spec
+      head_pod_spec.host_network = True
+      head_pod_spec.dns_policy = "ClusterFirstWithHostNet"
+
+      rm_container.restart_policy = "Always"
+      proxy_container.restart_policy = "Always"
+
+      init_containers = head_pod_spec.init_containers or []
+      init_containers.extend([rm_container, proxy_container])
+      head_pod_spec.init_containers = init_containers
+
+      # Inject JAX env vars into main container.
+      jax_env = [
+          client.V1EnvVar(
+              name="PATHWAYS_HEAD",
+              value_from=client.V1EnvVarSource(
+                  field_ref=client.V1ObjectFieldSelector(
+                      field_path=(
+                          "metadata.labels['jobset.sigs.k8s.io/coordinator']"
+                      )
+                  )
+              ),
+          ),
+          client.V1EnvVar(name="JAX_PLATFORMS", value="proxy"),
+          client.V1EnvVar(name="XCLOUD_ENVIRONMENT", value="GCP"),
+          client.V1EnvVar(
+              name="JAX_BACKEND_TARGET",
+              value=f"grpc://$(PATHWAYS_HEAD):{PATHWAYS_PROXY_PORT}",
+          ),
+      ]
+      containers = head_pod_spec.containers or []
+      for c in containers:
+        if c.name == main_container_name:
+          env = c.env or []
+          env.extend(jax_env)
+          c.env = env
+          break
+      head_pod_spec.containers = containers
+
+      annotations = user_pod_template.get("metadata", {}).get("annotations", {})
+      labels = user_pod_template.get("metadata", {}).get("labels", {})
+    else:
+      # Headless mode.
+      head_pod_spec = client.V1PodSpec(
+          host_network=True,
+          dns_policy="ClusterFirstWithHostNet",
+          containers=[rm_container, proxy_container],
+      )
+      annotations = {}
+      labels = {}
+
+    if not head_pod_spec.restart_policy:
+      head_pod_spec.restart_policy = "Never"
+
+    # Default annotations
+    job_annotations = {
+        "alpha.jobset.sigs.k8s.io/exclusive-topology": "kubernetes.io/hostname"
+    }
+    job_annotations.update(annotations)
+
+    head_job_template = client.V1JobTemplateSpec(
+        metadata=client.V1ObjectMeta(annotations=job_annotations),
+        spec=client.V1JobSpec(
+            backoff_limit=0,
+            completion_mode="Indexed",
+            completions=1,
+            parallelism=1,
+            template=client.V1PodTemplateSpec(
+                metadata=client.V1ObjectMeta(
+                    annotations=job_annotations, labels=labels
+                ),
+                spec=head_pod_spec,
+            ),
+        ),
+    )
+    return head_job_template
+
+  def _compile_config(self) -> dict[str, Any]:
+    """Compiles the JobSet configuration into a dictionary."""
+    with client.ApiClient() as api_client:
+      serialized_head = api_client.sanitize_for_serialization(
+          self._head_job_template
+      )
+      serialized_worker = api_client.sanitize_for_serialization(
+          self._worker_job_template
+      )
+
+    replicated_jobs = [
+        {
+            "name": PATHWAYS_HEAD_JOB_NAME,
+            "replicas": 1,
+            "template": serialized_head,
+        },
+        {
+            "name": PATHWAYS_WORKER_JOB_NAME,
+            "replicas": self._worker_replicas,
+            "template": serialized_worker,
+        },
+    ]
+
+    jobset_config = {
+        "apiVersion": f"jobset.sigs.k8s.io/{self._jobset_api_version}",
+        "kind": "JobSet",
+        "metadata": {
+            "name": self._name,
+            "namespace": self._namespace,
+        },
+        "spec": {
+            "failurePolicy": {"maxRestarts": self._max_restarts},
+            "replicatedJobs": replicated_jobs,
+        },
+    }
+    if self._labels:
+      jobset_config["metadata"]["labels"] = self._labels
+    if self._annotations:
+      jobset_config["metadata"]["annotations"] = self._annotations
+    if self._success_policy:
+      jobset_config["spec"]["successPolicy"] = self._success_policy
+
+    return jobset_config
+
+  def to_dict(self) -> dict[str, Any]:
+    """Returns the JobSet configuration as a dictionary."""
+    return self._compile_config()

{pathwaysutils-0.1.8 → pathwaysutils-0.1.9}/pathwaysutils/experimental/shared_pathways_service/deploy_pathways_service.py

@@ -102,7 +102,7 @@ def get_tpu_config(tpu_type: str) -> TPUConfig:
           machine_type="ct5p-hightpu-4t",
           chips_per_vm=4,
           accelerator_label="tpu-v5p-slice",
-          instance_prefix="tpuv5p",
+          instance_prefix="tpuv5",
       ),
       "v6e": TPUConfig(
           machine_type="ct6e-standard-4t",
@@ -113,7 +113,7 @@ def get_tpu_config(tpu_type: str) -> TPUConfig:
       "tpu7x": TPUConfig(
           machine_type="tpu7x-standard-4t",
           chips_per_vm=4,
-          accelerator_label="tpu-v7-slice",
+          accelerator_label="tpu7x",
           instance_prefix="tpu7x",
       ),
   }