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

xpk/core/capacity.py

@@ -16,8 +16,8 @@ limitations under the License.
 
 import enum
 
-from ..utils.console import xpk_print
-from .commands import run_command_with_updates
+from ..utils.console import xpk_print, xpk_exit
+from .commands import run_command_with_updates, run_command_for_value
 
 AUTOPROVISIONING_CONFIG_VALUE = 'AUTOPROVISION'
 AUTOPROVISIONING_CONFIG_MINIMUM_KEY = 'minimum_chips'
@@ -36,6 +36,7 @@ class CapacityType(enum.Enum):
   RESERVATION = 'reservation'
   SPOT = 'spot'
   UNKNOWN = 'unknown'
+  FLEX_START = 'flex_start'
 
 
 def print_reservations(args) -> int:
@@ -84,6 +85,9 @@ def get_capacity_type(args) -> tuple[CapacityType, int]:
   if args.spot:
     capacity_type = CapacityType.SPOT
     num_types += 1
+  if args.flex:
+    capacity_type = CapacityType.FLEX_START
+    num_types += 1
 
   # Check that the number of user arguments provided is valid.
   if num_types == 0:
@@ -91,14 +95,62 @@ def get_capacity_type(args) -> tuple[CapacityType, int]:
   elif num_types != 1:
     xpk_print(
         'ERROR: User specified more than one of the following arguments. Please'
-        ' specify only one of `--reservation=$RESERVATION_NAME`, `--on-demand`'
-        ' or `--spot`.'
+        ' specify only one of `--reservation=$RESERVATION_NAME`, `--on-demand`,'
+        ' `--flex` or `--spot`.'
     )
     return_code = 1
 
   return capacity_type, return_code
 
 
+def get_reservation_maintenance_interval(
+    reservation: str, zone: str, project: str
+) -> str:
+  """Get reservation maintenance interval.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      f'gcloud beta compute reservations describe {reservation}'
+      f' --project={project} --zone={zone} --format="value(specificReservation.instanceProperties.maintenanceInterval)"'
+  )
+  return_code, output = run_command_for_value(
+      command, 'Get reservation maintenance interval', None
+  )
+  if return_code != 0:
+    xpk_print(f'Get reservation maintenance interval ERROR {return_code}')
+    xpk_exit(1)
+  return output.strip()
+
+
+def get_reservation_placement_policy(
+    reservation: str, zone: str, project: str
+) -> str:
+  """Get reservation placement policy.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      f'gcloud beta compute reservations describe {reservation}'
+      f' --project={project} --zone={zone} --format="value(resourcePolicies.policy)"'
+  )
+  return_code, output = run_command_for_value(
+      command, 'Get reservation placement policy', None
+  )
+  if return_code != 0:
+    xpk_print(f'Get reservation placement policy ERROR {return_code}')
+    xpk_exit(1)
+  return output.strip()
+
+
 def verify_reservation_exists(args) -> int:
   """Verify the reservation exists.
 
@@ -121,9 +173,9 @@ def verify_reservation_exists(args) -> int:
 
 
 def get_capacity_arguments_from_capacity_type(
-    args, capacity_type: CapacityType
+    args, capacity_type: CapacityType, max_nodes: int
 ) -> tuple[str, int]:
-  """Determine the TPU Nodepool creation capacity arguments needed.
+  """Determine the Nodepool creation capacity arguments needed.
 
   Args:
     args: user provided arguments for running the command.
@@ -141,6 +193,12 @@ def get_capacity_arguments_from_capacity_type(
       capacity_args = ''
     case CapacityType.SPOT:
       capacity_args = '--spot'
+    case CapacityType.FLEX_START:
+      capacity_args = (
+          ' --flex-start --enable-queued-provisioning --enable-autoscaling'
+          ' --location-policy=ANY --reservation-affinity=none'
+          f' --no-enable-autorepair --max-nodes={max_nodes}'
+      )
     case CapacityType.RESERVATION:
       capacity_args = (
           f'--reservation-affinity=specific --reservation={args.reservation}'
@@ -173,6 +231,8 @@ def get_capacity_node_selectors_from_capacity_type(
   match capacity_type:
     case CapacityType.ON_DEMAND.name:
       node_selector = ''
+    case CapacityType.FLEX_START.name:
+      node_selector = 'cloud.google.com/gke-queued="true"'
     case CapacityType.SPOT.name:
       node_selector = 'cloud.google.com/gke-spot="true"'
     case CapacityType.RESERVATION.name:

xpk/core/cluster.py

@@ -14,28 +14,37 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import yaml
 from google.api_core.exceptions import PermissionDenied
 from google.cloud import resourcemanager_v3
 from kubernetes import client as k8s_client
 from kubernetes import config
 from kubernetes.client.exceptions import ApiException
-from .resources import get_cluster_system_characteristics
 
 from ..utils.console import xpk_exit, xpk_print
-from .capacity import H100_DEVICE_TYPE
+from .capacity import B200_DEVICE_TYPE, H100_DEVICE_TYPE, H200_DEVICE_TYPE
 from .commands import (
     run_command_for_value,
     run_command_with_updates,
     run_command_with_updates_retry,
 )
-from .gcloud_context import add_zone_and_project, get_gke_server_config, zone_to_region
+from .gcloud_context import (
+    add_zone_and_project,
+    get_gke_server_config,
+    zone_to_region,
+)
 from .nodepool import upgrade_gke_nodepools_version
+from .resources import get_cluster_system_characteristics
 from .system_characteristics import SystemCharacteristics
 
 JOBSET_VERSION = 'v0.8.0'
-PATHWAYS_JOB_VERSION = 'v0.1.0'
-INSTALLER_NCC_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
-INSTALLER_NCC_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
+PATHWAYS_JOB_VERSION = 'v0.1.2'
+INSTALLER_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-tcpx-installer.yaml'
+INSTALLER_NCCL_TCPXO = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpxo/nccl-tcpxo-installer.yaml'
+INSTALLER_NCCL_RDMA = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-rdma/nccl-rdma-installer.yaml'
+CONFIG_NCCL_TCPX = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/gpudirect-tcpx/nccl-config.yaml'
+NRI_DEVICE_INJECTOR = 'https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nri_device_injector/nri-device-injector.yaml'
+MGLRU_DISABLE = 'https://raw.githubusercontent.com/GoogleCloudPlatform/cluster-toolkit/refs/heads/main/examples/gke-a3-ultragpu/mglru-disable.yaml'
 
 DEFAULT_NAMESPACE = 'default'
 XPK_SA = 'xpk-sa'
@@ -112,9 +121,11 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
     0 if successful and 1 otherwise.
   """
   if system.device_type == H100_DEVICE_TYPE:
-    command = f'kubectl apply -f {INSTALLER_NCC_TCPX}'
+    command = f'kubectl apply -f {INSTALLER_NCCL_TCPX}'
+  elif system.device_type in [H200_DEVICE_TYPE, B200_DEVICE_TYPE]:
+    command = f'kubectl apply -f {INSTALLER_NCCL_RDMA}'
   else:
-    command = f'kubectl apply -f {INSTALLER_NCC_TCPXO}'
+    command = f'kubectl apply -f {INSTALLER_NCCL_TCPXO}'
 
   return_code = run_command_with_updates(
       command, 'Install NCCL Plugin On Cluster', args
@@ -126,9 +137,108 @@ def install_nccl_on_cluster(args, system: SystemCharacteristics) -> int:
     )
     return 1
 
+  if system.device_type == H100_DEVICE_TYPE:
+    command = f'kubectl apply -f {CONFIG_NCCL_TCPX}'
+
+    return_code = run_command_with_updates(
+        command, 'Install NCCL Config On Cluster', args
+    )
+
+    if return_code != 0:
+      xpk_print(
+          f'Install NCCL Config On Cluster request returned ERROR {return_code}'
+      )
+      return 1
+
+  return 0
+
+
+def disable_mglru_on_cluster(args) -> int:
+  """Disable MGLRU on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = f'kubectl apply -f {MGLRU_DISABLE}'
+  return_code = run_command_with_updates(
+      command, 'Disable MGLRU On Cluster', args
+  )
+
+  if return_code != 0:
+    xpk_print('Disablig MGLRU On Cluster request returned ERROR')
+    return 1
+
+  return 0
+
+
+def install_nri_on_cluster(args) -> int:
+  """Install NRI Device Injector on the cluster.
+
+  Args:
+    args: user provided arguments for running the command.
+    system: system characteristics.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = f'kubectl apply -f {NRI_DEVICE_INJECTOR}'
+  return_code = run_command_with_updates(
+      command, 'Install NRI Device Injector On Cluster', args
+  )
+
+  if return_code != 0:
+    xpk_print(
+        'Install NRI Device Injector On Cluster request returned ERROR'
+        f' {return_code}'
+    )
+    return 1
+
   return 0
 
 
+def get_cluster_nodes_info(args) -> list[dict]:
+  """Get list of cluster's nodes descrition in yaml format
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    List of nodes info yaml objects.
+  """
+  xpk_print("Getting cluster's info...")
+  command = 'kubectl get nodes -o yaml'
+  err_code, val = run_command_for_value(
+      command=command,
+      task='Get cluster nodes info',
+      global_args=args,
+  )
+  if err_code != 0:
+    xpk_exit(err_code)
+  data = yaml.safe_load(val)
+  return data['items']  # pytype: disable=bad-return-type
+
+
+def count_nodes_on_cluster(args, system: SystemCharacteristics) -> int:
+  """Count cluster nodes by accelerator type"""
+  nodes_info = get_cluster_nodes_info(args)
+  accelerators = [
+      node['metadata']['labels']['cloud.google.com/gke-accelerator']
+      for node in nodes_info
+      if 'cloud.google.com/gke-accelerator' in node['metadata']['labels']
+  ]
+  if system.device_type != H200_DEVICE_TYPE:
+    xpk_print(
+        'Automatic node detection is not supported for device type:'
+        f' {system.device_type}'
+    )
+    xpk_exit(1)
+  num_nodes: int = sum(acc == system.gke_accelerator for acc in accelerators)
+  return num_nodes
+
+
 def get_cluster_network(args) -> str:
   xpk_print("Getting cluster's VPC network...")
   cluster_network_cmd = (
@@ -205,28 +315,60 @@ def update_cluster_with_pd_driver_if_necessary(args) -> int:
   return 0
 
 
-def is_driver_enabled_on_cluster(args, driver: str) -> bool:
+def update_cluster_with_lustre_driver_if_necessary(args) -> int:
+  """Updates a GKE cluster to enable Lustre CSI driver, if not enabled already.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and error code otherwise.
+  """
+  if is_driver_enabled_on_cluster(
+      args, driver='lustreCsiDriver'
+  ) and is_driver_enabled_on_cluster(
+      args, driver='lustreCsiDriver', config_key='enableLegacyLustrePort'
+  ):
+    return 0
+  cluster_update_return_code = update_gke_cluster_with_lustre_driver_enabled(
+      args
+  )
+  if cluster_update_return_code > 0:
+    xpk_print(
+        'Updating GKE cluster to enable PersistentDisk CSI driver failed!'
+    )
+    return cluster_update_return_code
+
+  return 0
+
+
+def is_driver_enabled_on_cluster(
+    args, driver: str, config_key: str = 'enabled', config_val: str = 'true'
+) -> bool:
   """Checks if the CSI driver is enabled on the cluster.
   Args:
     args: user provided arguments for running the command.
     driver (str) : name of the driver
+    config (str): the config to look for; by default looks for "enabled" parameter
+    config_val (str): the value indicating the enabled; default vale is "true"
   Returns:
     True if driver is enabled on the cluster and False otherwise.
   """
   command = (
       f'gcloud container clusters describe {args.cluster}'
       f' --project={args.project} --region={zone_to_region(args.zone)}'
-      f' --format="value(addonsConfig.{driver}Config.enabled)"'
+      f' --format="value(addonsConfig.{driver}Config.{config_key})"'
   )
   return_code, driver_enabled = run_command_for_value(
       command,
-      f'Checks if {driver} driver is enabled in cluster describe.',
+      f"Checks if {driver} driver's {config_key} is enabled in cluster"
+      ' describe.',
       args,
   )
   if return_code != 0:
     xpk_exit(return_code)
-  if driver_enabled.strip().lower() == 'true':
-    xpk_print(f'{driver} driver is enabled on the cluster, no update needed.')
+  if driver_enabled.strip().lower() == config_val.lower():
+    xpk_print(
+        f"{driver} driver's {config_key} config is {config_val} on the cluster."
+    )
     return True
   return False
 
@@ -313,6 +455,19 @@ def get_gpu_type_from_cluster(args) -> str:
   return ''
 
 
+def setup_k8s_service_accounts() -> None:
+  """
+  Creates/sets up SAs and the roles for them
+  """
+  default_sa = 'default'
+
+  create_xpk_k8s_service_account()
+
+  role_name = create_pod_reader_role()
+  create_role_binding(default_sa, role_name)
+  create_role_binding(XPK_SA, role_name)
+
+
 def create_xpk_k8s_service_account() -> None:
   k8s_core_client = k8s_client.CoreV1Api()
   sa = k8s_client.V1ServiceAccount(
@@ -331,6 +486,94 @@ def create_xpk_k8s_service_account() -> None:
     )
 
 
+def create_pod_reader_role() -> str:
+  """
+  Creates the 'pod-reader' Role in the default namespace.
+  """
+  k8s_rbac_client = k8s_client.RbacAuthorizationV1Api()
+  role_name = 'pod-reader'
+
+  role = k8s_client.V1Role(
+      metadata=k8s_client.V1ObjectMeta(
+          name=role_name, namespace=DEFAULT_NAMESPACE
+      ),
+      rules=[
+          k8s_client.V1PolicyRule(
+              api_groups=[''],
+              resources=['pods', 'services'],
+              verbs=['get', 'list', 'watch'],
+          ),
+          k8s_client.V1PolicyRule(
+              api_groups=['batch'],
+              resources=['jobs'],
+              verbs=['get', 'list', 'watch'],
+          ),
+      ],
+  )
+
+  xpk_print(
+      f'Attempting to create Role: {role_name} in namespace:'
+      f' {DEFAULT_NAMESPACE}'
+  )
+  try:
+    k8s_rbac_client.create_namespaced_role(DEFAULT_NAMESPACE, role, pretty=True)
+    xpk_print(f'Successfully created Role: {role_name}')
+    return role_name
+  except ApiException as e:
+    if e.status == 409:  # Conflict, meaning it already exists
+      xpk_print(f'Role: {role_name} already exists. Skipping its creation.')
+      return role_name
+    else:
+      xpk_print(f'Error creating Role {role_name}: {e}')
+      xpk_exit(1)
+
+
+def create_role_binding(sa: str, role_name: str) -> None:
+  """
+  Creates a RoleBinding to associate the Service Account
+  with the Role in the default namespace.
+  Assumes the Service Account and the Role already exist.
+  """
+  k8s_rbac_client = k8s_client.RbacAuthorizationV1Api()
+  role_binding_name = f'{sa}-{role_name}-binding'
+
+  role_binding = k8s_client.V1RoleBinding(
+      metadata=k8s_client.V1ObjectMeta(
+          name=role_binding_name, namespace=DEFAULT_NAMESPACE
+      ),
+      subjects=[
+          k8s_client.RbacV1Subject(
+              kind='ServiceAccount', name=sa, namespace=DEFAULT_NAMESPACE
+          )
+      ],
+      role_ref=k8s_client.V1RoleRef(
+          kind='Role', name=role_name, api_group='rbac.authorization.k8s.io'
+      ),
+  )
+
+  xpk_print(
+      f'Attempting to create RoleBinding: {role_binding_name} for Service'
+      f' Account: {XPK_SA} to Role: {role_name} in namespace:'
+      f' {DEFAULT_NAMESPACE}'
+  )
+  try:
+    k8s_rbac_client.create_namespaced_role_binding(
+        DEFAULT_NAMESPACE, role_binding, pretty=True
+    )
+    xpk_print(
+        f'Successfully created RoleBinding: {role_binding_name} for {XPK_SA}'
+    )
+  except ApiException as e:
+    if e.status == 409:  # Conflict, meaning it already exists
+      xpk_print(
+          f'RoleBinding: {role_binding_name} already exists. Skipping its'
+          ' creation.'
+      )
+    else:
+      xpk_print(f'Error creating RoleBinding {role_binding_name}: {e}')
+      xpk_exit(1)
+
+
 def update_gke_cluster_with_clouddns(args) -> int:
   """Run the GKE cluster update command for existing clusters and enable CloudDNS.
 
@@ -412,6 +655,32 @@ def update_gke_cluster_with_gcsfuse_driver_enabled(args) -> int:
   return 0
 
 
+def update_gke_cluster_with_lustre_driver_enabled(args) -> int:
+  """Run the GKE cluster update command for existing cluster and enable Lustre CSI driver.
+  Args:
+    args: user provided arguments for running the command.
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  command = (
+      'gcloud container clusters update'
+      f' {args.cluster} --project={args.project}'
+      f' --region={zone_to_region(args.zone)}'
+      ' --enable-legacy-lustre-port'
+      ' --quiet'
+  )
+  xpk_print(
+      'Updating GKE cluster to enable Lustre CSI driver, may take a while!'
+  )
+  return_code = run_command_with_updates(
+      command, 'GKE Cluster Update to enable Lustre CSI driver', args
+  )
+  if return_code != 0:
+    xpk_print(f'GKE Cluster Update request returned ERROR {return_code}')
+    return 1
+  return 0
+
+
 def upgrade_gke_control_plane_version(args, default_rapid_gke_version) -> int:
   """Upgrade GKE cluster's control plane version before updating nodepools to use CloudDNS.
 

xpk/core/config.py

@@ -15,16 +15,14 @@ limitations under the License.
 """
 
 import os
-import re
 
 import ruamel.yaml
 
 from ..utils import file
 from ..utils.console import xpk_print
-from .system_characteristics import AcceleratorType, SystemCharacteristics
 
 # This is the version for XPK PyPI package
-__version__ = 'v0.8.0'
+__version__ = 'v0.10.0'
 XPK_CURRENT_VERSION = __version__
 XPK_CONFIG_FILE = os.path.expanduser('~/.config/xpk/config.yaml')
 
@@ -117,65 +115,3 @@ class XpkConfig:
       return None
     val: dict[str, str] = config_yaml[CONFIGS_KEY]
     return val
-
-
-def parse_env_config(args, tensorboard_config, system: SystemCharacteristics):
-  """Parses the environment configurations to the jobset config.
-
-  Args:
-    args: user provided arguments for running the command.
-    tensorboard_config: configuration of Vertex Tensorboard.
-    system: system characteristics.
-  """
-  env = {}
-
-  env_pat = re.compile(r'(^[a-zA-Z_][a-zA-Z0-9_]*?)(?:=(.*))?$', re.M)
-  if args.env_file:
-    print('Setting container environment from', args.env_file)
-    with open(file=args.env_file, mode='r', encoding='utf-8') as f:
-      for match in env_pat.finditer(f.read()):
-        variable = match.group(1)
-        if match.group(2) is not None:
-          env[variable] = match.group(2)
-        else:
-          assert variable in os.environ, (
-              f'Variable {variable} is not set in the current '
-              'environment, a value must be specified.'
-          )
-          env[variable] = os.environ[variable]
-  if args.env:
-    for var in args.env:
-      match = env_pat.match(var)
-      assert match and match.group(2) is not None, (
-          'Invalid environment variable, format must be '
-          f'`--env VARIABLE=value`: {var}'
-      )
-      variable = match.group(1)
-      env[variable] = match.group(2)
-
-  if not args.use_pathways:
-    if args.debug_dump_gcs:
-      if 'XLA_FLAGS' in env:
-        raise ValueError(
-            'Conflict: XLA_FLAGS defined in both --debug_dump_gcs '
-            'and environment file. Please choose one way to define '
-            'XLA_FLAGS.'
-        )
-      env['XLA_FLAGS'] = '--xla_dump_to=/tmp/xla_dump/'
-
-    if tensorboard_config:
-      env['UPLOAD_DATA_TO_TENSORBOARD'] = True
-      for key, value in tensorboard_config.items():
-        env[key.upper()] = value
-
-  if system.accelerator_type == AcceleratorType['GPU']:
-    # For GPUs, it has two more spaces ahead of name and value respectively
-    env_format = '''
-                  - name: {key}
-                    value: "{value}"'''
-  else:
-    env_format = '''
-                - name: {key}
-                  value: "{value}"'''
-
-  args.env = ''.join(env_format.format(key=k, value=v) for k, v in env.items())

xpk/core/docker_manager.py

@@ -30,7 +30,7 @@ import time
 DockerRunCommandExitCode = 135
 dockerBuildErrorCode = 134
 ctk_dockerfile_path = "Dockerfile"
-ctk_build_ref = "v1.48.0"
+ctk_build_ref = "v1.57.1"
 ctk_docker_image = "xpk-ctk"
 ctk_container_name = "xpk-ctk-container"
 gcloud_cfg_mount_path = "/root/.config/gcloud"