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

xpk/core/gcloud_context.py

@@ -0,0 +1,196 @@
+"""
+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.
+"""
+
+import subprocess
+import sys
+from dataclasses import dataclass
+
+from ..utils.console import xpk_print
+from .commands import run_command_for_value
+
+
+def get_project():
+  """Get GCE project from `gcloud config get project`.
+
+  Returns:
+     The project name.
+  """
+  completed_command = subprocess.run(
+      ['gcloud', 'config', 'get', 'project'], check=True, capture_output=True
+  )
+  project_outputs = completed_command.stdout.decode().strip().split('\n')
+  if len(project_outputs) < 1 or project_outputs[-1] == '':
+    sys.exit(
+        'You must specify the project in the project flag or set it with'
+        " 'gcloud config set project <project>'"
+    )
+  return project_outputs[
+      -1
+  ]  # The project name lives on the last line of the output
+
+
+def get_zone():
+  """Get GCE zone from `gcloud config get compute/zone`.
+
+  Returns:
+     The zone name.
+  """
+  completed_command = subprocess.run(
+      ['gcloud', 'config', 'get', 'compute/zone'],
+      check=True,
+      capture_output=True,
+  )
+  zone_outputs = completed_command.stdout.decode().strip().split('\n')
+  if len(zone_outputs) < 1 or zone_outputs[-1] == '':
+    sys.exit(
+        "You must specify the zone in the zone flag or set it with 'gcloud"
+        " config set compute/zone <zone>'"
+    )
+  return zone_outputs[-1]  # The zone name lives on the last line of the output
+
+
+def add_zone_and_project(args):
+  """Obtains the zone and project names from gcloud configs if not defined.
+
+  Args:
+    args: user provided arguments for running the command.
+  """
+  if not args.project:
+    args.project = get_project()
+  if not args.zone:
+    args.zone = get_zone()
+  xpk_print(f'Working on {args.project} and {args.zone}')
+
+
+def zone_to_region(zone) -> str:
+  """Helper function converts zone name to region name.
+
+  Args:
+    zone: zone name.
+
+  Returns:
+     The region name.
+  """
+  zone_terms = zone.split('-')
+  return zone_terms[0] + '-' + zone_terms[1]  # pytype: disable=bad-return-type
+
+
+@dataclass
+class GkeServerConfig:
+  """Stores the valid gke versions based on gcloud recommendations."""
+
+  default_rapid_gke_version: str
+  valid_versions: set[str]
+
+
+def get_gke_server_config(args) -> tuple[int, GkeServerConfig | None]:
+  """Determine the GKE versions supported by gcloud currently.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    Tuple of
+    int: 0 if successful and 1 otherwise.
+    GkeServerConfig: stores valid gke version to use in node pool and cluster.
+  """
+  base_command = (
+      'gcloud container get-server-config'
+      f' --project={args.project} --region={zone_to_region(args.zone)}'
+  )
+  default_rapid_gke_version_cmd = (
+      base_command
+      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      ' --format="value(channels.defaultVersion)"'
+  )
+  valid_versions_cmd = (
+      base_command
+      + ' --flatten="channels" --filter="channels.channel=RAPID"'
+      ' --format="value(channels.validVersions)"'
+  )
+  base_command_description = 'Determine server supported GKE versions for '
+
+  server_config_commands_and_descriptions = [
+      (
+          default_rapid_gke_version_cmd,
+          base_command_description + 'default rapid gke version',
+      ),
+      (
+          valid_versions_cmd,
+          base_command_description + 'valid versions',
+      ),
+  ]
+  command_outputs = []
+
+  for command, command_description in server_config_commands_and_descriptions:
+    return_code, cmd_output = run_command_for_value(
+        command,
+        command_description,
+        args,
+        hide_error=True,
+    )
+    if return_code != 0:
+      xpk_print(f'Unable to get server config for {command_description}.')
+      return return_code, None
+    command_outputs.append(cmd_output)
+
+  return 0, GkeServerConfig(
+      default_rapid_gke_version=command_outputs[0].strip(),
+      valid_versions=set(command_outputs[1].split(';')),
+  )
+
+
+def get_gke_control_plane_version(
+    args, gke_server_config: GkeServerConfig
+) -> tuple[int, str | None]:
+  """Determine gke control plane version for cluster creation.
+
+  Args:
+    args: user provided arguments for running the command.
+    gke_server_config: holds valid gke versions and recommended default version.
+
+  Returns:
+    Tuple of
+    int: 0 if successful and 1 otherwise.
+    str: gke control plane version to use.
+  """
+
+  # Override with user provide gke version if specified.
+  if args.gke_version is not None:
+    master_gke_version = args.gke_version
+  else:
+    master_gke_version = gke_server_config.default_rapid_gke_version
+
+  is_valid_version = master_gke_version in gke_server_config.valid_versions
+
+  if not is_valid_version:
+    xpk_print(
+        f'Planned GKE Version: {master_gke_version}\n Valid Versions:'
+        f'\n{gke_server_config.valid_versions}\nRecommended / Default GKE'
+        f' Version: {gke_server_config.default_rapid_gke_version}'
+    )
+    xpk_print(
+        f'Error: Planned GKE Version {master_gke_version} is not valid.'
+        f'Checks failed: Is Version Valid: {is_valid_version}'
+    )
+    xpk_print(
+        'Please select a gke version from the above list using --gke-version=x'
+        ' argument or rely on the default gke version:'
+        f' {gke_server_config.default_rapid_gke_version}'
+    )
+    return 1, None
+
+  return 0, master_gke_version

xpk/core/gcluster_manager.py

@@ -15,8 +15,8 @@ limitations under the License.
 """
 
 from .docker_manager import CommandRunner
-from ..utils.console import xpk_print
-
+from ..utils.console import xpk_exit, xpk_print
+from .remote_state.remote_state_client import RemoteStateClient
 
 xpk_gcloud_cfg_path = '~/gcloud/cfg'
 xpk_deployment_dir = '/deployment'
@@ -44,8 +44,10 @@ class GclusterManager:
   def __init__(
       self,
       gcluster_command_runner: CommandRunner,
+      remote_state_client: RemoteStateClient | None,
   ) -> None:
     self.gcluster_command_runner = gcluster_command_runner
+    self.remote_state_client = remote_state_client
 
   def _run_create_deployment_cmd(
       self, blueprint_container_path: str, prefix: str = ''
@@ -156,3 +158,19 @@ class GclusterManager:
     xpk_print('Staging blueprint completed!')
     xpk_print(f"File path in gcluster's working directory: {staged_blueprint}")
     return staged_blueprint
+
+  def upload_state(self) -> None:
+    xpk_print('Uploading state.')
+    if self.remote_state_client is None:
+      xpk_print('No remote state defined')
+      xpk_exit(1)
+    self.remote_state_client.upload_state()
+
+  def download_state(self) -> None:
+    if self.remote_state_client is None:
+      xpk_print('No remote state defined')
+      xpk_exit(1)
+
+    if self.remote_state_client.check_remote_state_exists():
+      self.remote_state_client.download_state()
+    xpk_print('Remote state not found.')

xpk/core/gcsfuse.py

@@ -0,0 +1,50 @@
+"""
+Copyright 2024 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 ..utils import templates
+
+FUSE_PV_PATH = "/../templates/fuse-pv.yaml"
+FUSE_PVC_PATH = "/../templates/fuse-pvc.yaml"
+
+
+def create_pv(name: str, size: int, bucket: str) -> dict:
+  data = templates.load(FUSE_PV_PATH)
+  data["metadata"]["name"] = f"{name}-pv"
+  data["spec"]["capacity"]["storage"] = f"{size}Gi"
+  data["spec"]["csi"]["volumeHandle"] = bucket
+  return data
+
+
+def create_pvc(name: str, size: int) -> dict:
+  data = templates.load(FUSE_PVC_PATH)
+  data["metadata"]["name"] = f"{name}-pvc"
+  data["spec"]["resources"]["requests"]["storage"] = f"{size}Gi"
+  data["spec"]["volumeName"] = f"{name}-pv"
+  return data
+
+
+def manifest(name: str, bucket: str, size: int) -> list[dict]:
+  """Creates GCS FUSE manifest file.
+
+  Args:
+      path (str): path to the file where the manifest will be created
+      name (str): base name of the volumes
+      bucket (str): name of the storage bucket
+      size (str): size of the storage
+  """
+  pv = create_pv(name, size, bucket)
+  pvc = create_pvc(name, size)
+  return [pv, pvc]

xpk/core/kjob.py

@@ -14,11 +14,33 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from ..core.blueprint.blueprint_generator import get_subnetworks_for_a3mega, get_subnetworks_for_a3ultra
+from ..core.capacity import H100_MEGA_DEVICE_TYPE, H200_DEVICE_TYPE
 from argparse import Namespace
-from ..utils.console import xpk_print
+import yaml
+from .workload_decorators.tcpxo_decorator import get_tcpxo_deamon_entry
+from ..utils.console import xpk_print, xpk_exit
+
+from ..utils import templates
+from kubernetes import client as k8s_client
+from kubernetes.client import ApiClient
+from kubernetes.client.rest import ApiException
+from .cluster import setup_k8s_env, XPK_SA, DEFAULT_NAMESPACE
+from .storage import get_auto_mount_storages, get_auto_mount_gcsfuse_storages
 from .commands import run_command_for_value, run_kubectl_apply, run_command_with_updates
+from .config import XpkConfig, KJOB_SHELL_IMAGE, KJOB_SHELL_INTERACTIVE_COMMAND, KJOB_SHELL_WORKING_DIRECTORY, KJOB_BATCH_IMAGE, KJOB_BATCH_WORKING_DIRECTORY
+from .resources import get_cluster_system_characteristics, SystemCharacteristics, AcceleratorType
 from enum import Enum
 
+from ..core.workload_decorators import tcpxo_decorator
+
+from ..core.workload_decorators import rdma_decorator
+
+KJOB_API_GROUP_NAME = "kjobctl.x-k8s.io"
+KJOB_API_GROUP_VERSION = "v1alpha1"
+KJOB_API_VOLUME_BUNDLE_PLURAL = "volumebundles"
+VOLUME_BUNDLE_TEMPLATE_PATH = "/../templates/volume_bundle.yaml"
+
 
 class AppProfileDefaults(Enum):
   NAME = "xpk-def-app-profile"
@@ -30,12 +52,14 @@ class JobTemplateDefaults(Enum):
   COMPLETIONS = 1
   CONTAINER_NAME = "xpk-batch-container"
   IMAGE = "ubuntu:22.04"
+  WORKING_DIRECTORY = "/"
 
 
 class PodTemplateDefaults(Enum):
   NAME = "xpk-def-pod"
   CONTAINER_NAME = "xpk-interactive-container"
   IMAGE = "busybox:1.28"
+  WORKING_DIRECTORY = "/"
   INTERACTIVE_COMMAND = "/bin/sh"
 
 
@@ -52,10 +76,29 @@ job_template_yaml = """
       completionMode: Indexed
       template:
         spec:
+          dnsPolicy: ClusterFirstWithHostNet
+          tolerations:
+            - operator: "Exists"
+              key: nvidia.com/gpu
           containers:
             - name: {container_name}
               image: {image}
-          restartPolicy: OnFailure"""
+              workingDir: {working_directory}
+              {resources}
+          {node_selector}
+          priorityClassName: {priority}
+          restartPolicy: OnFailure
+          serviceAccountName: {service_account}
+"""
+job_node_selector_template = """
+          nodeSelector:
+            cloud.google.com/gke-accelerator: {gpu_name}
+"""
+job_resources_template = """
+              resources:
+                limits:
+                  nvidia.com/gpu: {gpu_per_node}
+"""
 
 app_profile_yaml = """
 apiVersion: kjobctl.x-k8s.io/v1alpha1
@@ -70,6 +113,7 @@ spec:
       requiredFlags: []
     - name: Interactive
       template: {interactive_template}
+  volumeBundles: {volume_bundles}
 """
 
 pod_template_yaml = """
@@ -80,12 +124,53 @@ metadata:
   namespace: default
 template:
   spec:
+    tolerations:
+      - effect: NoSchedule
+        key: components.gke.io/gke-managed-components
+        operator: Equal
+        value: "true"
     containers:
       - name: {container_name}
         image: {image}
         command: [{interactive_command}]
+        workingDir: {working_directory}
+    initContainers:
+      - name: init
+        image: {image}
+        command: ['/bin/mkdir', '-p', '{working_directory}']
+    serviceAccountName: {service_account}
 """
 
+Kueue_TAS_annotation = "kueue.x-k8s.io/podset-preferred-topology=cloud.google.com/gce-topology-host"
+
+default_interface_annotation = "networking.gke.io/default-interface=eth0"
+
+
+def get_a3ultra_pod_template_annotations(args: Namespace) -> tuple[str, str]:
+  sub_networks = get_subnetworks_for_a3ultra(args.cluster)
+  interfaces_key, interfaces_value = rdma_decorator.get_interfaces_entry(
+      sub_networks
+  )
+
+  return (
+      default_interface_annotation,
+      f"{interfaces_key}=$'{interfaces_value}'",
+  )
+
+
+def get_a3mega_pod_template_annotations(
+    args: Namespace,
+) -> tuple[str, str, str]:
+  """Adds or updates annotations in the Pod template."""
+  sub_networks = get_subnetworks_for_a3mega(args.cluster)
+  tcpxo_deamon_key, tcpxo_deamon_paths = get_tcpxo_deamon_entry()
+  interfaces_key, interfaces_value = tcpxo_decorator.get_interfaces_entry(
+      sub_networks
+  )
+  tcpxo = f"{tcpxo_deamon_key}=$'{tcpxo_deamon_paths}'"
+  interfaces = f"{interfaces_key}=$'{interfaces_value}'"
+  return tcpxo, interfaces, default_interface_annotation
+
 
 def verify_kjob_installed(args: Namespace) -> int:
   """Check if kjob is installed. If not provide user with proper communicate and exit.
@@ -112,7 +197,25 @@ def verify_kjob_installed(args: Namespace) -> int:
   return 0
 
 
-def create_app_profile_instance(args: Namespace) -> int:
+def get_pod_template_interactive_command() -> str:
+  """Gets the interactive command for PodTemplate from config otherwise the default value.
+
+  Args:
+    args - user provided arguments
+  Returns:
+    str - PodTemplate's interactive command
+  """
+  config = XpkConfig()
+  pod_command = config.get(KJOB_SHELL_INTERACTIVE_COMMAND)
+  if pod_command is None or len(pod_command) == 0:
+    pod_command = PodTemplateDefaults.INTERACTIVE_COMMAND.value
+
+  return pod_command
+
+
+def create_app_profile_instance(
+    args: Namespace, volume_bundles: list[str]
+) -> int:
   """Create new AppProfile instance on cluster with default settings.
 
   Args:
@@ -125,13 +228,29 @@ def create_app_profile_instance(args: Namespace) -> int:
           name=AppProfileDefaults.NAME.value,
           batch_template=JobTemplateDefaults.NAME.value,
           interactive_template=PodTemplateDefaults.NAME.value,
+          volume_bundles=volume_bundles,
       ),
       task="Creating AppProfile",
       args=args,
   )
 
 
-def create_job_template_instance(args: Namespace) -> int:
+def decorate_job_template_with_gpu(yml_string: str, gpu_type: str) -> str:
+  job_spec = yaml.safe_load(yml_string)["template"]
+  if gpu_type == H100_MEGA_DEVICE_TYPE:
+    job_spec = tcpxo_decorator.decorate_kjob_template(job_spec)
+  if gpu_type == H200_DEVICE_TYPE:
+    job_spec = rdma_decorator.decorate_kjob_template(job_spec)
+  job_template_dict = yaml.safe_load(yml_string)
+  job_template_dict["template"] = job_spec
+  return yaml.dump(job_template_dict, sort_keys=False)
+
+
+def create_job_template_instance(
+    args: Namespace,
+    system: SystemCharacteristics | None,
+    service_account: str,
+) -> int:
   """Create new JobTemplate instance on cluster with default settings.
 
   Args:
@@ -139,20 +258,49 @@ def create_job_template_instance(args: Namespace) -> int:
   Returns:
     exit_code > 0 if creating JobTemplate fails, 0 otherwise
   """
+  config = XpkConfig()
+  job_image = config.get(KJOB_BATCH_IMAGE)
+  if job_image is None or len(job_image) == 0:
+    job_image = JobTemplateDefaults.IMAGE.value
+  working_directory = config.get(KJOB_BATCH_WORKING_DIRECTORY)
+  if working_directory is None or len(working_directory) == 0:
+    working_directory = JobTemplateDefaults.WORKING_DIRECTORY.value
+  resources = (
+      job_resources_template.format(gpu_per_node=system.chips_per_vm)
+      if system is not None
+      and system.accelerator_type == AcceleratorType["GPU"]
+      else ""
+  )
+
+  node_selector = (
+      job_node_selector_template.format(gpu_name=system.gke_accelerator)
+      if system is not None
+      and system.accelerator_type == AcceleratorType["GPU"]
+      else ""
+  )
+  yml_string = job_template_yaml.format(
+      name=JobTemplateDefaults.NAME.value,
+      parallelism=JobTemplateDefaults.PARALLELISM.value,
+      completions=JobTemplateDefaults.COMPLETIONS.value,
+      container_name=JobTemplateDefaults.CONTAINER_NAME.value,
+      image=job_image,
+      working_directory=working_directory,
+      resources=resources,
+      node_selector=node_selector,
+      priority=args.priority if hasattr(args, "priority") else "medium",
+      service_account=service_account,
+  )
+  if system is not None and system.accelerator_type == AcceleratorType["GPU"]:
+    yml_string = decorate_job_template_with_gpu(yml_string, system.device_type)
+
   return run_kubectl_apply(
-      yml_string=job_template_yaml.format(
-          name=JobTemplateDefaults.NAME.value,
-          parallelism=JobTemplateDefaults.PARALLELISM.value,
-          completions=JobTemplateDefaults.COMPLETIONS.value,
-          container_name=JobTemplateDefaults.CONTAINER_NAME.value,
-          image=JobTemplateDefaults.IMAGE.value,
-      ),
+      yml_string,
       task="Creating JobTemplate",
       args=args,
   )
 
 
-def create_pod_template_instance(args: Namespace) -> int:
+def create_pod_template_instance(args: Namespace, service_account: str) -> int:
   """Create new PodTemplate instance on cluster with default settings.
 
   Args:
@@ -160,28 +308,49 @@ def create_pod_template_instance(args: Namespace) -> int:
   Returns:
     exit_code > 0 if creating PodTemplate fails, 0 otherwise
   """
+  config = XpkConfig()
+  pod_image = config.get(KJOB_SHELL_IMAGE)
+  if pod_image is None or len(pod_image) == 0:
+    pod_image = PodTemplateDefaults.IMAGE.value
+  working_directory = config.get(KJOB_SHELL_WORKING_DIRECTORY)
+  if working_directory is None or len(working_directory) == 0:
+    working_directory = PodTemplateDefaults.WORKING_DIRECTORY.value
+
   return run_kubectl_apply(
       yml_string=pod_template_yaml.format(
           name=PodTemplateDefaults.NAME.value,
           container_name=PodTemplateDefaults.CONTAINER_NAME.value,
-          image=PodTemplateDefaults.IMAGE.value,
-          interactive_command=PodTemplateDefaults.INTERACTIVE_COMMAND.value,
+          image=pod_image,
+          working_directory=working_directory,
+          interactive_command=get_pod_template_interactive_command(),
+          service_account=service_account,
       ),
       task="Creating PodTemplate",
       args=args,
   )
 
 
-def prepare_kjob(args) -> int:
-  job_err_code = create_job_template_instance(args)
+def prepare_kjob(args: Namespace) -> int:
+  system = get_cluster_system_characteristics(args)
+
+  k8s_api_client = setup_k8s_env(args)
+  storages = get_auto_mount_storages(k8s_api_client)
+
+  service_account = ""
+  if len(storages) > 0:
+    service_account = XPK_SA
+
+  job_err_code = create_job_template_instance(args, system, service_account)
   if job_err_code > 0:
     return job_err_code
 
-  pod_err_code = create_pod_template_instance(args)
+  pod_err_code = create_pod_template_instance(args, service_account)
   if pod_err_code > 0:
     return pod_err_code
 
-  return create_app_profile_instance(args)
+  volume_bundles = [item.name for item in storages]
+
+  return create_app_profile_instance(args, volume_bundles)
 
 
 def apply_kjob_crds(args: Namespace) -> int:
@@ -203,3 +372,73 @@ def apply_kjob_crds(args: Namespace) -> int:
     return return_code
   xpk_print("Creating kjob CRDs succeeded")
   return 0
+
+
+def create_volume_bundle_instance(
+    k8s_api_client: ApiClient,
+    name: str,
+    manifest: list[dict],
+    readonly: bool,
+    mount_point: str,
+) -> None:
+  """
+  Creates a new VolumeBundle resource in the Kubernetes cluster.
+
+  This function reads a VolumeBundle template from a YAML file, populates it with
+  values from the provided arguments, and then creates the VolumeBundle object
+  in the cluster.
+
+  Args:
+      k8s_api_client: An ApiClient object for interacting with the Kubernetes API.
+      args: An argparse Namespace object containing the arguments for creating
+            the Storage resource.
+  """
+  data = templates.load(VOLUME_BUNDLE_TEMPLATE_PATH)
+  data["metadata"]["name"] = name
+  spec = data["spec"]
+  spec["volumes"] = []
+  spec["containerVolumeMounts"] = []
+
+  for obj in manifest:
+    if obj["kind"] == "PersistentVolumeClaim":
+      spec["volumes"].append({
+          "name": obj["metadata"]["name"],
+          "persistentVolumeClaim": {
+              "claimName": obj["metadata"]["name"],
+              "readOnly": readonly,
+          },
+      })
+      spec["containerVolumeMounts"].append({
+          "name": obj["metadata"]["name"],
+          "mountPath": mount_point,
+      })
+
+  data["spec"] = spec
+
+  api_instance = k8s_client.CustomObjectsApi(k8s_api_client)
+  try:
+    api_instance.create_namespaced_custom_object(
+        namespace=DEFAULT_NAMESPACE,
+        group=KJOB_API_GROUP_NAME,
+        version=KJOB_API_GROUP_VERSION,
+        plural=KJOB_API_VOLUME_BUNDLE_PLURAL,
+        body=data,
+    )
+    xpk_print(
+        f"Created {KJOB_API_VOLUME_BUNDLE_PLURAL}.{KJOB_API_GROUP_NAME} object:"
+        f" {data['metadata']['name']}"
+    )
+  except ApiException as e:
+    if e.status == 409:
+      xpk_print(f"VolumeBundle: {name} already exists. Skipping its creation")
+    else:
+      xpk_print(f"Encountered error during VolumeBundle creation: {e}")
+      xpk_exit(1)
+
+
+def get_gcsfuse_annotation(args: Namespace) -> str | None:
+  k8s_api_client = setup_k8s_env(args)
+  gcsfuse_storages = get_auto_mount_gcsfuse_storages(k8s_api_client)
+  if len(gcsfuse_storages) > 0:
+    return "gke-gcsfuse/volumes=true"
+  return None