xpk 0.7.1__py3-none-any.whl → 0.8.0__py3-none-any.whl

xpk/core/pathways.py

@@ -14,124 +14,24 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from .cluster import XPK_SA
+from ..core.commands import run_command_for_value, run_command_with_updates, run_commands
 from ..core.docker_container import get_user_workload_container
 from ..core.gcloud_context import zone_to_region
 from ..core.nodepool import get_all_nodepools_programmatic
 from ..utils.console import xpk_exit, xpk_print
 from .config import AcceleratorType
-from .storage import Storage, get_storage_volumes_yaml, GCS_FUSE_ANNOTATION
 from .system_characteristics import SystemCharacteristics
 
-PathwaysExpectedInstancesMap = {
-    'v6e': 'tpuv6e',
-    'v5p': 'tpuv5',
-    'v5litepod': 'tpuv5e',
-    'v4': 'tpuv4',
-    'v3': 'tpuv3',
-}
-
-
-def get_pathways_worker_args(args) -> str:
-  """Arguments for the Pathways workers.
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    str: yaml containing arguments for the Pathways workers.
-  """
-  yaml = """- --server_port=29001
-                - --resource_manager_address={rm_address}
-                - --gcs_scratch_location={args.pathways_gcs_location}"""
-  if args.use_pathways:
-    if args.custom_pathways_worker_args:
-      yaml = append_custom_pathways_args(yaml, args.custom_pathways_worker_args)
-    return yaml.format(args=args, rm_address=get_rm_address(args))
-  else:
-    return ''
-
-
-def get_pathways_proxy_args(args) -> str:
-  """Arguments for the Pathways proxy.
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    str: yaml containing arguments for the Pathways proxy.
-  """
-  yaml = """- --server_port=29000
-                - --resource_manager_address={rm_address}
-                - --gcs_scratch_location={args.pathways_gcs_location}"""
-
-  if args.use_pathways:
-    if args.custom_pathways_proxy_server_args:
-      yaml = append_custom_pathways_args(
-          yaml, args.custom_pathways_proxy_server_args
-      )
-    return yaml.format(args=args, rm_address=get_rm_address(args))
-  else:
-    return ''
-
-
-def get_pathways_sidecar_container(args) -> str:
-  """This is a sidecar container that runs the remote python server.
-
-      It is a special case of the initContainer (designated by restartPolicy:
-      Always)
-      See https://kubernetes.io/docs/concepts/workloads/pods/sidecar-containers/
-      for more details.
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    str: yaml containing arguments for the Pathways sidecar container.
-  """
-  yaml = """initContainers:
-              - name: remote-python-sidecar
-                image: {args.remote_python_sidecar_image}
-                imagePullPolicy: Always
-                securityContext:
-                  privileged: true
-                volumeMounts:
-                - mountPath: /tmp  # Shared volume mount with the main container.
-                  name: shared-tmp
-                restartPolicy: Always
-                ports:
-                - containerPort: 50051
-                env:
-                - name: GRPC_SERVER_ADDRESS
-                  value: '0.0.0.0:50051'"""
-  if args.use_pathways and args.remote_python_sidecar_image is not None:
-    return yaml.format(args=args)
-  else:
-    return ''
-
 
 def add_pw_resource_flavors(args):
   """Add resource flavors required for Pathways enabled clusters."""
   resource_flavor_yaml = """apiVersion: kueue.x-k8s.io/v1beta1
 kind: ResourceFlavor
-metadata:
-  name: cpu-rm
-spec:
-  nodeLabels:
-    cloud.google.com/gke-nodepool: cpu-rm-np
----
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: ResourceFlavor
-metadata:
-  name: cpu-proxy
-spec:
-  nodeLabels:
-    cloud.google.com/gke-nodepool: cpu-proxy-np
----
-apiVersion: kueue.x-k8s.io/v1beta1
-kind: ResourceFlavor
 metadata:
   name: cpu-user
 spec:
   nodeLabels:
-    cloud.google.com/gke-nodepool: cpu-user-np
+    cloud.google.com/gke-nodepool: cpu-np
 ---"""
   if args.enable_pathways:
     return resource_flavor_yaml
@@ -142,18 +42,6 @@ def add_pw_resources_to_kueue(args):
   """Add resource flavors required for Pathways, to the cluster queue."""
   resources_yaml = """- coveredResources: ["cpu", "memory"]
     flavors:
-    - name: cpu-rm
-      resources:
-      - name: "cpu"
-        nominalQuota: 480
-      - name: "memory"
-        nominalQuota: 2000G
-    - name: cpu-proxy
-      resources:
-      - name: "cpu"
-        nominalQuota: 480
-      - name: "memory"
-        nominalQuota: 2000G
     - name: cpu-user
       resources:
       - name: "cpu"
@@ -175,6 +63,10 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool:
   Returns:
     True once conditions satisfy and variables are set. Exits otherwise.
   """
+  # Ensure that PathwaysJob is installed and available on the cluster.
+  if not check_if_pathways_job_is_installed(args):
+    xpk_exit(1)
+
   # Ensure command is provided if not using Pathways in headless mode
   if args.command is None and not args.headless:
     xpk_print(
@@ -187,7 +79,7 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool:
 
   # Ensure the cluster and CPU nodepools were created with create-pathways
   all_node_pools = get_all_nodepools_programmatic(args)
-  desired_pw_cpu_node_pools = {'cpu-user-np', 'cpu-rm-np', 'cpu-proxy-np'}
+  desired_pw_cpu_node_pools = {'cpu-np'}
   if not desired_pw_cpu_node_pools.issubset(set(all_node_pools[0])):
     xpk_print(
         'Cluster needs to be created with `xpk create-pathways` to run'
@@ -209,6 +101,35 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool:
   return True
 
 
+def check_if_pathways_job_is_installed(args) -> bool:
+  """Check if PathwaysJob is installed on the cluster.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'kubectl get pods -n pathways-job-system --no-headers -o'
+      ' custom-columns=NAME:.metadata.name'
+  )
+  task = f'Check if PathwaysJob is installed on {args.cluster}'
+  return_code, return_msg = run_command_for_value(command, task, args)
+  # return_msg contains the name of the controller pod, if found.
+  xpk_print('check_if_pathways_job_is_installed', return_code, return_msg)
+
+  if return_code != 0:
+    xpk_print(f'{task} returned with ERROR {return_code}.\n')
+    return False
+  if not return_msg:
+    xpk_print(
+        'You are using a new version of XPK, which uses PathwaysJob'
+        ' for Pathways workloads. Please update the cluster using'
+        ' `cluster create-pathways` to enjoy the upgrade!'
+    )
+    return False
+  return True
+
+
 def get_pathways_unified_query_link(args) -> str:
   """Get the unified query link for the pathways workload."""
   query_params = (
@@ -223,58 +144,106 @@ def get_pathways_unified_query_link(args) -> str:
   return f'https://console.cloud.google.com/logs/query;query={query_params}'
 
 
-def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
-  """Arguments for the Pathways resource manager.
-  Args:
-    args: user provided arguments for running the command.
+def append_custom_pathways_flags(custom_args, prev_indentation=8) -> str:
+  """Append custom Pathways args to Pathways components using a YAML with proper indentation.
 
   Returns:
-    str: yaml containing arguments for the Pathways resource manager.
+    yaml (string): yaml with additional args appended.
   """
-  yaml = """- --server_port=29001
-                - --gcs_scratch_location={args.pathways_gcs_location}
-                - --node_type=resource_manager
-                - --instance_count={instance_count}
-                - --instance_type={instance_type}"""
-  if args.use_pathways:
-    if args.custom_pathways_server_args:
-      yaml = append_custom_pathways_args(yaml, args.custom_pathways_server_args)
-    return yaml.format(
-        args=args,
-        instance_count=args.num_slices,
-        instance_type=f'{get_pathways_expected_tpu_type(system.device_type)}:{system.topology}',
+  yaml = """"""
+  indentation = ' ' * (prev_indentation + 2)
+  if custom_args:
+    custom_args = custom_args.split(' ')
+    for arg in custom_args:
+      yaml += '\n' + indentation + '- ' + arg
+  return yaml
+
+
+def append_custom_pathways_proxy_server(args) -> str:
+  """Append custom Pathways proxy server component using a YAML with proper indentation.
+
+  Returns:
+      yaml (string): yaml with custom proxy server appended.
+  """
+  yaml = """"""
+  if args.proxy_server_image or args.custom_pathways_proxy_server_args:
+    yaml = """- componentType: proxy_server"""
+  indentation = (
+      ' ' * 8
+  )  # Currently 8, based on the YAML, may need to update in the future.
+  if args.proxy_server_image:
+    yaml += '\n' + indentation + 'image: ' + args.proxy_server_image
+  if args.custom_pathways_proxy_server_args:
+    yaml += '\n' + indentation + 'customFlags: '
+    yaml += append_custom_pathways_flags(
+        args.custom_pathways_proxy_server_args, len(indentation)
     )
-  else:
-    return ''
+  return yaml
 
 
-def append_custom_pathways_args(yaml, custom_args) -> str:
-  """Append custom Pathways args to the YAML with proper indentation.
+def append_custom_pathways_server(args) -> str:
+  """Append custom Pathways server component using a YAML with proper indentation.
 
-  Args:
-      yaml (string): existing yaml containing args
+  Returns:
+      yaml (string): yaml with custom pathways server appended.
+  """
+  yaml = """"""
+  if args.server_image or args.custom_pathways_server_args:
+    yaml = """- componentType: pathways_server"""
+  indentation = (
+      ' ' * 8
+  )  # Currently 8, based on the YAML, may need to update in the future.
+  if args.server_image:
+    yaml += '\n' + indentation + 'image: ' + args.server_image
+  if args.custom_pathways_server_args:
+    yaml += '\n' + indentation + 'customFlags: '
+    yaml += append_custom_pathways_flags(
+        args.custom_pathways_server_args, len(indentation)
+    )
+  return yaml
+
+
+def append_custom_pathways_worker(args) -> str:
+  """Append custom Pathways worker component using a YAML with proper indentation.
+
+  Returns:
+      yaml (string): yaml with custom pathways server appended.
+  """
+  yaml = """"""
+  if args.server_image or args.custom_pathways_worker_args:
+    yaml = """- componentType: pathways_worker"""
+  indentation = (
+      ' ' * 8
+  )  # Currently 8, based on the YAML, may need to update in the future.
+  if args.server_image:
+    yaml += '\n' + indentation + 'image: ' + args.server_image
+  if args.custom_pathways_worker_args:
+    yaml += '\n' + indentation + 'customFlags: '
+    yaml += append_custom_pathways_flags(
+        args.custom_pathways_worker_args, len(indentation)
+    )
+  return yaml
+
+
+def append_custom_colocated_python_sidecar(args) -> str:
+  """Append custom Pathways colocated python sidecar component using a YAML with proper indentation.
 
   Returns:
-      yaml (string): yaml with additional args appended.
+      yaml (string): yaml with custom pathways server appended.
   """
-  second_line = yaml.split('\n')[1]
-  if (
-      not second_line
-  ):  # to cover edge case if only one arg remains, we would have to look at the entire YAML in this case.
-    return yaml
-  # Calculate the indentation based on the second line of existing YAML.
-  indentation = ' ' * (len(second_line) - len(second_line.lstrip()))
-  custom_args = custom_args.split(' ')
-  for arg in custom_args:
-    yaml += '\n' + indentation + '- ' + arg
+  yaml = """"""
+  if args.colocated_python_sidecar_image:
+    yaml = """- componentType: colocated_python_sidecar"""
+    indentation = (
+        ' ' * 8
+    )  # Currently 8, based on the YAML, may need to update in the future.
+    yaml += '\n' + indentation + 'image: ' + args.colocated_python_sidecar_image
   return yaml
 
 
 def get_user_workload_for_pathways(
     args,
     system: SystemCharacteristics,
-    pod_failure_policy,
-    storages: list[Storage],
 ) -> str:
   """
   Create a user workload container for Pathways.
@@ -289,63 +258,32 @@ def get_user_workload_for_pathways(
     str:
       Pathways server port as a YAML string
   """
-  user_workload_yaml = """- name: main
-      replicas: 1
-      template:
-        metadata:
-          labels:
-            xpk.google.com/workload: {args.workload}
-        spec:
-          backoffLimit: 0
-          completions: 1
-          parallelism: 1
-          {pod_failure_policy}
-          template:
-            metadata:
-              annotations:
-                {gcs_fuse_annotation}
-            spec:
-              containers:
+  user_workload_yaml = """
+          metadata:
+          spec:
+            containers:
               {container}
-              serviceAccountName: {service_account}
-              nodeSelector:
-                cloud.google.com/gke-nodepool: cpu-user-np
-              hostNetwork: true
-              dnsPolicy: ClusterFirstWithHostNet
-              restartPolicy: Never
-              volumes:
-              - hostPath:
-                  path: /tmp
-                  type: DirectoryOrCreate
-                name: shared-tmp
-              {storage_volumes}"""
+            nodeSelector:
+              cloud.google.com/gke-nodepool: cpu-np
+            hostNetwork: true
+            dnsPolicy: ClusterFirstWithHostNet
+            restartPolicy: Never
+            volumes:
+            - hostPath:
+                path: /tmp
+                type: DirectoryOrCreate
+              name: shared-tmp
+    """
   if args.headless:
     return ''
   else:
     container, _ = get_user_workload_container(args, system)
-    storage_volumes = get_storage_volumes_yaml(storages)
     return user_workload_yaml.format(
         args=args,
         container=container,
-        storage_volumes=storage_volumes,
-        pod_failure_policy=pod_failure_policy,
-        service_account=XPK_SA,
-        gcs_fuse_annotation=GCS_FUSE_ANNOTATION,
     )
 
 
-def get_rm_address(args) -> str:
-  """Generates the Pathways resource manager address.
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    str: Fully qualified RM address.
-  """
-  rm_address = f'{args.workload}-rm-0-0.{args.workload}:29001'
-  return rm_address
-
-
 def get_proxy_address(args) -> str:
   """Generates the Pathways proxy address.
   Args:
@@ -354,24 +292,41 @@ def get_proxy_address(args) -> str:
   Returns:
     str: Fully qualified proxy address.
   """
-  proxy_address = f'grpc://{args.workload}-proxy-0-0.{args.workload}:29000'
+  proxy_address = (
+      f'grpc://{args.workload}-pathways-head-0-0.{args.workload}:29000'
+  )
   return proxy_address
 
 
-def get_pathways_expected_tpu_type(device_type: str) -> str:
-  """Returns the device type expected by Pathways
+def try_to_delete_pathwaysjob_first(args, workloads) -> bool:
+  """Function to delete PathwaysJob workload. This is needed as PathwaysJob
+  owns the JobSet it creates.
+
   Args:
-    device_type: the system characteristic device type
+    args: user provided arguments for running the command.
+    workloads: list of workloads that match the delete filter.
 
   Returns:
-    str: the device type expected by pathways.
+    True if successful and False otherwise.
   """
-  raw_type = device_type.split('-')[0].lower()
-  pathways_expected_instance = PathwaysExpectedInstancesMap[raw_type]
-  if not pathways_expected_instance:
-    xpk_print(
-        f'Passed in device_type {device_type} is incorrect. Please pass in a'
-        ' valid device type'
+  commands = []
+  task_names = []
+  for workload in workloads:
+    args.workload = workload
+    command = f'kubectl delete pathwaysjob {workload} -n default'
+    task_name = f'PathwaysWorkloadDelete-{workload}'
+    commands.append(command)
+    task_names.append(task_name)
+
+  # Not batching deletion for single workload
+  if len(workloads) == 1:
+    return_code = run_command_with_updates(commands[0], 'Delete Workload', args)
+  else:
+    return_code = run_commands(
+        commands, 'Delete Workload', task_names, batch=100
     )
-    xpk_exit(1)
-  return pathways_expected_instance
+
+  if return_code != 0:
+    xpk_print(f'Delete Workload request returned ERROR {return_code}')
+    return False
+  return True

xpk/core/scheduling.py

@@ -229,6 +229,21 @@ def create_accelerator_label(accelerator_type, system) -> str:
   )
 
 
+def create_tpu_machine_type(accelerator_type, system) -> str:
+  """Generates TPU machine type..
+
+  Args:
+    accelerator_type: type of accelerator.
+    system: system characteristics.
+
+  Returns:
+    The accelerator label.
+  """
+  if accelerator_type == AcceleratorType['TPU']:
+    return f'{system.gce_machine_type}'
+  return ''
+
+
 def create_machine_label(
     accelerator_type, system, autoprovisioning_enabled: bool = False
 ) -> str:
@@ -251,3 +266,24 @@ def create_machine_label(
         f' {system.topology}'
     )
   return ''
+
+
+def create_tpu_topology(
+    accelerator_type, system, autoprovisioning_enabled: bool = False
+) -> str:
+  """Generates TPU topology.
+
+  Args:
+    accelerator_type: type of accelerator.
+    system: system characteristics.
+    autoprovisioning_enabled: describes autoprovisioning enablement.
+
+  Returns:
+    The machine label.
+  """
+  if (
+      accelerator_type == AcceleratorType['TPU']
+      and not autoprovisioning_enabled
+  ):
+    return f'{system.topology}'
+  return ''

xpk/core/storage.py

@@ -45,8 +45,18 @@ STORAGE_CRD_PLURAL = "storages"
 STORAGE_CRD_NAME = f"{XPK_API_GROUP_NAME}.{STORAGE_CRD_PLURAL}"
 GCS_FUSE_TYPE = "gcsfuse"
 GCP_FILESTORE_TYPE = "gcpfilestore"
+PARALLELSTORE_TYPE = "parallelstore"
+GCE_PD_TYPE = "pd"
 MANIFESTS_PATH = os.path.abspath("xpkclusters/storage-manifests")
-GCS_FUSE_ANNOTATION = 'gke-gcsfuse/volumes: "true"'
+GCS_FUSE_ANNOTATIONS = {
+    "gke-gcsfuse/volumes": "true",
+    "gke-gcsfuse/cpu-limit": "0",
+    "gke-gcsfuse/memory-limit": "0",
+    "gke-gcsfuse/ephemeral-storage-limit": "0",
+}
+PARALLELSTORE_ANNOTATIONS = {
+    "gke-parallelstore/volumes": "true",
+}
 
 
 @dataclass
@@ -210,6 +220,24 @@ def get_auto_mount_gcsfuse_storages(k8s_api_client: ApiClient) -> list[Storage]:
   return list(filter(lambda storage: storage.type == GCS_FUSE_TYPE, storages))
 
 
+def get_auto_mount_parallelstore_storages(
+    k8s_api_client: ApiClient,
+) -> list[Storage]:
+  """
+  Retrieves all GCS Fuse Storage resources that have --auto-mount flag set to true.
+
+  Args:
+      k8s_api_client: An ApiClient object for interacting with the Kubernetes API.
+
+  Returns:
+      A list of GCS Fuse Storage objects that have `auto_mount` set to True.
+  """
+  storages: list[Storage] = get_auto_mount_storages(k8s_api_client)
+  return list(
+      filter(lambda storage: storage.type == PARALLELSTORE_TYPE, storages)
+  )
+
+
 def get_storages(
     k8s_api_client: ApiClient, requested_storages: list[str]
 ) -> list[Storage]:
@@ -314,6 +342,29 @@ def install_storage_crd(k8s_api_client: ApiClient) -> None:
       xpk_exit(1)
 
 
+def get_storage_annotations(storages: list[Storage]) -> list[str]:
+  """
+  Generates the storage annotations for workloads in the format of a YAML snippet.
+
+  Args:
+      storages: A list of Storage objects
+      offset: An integer specifying the depth of the YAML file
+
+  Returns:
+      A string containing the YAML representation of the storage annotations.
+  """
+  annotations = []
+  if any(storage.type == GCS_FUSE_TYPE for storage in storages):
+    for key, value in GCS_FUSE_ANNOTATIONS.items():
+      annotations.append(f'{key}: "{value}"')
+
+  if any(storage.type == PARALLELSTORE_TYPE for storage in storages):
+    for key, value in PARALLELSTORE_ANNOTATIONS.items():
+      annotations.append(f'{key}: "{value}"')
+
+  return annotations
+
+
 def get_storage_volume_mounts_yaml(storages: list[Storage]) -> str:
   """
   Generates the YAML representation of the volumeMounts section for the given Storages.
@@ -360,26 +411,29 @@ def get_storage_volumes_yaml(storages: list[Storage]) -> str:
   return yaml_str
 
 
-def get_storage_volume_mounts_yaml_for_gpu(storages: list[Storage]) -> str:
+def get_storage_volume_mounts_for_gpu(
+    storages: list[Storage],
+) -> list[dict]:
   """
   Generates the YAML representation of the volumeMounts section for the given Storages.
 
-  This function creates the YAML snippet that defines how the storage volumes
+  This function creates the list of storage specifications that define how the storage volumes
   should be mounted within a Pod's containers.
 
   Args:
       storages: A list of Storage objects.
 
   Returns:
-      A string containing the YAML representation of the volumeMounts section.
-  """
-  yaml_str = ""
-  for storage in storages:
-    yaml_str += f"""- name: {storage.pv}
-                  mountPath: {storage.mount_point}
-                  readOnly: {storage.readonly}
-            """
-  return yaml_str
+      A list containing the dictionary representation of the volumeMounts section.
+  """
+  return [
+      {
+          "name": storage.pv,
+          "mountPath": storage.mount_point,
+          "readOnly": storage.readonly,
+      }
+      for storage in storages
+  ]
 
 
 def get_storage_volumes_yaml_for_gpu(storages: list[Storage]) -> str:

xpk/core/system_characteristics.py

@@ -173,6 +173,15 @@ UserFacingNameToSystemCharacteristics = {
         AcceleratorType['GPU'],
         'a100-40gb-8',
     ),
+    'b200-8': SystemCharacteristics(
+        'N/A',
+        1,
+        'nvidia-b200',
+        'a4-highgpu-8g',
+        8,
+        AcceleratorType['GPU'],
+        'b200-8',
+    ),
     'h200-141gb-8': SystemCharacteristics(
         'N/A',
         1,