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

xpk/core/pathways.py

@@ -14,13 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from .cluster import XPK_SA
+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 .core import (
-    AcceleratorType,
-    get_all_nodepools_programmatic,
-    get_user_workload_container,
-    zone_to_region,
-)
+from .config import AcceleratorType
+from .storage import Storage, get_storage_volumes_yaml, GCS_FUSE_ANNOTATION
 from .system_characteristics import SystemCharacteristics
 
 PathwaysExpectedInstancesMap = {
@@ -41,9 +41,11 @@ def get_pathways_worker_args(args) -> str:
     str: yaml containing arguments for the Pathways workers.
   """
   yaml = """- --server_port=29001
-              - --resource_manager_address={rm_address}
-              - --gcs_scratch_location={args.pathways_gcs_location}"""
+                - --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 ''
@@ -58,15 +60,53 @@ def get_pathways_proxy_args(args) -> str:
     str: yaml containing arguments for the Pathways proxy.
   """
   yaml = """- --server_port=29000
-              - --resource_manager_address={rm_address}
-              - --gcs_scratch_location={args.pathways_gcs_location}"""
+                - --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
@@ -105,9 +145,9 @@ def add_pw_resources_to_kueue(args):
     - name: cpu-rm
       resources:
       - name: "cpu"
-        nominalQuota: 80
+        nominalQuota: 480
       - name: "memory"
-        nominalQuota: 160G
+        nominalQuota: 2000G
     - name: cpu-proxy
       resources:
       - name: "cpu"
@@ -166,23 +206,17 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool:
   # Set the job which determines the life of other Pathways jobs
   args.targetReplicatedJob = 'proxy' if args.headless else 'main'
 
-  # Always report user code failures back to JobSet.
-  args.restart_on_user_code_failure = True
-
   return True
 
 
 def get_pathways_unified_query_link(args) -> str:
   """Get the unified query link for the pathways workload."""
-  pw_suffixes = ['main', 'rm', 'proxy']
-  pw_pod_names = [f'"{args.workload}-{suffix}-0"' for suffix in pw_suffixes]
-  pw_pod_names_query = '%20OR%20'.join(pw_pod_names + ['worker-0-0'])
   query_params = (
       'resource.type%3D"k8s_container"%0A'
       f'resource.labels.project_id%3D"{args.project}"%0A'
       f'resource.labels.location%3D"{zone_to_region(args.zone)}"%0A'
       f'resource.labels.cluster_name%3D"{args.cluster}"%0A'
-      f'resource.labels.pod_name:{pw_pod_names_query}%0A'
+      f'resource.labels.pod_name:"{args.workload}-"%0A'
       'severity>%3DDEFAULT'
   )
 
@@ -198,11 +232,13 @@ def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
     str: yaml containing arguments for the Pathways resource manager.
   """
   yaml = """- --server_port=29001
-              - --gcs_scratch_location={args.pathways_gcs_location}
-              - --node_type=resource_manager
-              - --instance_count={instance_count}
-              - --instance_type={instance_type}"""
+                - --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,
@@ -212,7 +248,34 @@ def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
     return ''
 
 
-def get_user_workload_for_pathways(args, system: SystemCharacteristics) -> str:
+def append_custom_pathways_args(yaml, custom_args) -> str:
+  """Append custom Pathways args to the YAML with proper indentation.
+
+  Args:
+      yaml (string): existing yaml containing args
+
+  Returns:
+      yaml (string): yaml with additional args 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
+  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.
   Don't create one for Pathways headless mode.
@@ -227,32 +290,48 @@ def get_user_workload_for_pathways(args, system: SystemCharacteristics) -> 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
-        template:
-          spec:
-            containers:
+      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:
               {container}
-            nodeSelector:
-              cloud.google.com/gke-nodepool: cpu-user-np
-            restartPolicy: OnFailure
-            volumes:
-            - hostPath:
-                path: /tmp
-                type: DirectoryOrCreate
-              name: shared-tmp"""
+              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}"""
   if args.headless:
     return ''
   else:
     container, _ = get_user_workload_container(args, system)
-    return user_workload_yaml.format(args=args, container=container)
+    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:

xpk/core/remote_state/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright 2025 Google LLC
+
+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.
+"""

xpk/core/remote_state/fuse_remote_state.py

@@ -0,0 +1,99 @@
+"""
+Copyright 2025 Google LLC
+
+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.
+"""
+
+from .remote_state_client import RemoteStateClient
+from ...utils.gcs_utils import upload_directory_to_gcs, check_file_exists, download_bucket_to_dir, upload_file_to_gcs
+from ...utils.console import xpk_print
+from google.cloud.storage import Client
+import os
+
+
+class FuseStateClient(RemoteStateClient):
+  """FuseStateClient is a class for managing remote xpk state stored in GCS Fuse."""
+
+  def __init__(
+      self,
+      bucket: str,
+      state_directory: str,
+      cluster: str,
+      deployment_name: str,
+      prefix: str,
+  ) -> None:
+    self.bucket = bucket
+    self.state_dir = state_directory
+    self.storage_client = Client()
+    self.cluster = cluster
+    self.prefix = prefix
+    self.deployment_name = deployment_name
+
+  def _get_bucket_path(self) -> str:
+    return (
+        f'xpk_terraform_state/{self.prefix}/blueprints/{self.deployment_name}/'
+    )
+
+  def _get_bucket_path_blueprint(self) -> str:
+    return f'xpk_terraform_state/{self.prefix}/blueprints/'
+
+  def _get_deployment_filename(self) -> str:
+    return f'{self.deployment_name}.yaml'
+
+  def _get_blueprint_path(self) -> str:
+    blueprint_dir = '/'.join(self.state_dir.split('/')[:-1])
+    return os.path.join(blueprint_dir, self.deployment_name) + '.yaml'
+
+  def upload_state(self) -> None:
+    xpk_print(
+        f'Uploading dependecies from directory {self.state_dir} to bucket:'
+        f' {self.bucket}. Path within bucket is: {self._get_bucket_path()}'
+    )
+    upload_directory_to_gcs(
+        storage_client=self.storage_client,
+        bucket_name=self.bucket,
+        bucket_path=self._get_bucket_path(),
+        source_directory=self.state_dir,
+    )
+    blueprint_bucket_path = (
+        self._get_bucket_path_blueprint() + self._get_deployment_filename()
+    )
+    xpk_print(
+        f'Uploading blueprint file: {self._get_blueprint_path()} to bucket'
+        f' {self.bucket}. Path within bucket is: {blueprint_bucket_path}'
+    )
+    upload_file_to_gcs(
+        storage_client=self.storage_client,
+        bucket_name=self.bucket,
+        bucket_path=blueprint_bucket_path,
+        file=self._get_blueprint_path(),
+    )
+
+  def download_state(self) -> None:
+    xpk_print(
+        f'Downloading from bucket: {self.bucket}, from path:'
+        f' {self._get_bucket_path()} to directory: {self.state_dir}'
+    )
+    download_bucket_to_dir(
+        self.storage_client,
+        self.bucket,
+        self._get_bucket_path(),
+        destination_directory=self.state_dir,
+    )
+
+  def check_remote_state_exists(self) -> bool:
+    return check_file_exists(
+        self.storage_client,
+        self.bucket,
+        self._get_bucket_path_blueprint() + self._get_deployment_filename(),
+    )

xpk/core/remote_state/remote_state_client.py

@@ -0,0 +1,38 @@
+"""
+Copyright 2025 Google LLC
+
+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.
+"""
+
+from abc import ABC, abstractmethod
+
+
+class RemoteStateClient(ABC):
+  """This is a base class that defines methods a class for managing remote cluster state.
+  Args:
+      ABC (_type_): _description_
+  """
+
+  @abstractmethod
+  def upload_state(self) -> None:
+    """Upload state to remote storage"""
+    return None
+
+  @abstractmethod
+  def download_state(self) -> None:
+    """Download state from remote storage"""
+    return None
+
+  @abstractmethod
+  def check_remote_state_exists(self) -> bool:
+    return False

xpk/core/resources.py

@@ -0,0 +1,238 @@
+"""
+Copyright 2025 Google LLC
+
+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.
+"""
+
+from dataclasses import dataclass
+
+from ..utils.console import xpk_print
+from ..utils.file import write_tmp_file
+from .capacity import (
+    AUTOPROVISIONING_CONFIG_MAXIMUM_KEY,
+    AUTOPROVISIONING_CONFIG_MINIMUM_KEY,
+    AUTOPROVISIONING_CONFIG_VALUE,
+    CAPACITY_TYPE_CONFIG_KEY,
+    RESERVATION_CONFIG_KEY,
+    CapacityType,
+    get_capacity_type,
+)
+from .commands import run_command_for_value, run_commands
+from .config import XPK_CURRENT_VERSION
+from .system_characteristics import AcceleratorType, get_system_characteristics_by_device_type, SystemCharacteristics
+
+CLUSTER_RESOURCES_CONFIGMAP = 'resources-configmap'
+CLUSTER_METADATA_CONFIGMAP = 'metadata-configmap'
+
+CLUSTER_CONFIGMAP_YAML = """kind: ConfigMap
+apiVersion: v1
+metadata:
+  name: {name}
+data:
+  {data}
+"""
+
+
+@dataclass
+class AutoprovisioningConfig:
+  config_filename: str
+  minimum_chips: int
+  maximum_chips: int
+
+
+def get_cluster_configmap(args, configmap_name) -> dict[str, str] | None:
+  """Run the Get GKE Cluster ConfigMap request.
+
+  Args:
+    args: user provided arguments for running the command.
+    configmap_name: name of the configmap.
+
+  Returns:
+    key:value pairs stored in cluster ConfigMap.
+  """
+  command = (
+      'kubectl get configmap'
+      f' {configmap_name} -o=custom-columns="ConfigData:data" --no-headers=true'
+  )
+
+  return_code, return_value = run_command_for_value(
+      command, 'GKE Cluster Get ConfigMap', args
+  )
+  if return_code != 0:
+    xpk_print(f'GKE Cluster Get ConfigMap request returned ERROR {return_code}')
+    return None
+
+  config_map = {}
+  return_value = return_value.strip()
+
+  if return_value:
+    # Format of ConfigMap: map[key1:value1 key2:value2]
+    return_value = return_value[return_value.index('map') :]
+    configs = return_value[4:-1].split(' ')
+
+    for config in configs:
+      key, value = config.strip().split(':')
+      config_map[key] = value
+  return config_map
+
+
+def create_cluster_configmaps(
+    args,
+    system,
+    tensorboard_config: dict,
+    autoprovisioning_config: AutoprovisioningConfig | None,
+) -> int:
+  """Run the Create GKE Cluster ConfigMap request.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+    tensorboard_config: map that contains Vertex Tensorboard name, id and location
+    autoprovisioning_config: Config used in autoprovisioning.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  configmap_yml = {}
+
+  # ConfigMap to store resources available in the cluster.
+  device_type = system.device_type
+  if system.accelerator_type == AcceleratorType['GPU']:
+    resources_data = f'{device_type}: "{int(args.num_nodes)}"'
+  elif (
+      not args.enable_pathways
+      and args.enable_autoprovisioning
+      and autoprovisioning_config
+  ):
+    # Currently autoprovisioning is not supported with Pathways.
+    # Auto provisioning will have variable topologies for a gke accelerator type.
+    resources_data = (
+        f'{system.gke_accelerator}: {AUTOPROVISIONING_CONFIG_VALUE}'
+    )
+    resources_data += (
+        f'\n  {AUTOPROVISIONING_CONFIG_MINIMUM_KEY}:'
+        f' "{autoprovisioning_config.minimum_chips}"'
+    )
+    resources_data += (
+        f'\n  {AUTOPROVISIONING_CONFIG_MAXIMUM_KEY}:'
+        f' "{autoprovisioning_config.maximum_chips}"'
+    )
+  else:
+    resources_data = (
+        f'{device_type}: "{int(args.num_slices) * system.vms_per_slice}"'
+    )
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  resources_yml = CLUSTER_CONFIGMAP_YAML.format(
+      args=args, name=resources_configmap_name, data=resources_data
+  )
+  configmap_yml[resources_configmap_name] = resources_yml
+
+  # ConfigMap to store cluster metadata.
+  # XPK Version.
+  metadata = f'xpk_version: {XPK_CURRENT_VERSION}'
+  # Vertex Tensorboard information
+  for key, value in tensorboard_config.items():
+    metadata += f'\n  {key}: "{value}"'
+  # Capacity Type.
+  capacity_type, return_code = get_capacity_type(args)
+  if return_code != 0:
+    xpk_print('Unable to determine capacity type.')
+    return return_code
+  metadata += f'\n  {CAPACITY_TYPE_CONFIG_KEY}: {capacity_type.name}'
+  # Reservation ID if applicable.
+  if capacity_type == CapacityType.RESERVATION:
+    metadata += f'\n  {RESERVATION_CONFIG_KEY}: {args.reservation}'
+  metadata_configmap_name = f'{args.cluster}-{CLUSTER_METADATA_CONFIGMAP}'
+  metadata_yml = CLUSTER_CONFIGMAP_YAML.format(
+      args=args, name=metadata_configmap_name, data=metadata
+  )
+  configmap_yml[metadata_configmap_name] = metadata_yml
+  return create_or_update_cluster_configmap(configmap_yml)
+
+
+def create_or_update_cluster_configmap(configmap_yml: dict) -> int:
+  """
+  Args:
+    configmap_yml: dict containing ConfigMap name and yml string.
+
+  Returns:
+    0 if successful, 1 otherwise.
+  """
+  commands = []
+  task_names = []
+  for configmap_name, yml_string in configmap_yml.items():
+    tmp = write_tmp_file(yml_string)
+    command = f'kubectl apply -f {str(tmp.file.name)}'
+    commands.append(command)
+    task_name = f'ConfigMap CreateOrUpdate-{configmap_name}'
+    task_names.append(task_name)
+
+  return_code = run_commands(
+      commands, 'GKE Cluster CreateOrUpdate ConfigMap(s)', task_names
+  )
+  if return_code != 0:
+    xpk_print(
+        'GKE Cluster Create/Update ConfigMap(s) request returned ERROR'
+        f' {return_code}'
+    )
+    return 1
+  return 0
+
+
+def check_cluster_resources(args, system) -> tuple[bool, bool]:
+  """Check if cluster has resources of a specified device_type/gke_accelerator.
+  This check will be skipped if <args.cluster>-<_CLUSTER_RESOURCES_CONFIGMAP> ConfigMap doesn't exist for the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    Tuple of bool, bool
+    True if resources in the cluster should be checked, False otherwise.
+    True if device_type/gke_accelerator exists in the cluster, False otherwise.
+  """
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  resources_config_map = get_cluster_configmap(args, resources_configmap_name)
+  if resources_config_map is None:
+    xpk_print(
+        f'No ConfigMap exist for cluster with the name {resources_config_map}.'
+        ' Cluster resources check will be skipped.'
+    )
+    return False, False
+  if system.device_type in resources_config_map:
+    return True, True
+  elif system.gke_accelerator in resources_config_map:
+    return True, True
+  return True, False
+
+
+def get_cluster_system_characteristics(args) -> SystemCharacteristics | None:
+  """Get systemCharcteristics based on the cluster resources configMap
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    returns system characteristics
+  """
+  resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
+  cluster_config_map = get_cluster_configmap(args, resources_configmap_name)
+
+  if cluster_config_map is None:
+    return None
+
+  for key in cluster_config_map:
+    system, result_code = get_system_characteristics_by_device_type(key)
+    if result_code == 0:
+      return system
+
+  return None