xpk 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

{xpk-0.4.0.dist-info → xpk-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: xpk
-Version: 0.4.0
+Version: 0.5.0
 Summary: xpk helps Cloud developers to orchestrate training jobs on accelerators on GKE.
 Author-email: Cloud TPU Team <cloud-tpu-eng@google.com>
 License: Apache-2.0
@@ -139,14 +139,6 @@ gcloud config set compute/zone $ZONE
 xpk .. --zone $ZONE --project $PROJECT_ID
 ```
 
-`Cluster Create` command will create a project-specific Service Account. Note that only one service 
-account will be created per project. This service account will be attached to the node pools instead of default 
-[Compute Engine Service Account](https://cloud.google.com/compute/docs/access/service-accounts#default_service_account). 
-All the required permissions will be assigned to this service account by XPK. Make sure you have 
-[Service Account Admin](https://cloud.google.com/iam/docs/understanding-roles#iam.serviceAccountAdmin) and 
-[Project IAM Admin](https://cloud.google.com/iam/docs/understanding-roles#resourcemanager.projectIamAdmin) 
-roles assigned to your user account.
-
 The cluster created is a regional cluster to enable the GKE control plane across
 all zones.
 
@@ -226,7 +218,9 @@ all zones.
     ```
 
 ### Create Vertex AI Tensorboard
-*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature.*
+*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature. Make sure you have
+[Vertex AI Administrator](https://cloud.google.com/vertex-ai/docs/general/access-control#aiplatform.admin) role
+assigned to your user account.*
 
 Vertex AI Tensorboard is a fully managed version of open-source Tensorboard. To learn more about Vertex AI Tensorboard, visit [this](https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-introduction). Note that Vertex AI Tensorboard is only available in [these](https://cloud.google.com/vertex-ai/docs/general/locations#available-regions) regions.
 
@@ -386,7 +380,9 @@ checkpointing so the job restarts near where it was interrupted.
   ```
 
 ### Create Vertex AI Experiment to upload data to Vertex AI Tensorboard
-*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature.*
+*Note: This feature is available in XPK >= 0.4.0. Enable [Vertex AI API](https://cloud.google.com/vertex-ai/docs/start/cloud-environment#enable_vertexai_apis) in your Google Cloud console to use this feature. Make sure you have
+[Vertex AI Administrator](https://cloud.google.com/vertex-ai/docs/general/access-control#aiplatform.admin) role
+assigned to your user account and to the [Compute Engine Service account](https://cloud.google.com/compute/docs/access/service-accounts#default_service_account) attached to the node pools in the cluster.*
 
 Vertex AI Experiment is a tool that helps to track and analyze an experiment run on Vertex AI Tensorboard. To learn more about Vertex AI Experiments, visit [this](https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments).
 

xpk-0.5.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+xpk.py,sha256=83W6YcVkv0x5X05DRdJ-Q_ycgruzBVsAErNPnfbxE9I,217498
+xpk-0.5.0.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+xpk-0.5.0.dist-info/METADATA,sha256=STLdsmG_h-dShES4OTEt-o-hDZ8SwJ9Q9yTy1pPgcEk,44142
+xpk-0.5.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+xpk-0.5.0.dist-info/entry_points.txt,sha256=lhrMqkTA09DLePaqxSMyW2RCLUKs2X1c84baGhMev_k,33
+xpk-0.5.0.dist-info/top_level.txt,sha256=aDe4N0jicmuWExx_6w0TxWQJaEuPSs9BnLU-3aF1GLo,4
+xpk-0.5.0.dist-info/RECORD,,

xpk.py

@@ -67,9 +67,8 @@ if (
 
 default_docker_image = 'python:3.10'
 default_script_dir = os.getcwd()
-default_gke_version = '1.29.1-gke.1589017'
 # This is the version for XPK PyPI package
-__version__ = '0.4.0'
+__version__ = '0.5.0'
 xpk_current_version = __version__
 
 h100_device_type = 'h100-80gb-8'
@@ -85,10 +84,7 @@ _LOCAL_QUEUE_NAME = 'multislice-queue'
 _DEFAULT_POOL_NAME = 'default-pool'
 _CLUSTER_RESOURCES_CONFIGMAP = 'resources-configmap'
 _CLUSTER_METADATA_CONFIGMAP = 'metadata-configmap'
-_XPK_SERVICE_ACCOUNT = 'xpk-sa'
-# Set to True to attach a service account to cluster & node pools
-_SERVICE_ACCOUNT_FEATURE_FLAG = xpk_current_version >= '0.4.0'
-_VERTEX_TENSORBOARD_FEATURE_FLAG = _SERVICE_ACCOUNT_FEATURE_FLAG
+_VERTEX_TENSORBOARD_FEATURE_FLAG = xpk_current_version >= '0.4.0'
 _DEFAULT_VERTEX_TENSORBOARD_NAME = 'tb-instance'
 
 
@@ -264,11 +260,12 @@ spec:
         labels:
           xpk.google.com/workload: {args.workload}
       spec:
-        backoffLimit: 0
+        backoffLimit: 4
         completions: {system.vms_per_slice}
         parallelism: {system.vms_per_slice}
         template:
           spec:
+            terminationGracePeriodSeconds: {args.termination_grace_period_seconds}
             containers:
             - args:
               {pathways_worker_args}
@@ -611,7 +608,6 @@ management:
 autoprovisioningLocations:
   {zones}
 {resource_limits}
-{service_account}
 """
 
 autoprovisioning_resource_limits = """
@@ -629,16 +625,6 @@ autoprovisioning_custom_resource_type = """
   maximum: {maximum}
 """
 
-# Add IAM roles to attach to service account used by node pools in the cluster
-IAMRoles = {
-    'Kubernetes Engine Admin': 'roles/container.admin',
-    'Artifact Registry Writer': 'roles/artifactregistry.writer',
-    'Monitoring Admin': 'roles/monitoring.admin',
-    'Logging Admin': 'roles/logging.admin',
-    'Storage Admin': 'roles/storage.admin',
-    'Vertex AI Administrator': 'roles/aiplatform.admin',
-}
-
 
 AcceleratorType = {'TPU': 1, 'GPU': 2, 'CPU': 3}
 
@@ -1891,6 +1877,14 @@ the corresponding Map in MaxText/accelerator_to_spec_map.py """
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 
+PathwaysExpectedInstancesMap = {
+    'v5p': 'v5',
+    'v5litepod': 'v5e',
+    'v4': 'v4',
+    'v3': 'v3',
+}
+
+
 def chunks(lst, n):
   """Return a list of n-sized chunks from lst.
 
@@ -2162,7 +2156,12 @@ def append_temporary_file(payload, file):
 
 
 def run_command_for_value(
-    command, task, global_args, dry_run_return_val='0', print_timer=False
+    command,
+    task,
+    global_args,
+    dry_run_return_val='0',
+    print_timer=False,
+    hide_error=False,
 ) -> tuple[int, str]:
   """Runs the command and returns the error code and stdout.
 
@@ -2173,6 +2172,8 @@ def run_command_for_value(
     task: user provided task name for running the command.
     global_args: user provided arguments for running the command.
     dry_run_return_val: return value of this command for dry run.
+    print_timer: print out the time the command is running.
+    hide_error: hide the error from the command output upon success.
 
   Returns:
     tuple[int, str]
@@ -2216,7 +2217,7 @@ def run_command_for_value(
       output = subprocess.check_output(
           command,
           shell=True,
-          stderr=subprocess.STDOUT,
+          stderr=subprocess.STDOUT if not hide_error else None,
       )
     except subprocess.CalledProcessError as e:
       xpk_print(f'Task {task} failed with {e.returncode}')
@@ -2397,142 +2398,6 @@ def zone_to_region(zone) -> str:
   return zone_terms[0] + '-' + zone_terms[1]
 
 
-def get_service_account_name(args) -> str:
-  """Get the name for the service account.
-  Args:
-    args: user provided arguments.
-
-  Returns:
-    the name of the service account.
-  """
-  return f'{args.project}-{_XPK_SERVICE_ACCOUNT}@{args.project}.iam.gserviceaccount.com'
-
-
-def check_if_service_account_exists(args) -> bool:
-  """Check if a service account with the given name exists in the project.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    True if service account exist, False otherwise.
-  """
-  service_account_name = get_service_account_name(args)
-  command = f'gcloud iam service-accounts describe {service_account_name}'
-  return_code = run_command_with_updates(
-      command, 'Service Account Describe', args, verbose=False
-  )
-  if return_code != 0:
-    xpk_print(
-        'Service Account Describe did not find the service account'
-        f' {service_account_name}.'
-    )
-    return False
-  return True
-
-
-def create_service_account(args) -> int:
-  """Creates a service account in the project.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  command = (
-      'gcloud iam service-accounts create'
-      f' {args.project}-{_XPK_SERVICE_ACCOUNT}     --description="Service'
-      ' Account for XPK"    '
-      f' --display-name="{args.project}-{_XPK_SERVICE_ACCOUNT}"'
-  )
-  return_code = run_command_with_updates(
-      command, 'Service Account Create', args
-  )
-  if return_code != 0:
-    xpk_print(f'Service Account Create request returned ERROR {return_code}')
-    xpk_print(
-        'Make sure you have Service Account Admin Role attached to your user.'
-    )
-    return 1
-  return 0
-
-
-def get_existing_roles_in_service_account(args) -> set:
-  """
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    set of IAM roles already attached to the service account.
-  """
-  roles = set()
-  service_account_name = get_service_account_name(args)
-  command = (
-      f'gcloud projects get-iam-policy {args.project}'
-      f' --filter="bindings.members:{service_account_name}"'
-      ' --flatten="bindings[].members" --format="table(bindings.role)"'
-  )
-  return_code, return_value = run_command_for_value(
-      command, 'Get IAM Roles For Service Account', args
-  )
-  if return_code != 0:
-    xpk_print(
-        'Get IAM Roles For Service Account request returned ERROR'
-        f' {return_code}'
-    )
-  else:
-    return_value = return_value.strip()
-    roles = set(return_value.split('\n'))
-    """Format of return_value is:
-          ROLE
-          roles/storage.admin
-          roles/logging.admin
-      removing `ROLE` from the list
-    """
-    if 'ROLE' in roles:
-      roles.remove('ROLE')
-  return roles
-
-
-def add_roles_to_service_account(args) -> int:
-  """Add IAM roles to service account.
-
-  Args:
-    args: user provided arguments for running the command.
-
-  Returns:
-    0 if successful and 1 otherwise.
-  """
-  service_account_name = get_service_account_name(args)
-  existing_roles = get_existing_roles_in_service_account(args)
-  xpk_print(f'IAM roles already attached to service account: {existing_roles}')
-
-  for name, role in IAMRoles.items():
-    if role in existing_roles:
-      continue
-
-    xpk_print(f'Adding {name} role to service account: {service_account_name}.')
-    command = (
-        f'gcloud projects add-iam-policy-binding {args.project}      '
-        f' --member="serviceAccount:{service_account_name}"      '
-        f' --role="{role}"       --condition=None'
-    )
-    return_code = run_command_with_updates(
-        command, 'Add IAM Role to Service Account', args, verbose=False
-    )
-    if return_code != 0:
-      xpk_print(
-          'Add IAM Role to Service Account request returned ERROR'
-          f' {return_code}'
-      )
-      xpk_print(
-          'Make sure you have Project IAM Admin Role attached to your user.'
-      )
-      return 1
-  return 0
-
-
 def get_total_chips_requested_from_args(
     args, system: SystemCharacteristics
 ) -> int:
@@ -2631,16 +2496,7 @@ def create_autoprovisioning_config(
       custom_resource_type=custom_resource_string,
   )
 
-  # Default service_account is the project's default service account.
-  service_account = ''
-  if _SERVICE_ACCOUNT_FEATURE_FLAG:
-    service_account_name = get_service_account_name(args)
-    service_account_exists = check_if_service_account_exists(args)
-    if service_account_exists:
-      service_account = f'serviceAccount: {service_account_name}'
-
   yml_string = autoprovisioning_config_file.format(
-      service_account=service_account,
       resource_limits=resource_limits,
       zones=f'- {args.zone}',
   )
@@ -2743,11 +2599,12 @@ def enable_autoprovisioning_on_cluster(
   return autoprovisioning_config, return_code
 
 
-def run_gke_cluster_create_command(args) -> int:
+def run_gke_cluster_create_command(args, gke_control_plane_version: str) -> int:
   """Run the Create GKE Cluster request.
 
   Args:
     args: user provided arguments for running the command.
+    gke_control_plane_version: version used if creating the cluster.
 
   Returns:
     0 if successful and 1 otherwise.
@@ -2772,7 +2629,7 @@ def run_gke_cluster_create_command(args) -> int:
       f' {args.cluster} --project={args.project}'
       f' --region={zone_to_region(args.zone)}'
       f' --node-locations={args.zone}'
-      f' --cluster-version={args.gke_version}'
+      f' --cluster-version={gke_control_plane_version}'
       f' --machine-type={machine_type}'
       ' --enable-autoscaling'
       ' --total-min-nodes 1 --total-max-nodes 1000'
@@ -2780,17 +2637,6 @@ def run_gke_cluster_create_command(args) -> int:
       f' {args.custom_cluster_arguments}'
   )
 
-  if _SERVICE_ACCOUNT_FEATURE_FLAG:
-    service_account_name = get_service_account_name(args)
-    service_account_exists = check_if_service_account_exists(args)
-    if service_account_exists:
-      command += f' --service-account={service_account_name}'
-    else:
-      xpk_print(
-          f'Service Account: {service_account_name} does not exist in the'
-          ' project. Will attach the default service account to the cluster.'
-      )
-
   device_type = args.tpu_type if args.tpu_type else args.device_type
   if device_type == h100_device_type:
     command += (
@@ -3376,11 +3222,12 @@ def get_all_clusters_programmatic(args) -> tuple[list[str], int]:
   return cluster_names, 0
 
 
-def create_cluster_if_necessary(args) -> int:
+def create_cluster_if_necessary(args, gke_control_plane_version: str) -> int:
   """Creates cluster if not present in the project.
 
   Args:
     args: user provided arguments for running the command.
+    gke_control_plane_version: version used if creating the cluster.
 
   Returns:
     0 if successful and 1 otherwise.
@@ -3390,10 +3237,10 @@ def create_cluster_if_necessary(args) -> int:
     xpk_print('Listing all clusters failed!')
     return 1
   if args.cluster in all_clusters:
-    xpk_print('Skipping cluster creation since it already exists')
+    xpk_print('Skipping cluster creation since it already exists.')
     return 0
   else:
-    return run_gke_cluster_create_command(args)
+    return run_gke_cluster_create_command(args, gke_control_plane_version)
 
 
 def get_all_nodepools_programmatic(args) -> tuple[list[str], int]:
@@ -3499,12 +3346,15 @@ def get_user_input(input_msg):
   return user_input in ('y', 'yes')
 
 
-def run_gke_node_pool_create_command(args, system) -> int:
+def run_gke_node_pool_create_command(
+    args, system, gke_node_pool_version
+) -> int:
   """Run the Create GKE Node Pool request.
 
   Args:
     args: user provided arguments for running the command.
     system: System characteristics based on device type/topology.
+    gke_node_pool_version: GKE version to use to create node pools.
 
   Returns:
     0 if successful and 1 otherwise.
@@ -3574,10 +3424,12 @@ def run_gke_node_pool_create_command(args, system) -> int:
         f' {args.custom_nodepool_arguments}'
     )
     if system.accelerator_type == AcceleratorType['TPU']:
-      command += f' --node-version={args.gke_version}'
+      command += f' --node-version={gke_node_pool_version}'
       command += f' --num-nodes={system.vms_per_slice}'
       command += ' --placement-type=COMPACT  --max-pods-per-node 15'
-      command += ' --scopes=storage-full,gke-default'
+      command += (
+          ' --scopes=storage-full,gke-default,"https://www.googleapis.com/auth/cloud-platform"'
+      )
       command += f' --tpu-topology={system.topology}'
       command += f' {args.custom_tpu_nodepool_arguments}'
     elif system.accelerator_type == AcceleratorType['GPU']:
@@ -3601,17 +3453,6 @@ def run_gke_node_pool_create_command(args, system) -> int:
       command += f' --num-nodes={system.vms_per_slice}'
       command += ' --scopes=storage-full,gke-default'
 
-    if _SERVICE_ACCOUNT_FEATURE_FLAG:
-      service_account_name = get_service_account_name(args)
-      service_account_exists = check_if_service_account_exists(args)
-      if service_account_exists:
-        command += f' --service-account={service_account_name}'
-      else:
-        xpk_print(
-            f'Service Account: {service_account_name} does not exist in the'
-            ' project. Will attach the default service account to the node'
-            ' pools.'
-        )
     task = f'NodepoolCreate-{node_pool_name}'
     commands.append(command)
     task_names.append(task)
@@ -3624,7 +3465,7 @@ def run_gke_node_pool_create_command(args, system) -> int:
         continue
       command = (
           'gcloud beta container node-pools create'
-          f' {node_pool_name} --node-version={args.gke_version}'
+          f' {node_pool_name} --node-version={gke_node_pool_version}'
           f' --cluster={args.cluster}'
           f' --project={args.project} --node-locations={args.zone}'
           f' --region={zone_to_region(args.zone)}'
@@ -4060,6 +3901,184 @@ def install_nccl_on_cluster(args) -> int:
   return 0
 
 
+@dataclass
+class GkeServerConfig:
+  """Stores the valid gke versions based on gcloud recommendations."""
+
+  default_rapid_gke_version: str
+  valid_master_versions: set[str]
+  valid_node_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_master_versions_cmd = (
+      base_command
+      + ' --flatten="channels" --format="value(validMasterVersions)"'
+  )
+  valid_node_versions_cmd = (
+      base_command + ' --flatten="channels" --format="value(validNodeVersions)"'
+  )
+  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_master_versions_cmd,
+          base_command_description + 'valid master versions',
+      ),
+      (
+          valid_node_versions_cmd,
+          base_command_description + 'valid node 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_master_versions=set(command_outputs[1].split(';')),
+      valid_node_versions=set(command_outputs[2].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_master_version = (
+      master_gke_version in gke_server_config.valid_master_versions
+  )
+  is_valid_node_version = (
+      master_gke_version in gke_server_config.valid_node_versions
+  )
+
+  if not is_valid_master_version or not is_valid_node_version:
+    xpk_print(
+        f'Planned GKE Version: {master_gke_version}\n Valid Master'
+        f' Versions:\n{gke_server_config.valid_master_versions}\nValid Node'
+        f' Versions:\n{gke_server_config.valid_node_versions}\nRecommended 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 Master Valid: {is_valid_master_version}'
+        f'\nIs Valid Node Version: {is_valid_node_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
+
+
+def get_gke_node_pool_version(
+    args, gke_server_config: GkeServerConfig
+) -> tuple[int, str | None]:
+  """Determine the gke node pool version for the node pool.
+
+  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.
+  """
+
+  # By default use the current gke master version for creating node pools.
+  command_description = 'Determine current gke master version'
+  command = (
+      f'gcloud beta container clusters describe {args.cluster}'
+      f' --region {zone_to_region(args.zone)} --project {args.project}'
+      ' --format="value(currentMasterVersion)"'
+  )
+
+  return_code, current_gke_master_version = run_command_for_value(
+      command, command_description, args
+  )
+  if return_code != 0:
+    xpk_print(
+        f'Unable to get server config for command: {command_description}.'
+    )
+    return return_code, None
+
+  # Override with user provide gke version if specified.
+  if args.gke_version is not None:
+    node_pool_gke_version = args.gke_version
+  else:
+    node_pool_gke_version = current_gke_master_version.strip()
+
+  is_supported_node_pool_version = (
+      node_pool_gke_version in gke_server_config.valid_node_versions
+  )
+  # In rare cases, user's provided gke version may be invalid, but gke will return an error if so.
+  # An example scenario is if the user provided gke version is greater than the master version.
+  if not is_supported_node_pool_version:
+    xpk_print(
+        f'Planned node pool version {node_pool_gke_version} is not supported in'
+        ' valid node_pool_gke_versions'
+        f' {gke_server_config.valid_node_versions}Please adjust the gke version'
+        ' using --gke-version=x or remove the arg and depend on xpk default of'
+        f' {current_gke_master_version}'
+    )
+    return 1, None
+  return 0, node_pool_gke_version
+
+
 ################### Subcommand Functions ###################
 def default_subcommand_function(
     _args,
@@ -4097,26 +4116,19 @@ def cluster_create(args) -> int:
   xpk_print(f'Starting cluster create for cluster {args.cluster}:', flush=True)
   add_zone_and_project(args)
 
-  if _SERVICE_ACCOUNT_FEATURE_FLAG:
-    service_account_name = get_service_account_name(args)
-    service_account_exists = check_if_service_account_exists(args)
-    if service_account_exists:
-      xpk_print(
-          f'Service Account: {service_account_name} already exist in the'
-          ' project. Will not create a new service account.'
-      )
-    else:
-      # create a service account in the project
-      create_service_account_code = create_service_account(args)
-      if create_service_account_code != 0:
-        xpk_exit(create_service_account_code)
+  return_code, gke_server_config = get_gke_server_config(args)
+  if return_code != 0:
+    xpk_exit(return_code)
 
-    # add IAM roles to the service account
-    add_roles_to_service_account_code = add_roles_to_service_account(args)
-    if add_roles_to_service_account_code != 0:
-      xpk_exit(add_roles_to_service_account_code)
+  return_code, gke_control_plane_version = get_gke_control_plane_version(
+      args, gke_server_config
+  )
+  if return_code != 0:
+    xpk_exit(return_code)
 
-  create_cluster_command_code = create_cluster_if_necessary(args)
+  create_cluster_command_code = create_cluster_if_necessary(
+      args, gke_control_plane_version
+  )
   if create_cluster_command_code != 0:
     xpk_exit(create_cluster_command_code)
 
@@ -4144,8 +4156,16 @@ def cluster_create(args) -> int:
     if create_cluster_network_config_code != 0:
       xpk_exit(create_cluster_network_config_code)
 
+  # Check the control plane version of the cluster and determine the node pool
+  # version to use.
+  return_code, gke_node_pool_version = get_gke_node_pool_version(
+      args, gke_server_config
+  )
+  if return_code != 0:
+    xpk_exit(return_code)
+
   run_gke_node_pool_create_command_code = run_gke_node_pool_create_command(
-      args, system
+      args, system, gke_node_pool_version
   )
   if run_gke_node_pool_create_command_code != 0:
     xpk_exit(run_gke_node_pool_create_command_code)
@@ -4818,7 +4838,7 @@ def get_volume_mounts(args, system: SystemCharacteristics) -> str:
   return regular_volume_mount_yaml
 
 
-def get_pathways_rm_args(args) -> str:
+def get_pathways_rm_args(args, system: SystemCharacteristics) -> str:
   """Arguments for the Pathways resource manager.
   Args:
     args: user provided arguments for running the command.
@@ -4833,13 +4853,56 @@ def get_pathways_rm_args(args) -> str:
               - --pathways_persistent_compilation_cache=false
               - --pathways_compilation_mode=compile_at_worker
               - --pathways_tmp_dir_pattern={args.pathways_gcs_location}
-              - --pathways_resource_manager_expected_num_worker_jobs={args.num_slices}"""
+              - --pathways_expected_instances={expected_instances}"""
   if args.use_pathways:
-    return yaml.format(args=args)
+    return yaml.format(
+        args=args,
+        expected_instances=compute_pathways_expected_instances(args, system),
+    )
   else:
     return ''
 
 
+def compute_pathways_expected_instances(
+    args, system: SystemCharacteristics
+) -> str:
+  """Computes the expected instances from the system characteristics.
+  Args:
+    args: user provided args.
+    system: system characteristics.
+
+  Returns:
+    str: formatted string representing the expected instances (eg:
+    "tpuv4:2x2x2,tpuv4:2x2x2" for 2 slices of v4-16).
+  """
+  expected_instances = ','.join([
+      f'tpu{get_pathways_expected_tpu_type(system.device_type)}:{system.topology}'
+      for _ in range(args.num_slices)
+  ])
+
+  xpk_print(f'Pathways expected instances are: {expected_instances}')
+  return expected_instances
+
+
+def get_pathways_expected_tpu_type(device_type: str) -> str:
+  """Returns the device type expected by Pathways
+  Args:
+    device_type: the system characteristic device type
+
+  Returns:
+    str: the device type expected by pathways.
+  """
+  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'
+    )
+    xpk_exit(1)
+  return pathways_expected_instance
+
+
 def get_pathways_worker_args(args) -> str:
   """Arguments for the Pathways workers.
   Args:
@@ -5516,7 +5579,7 @@ def workload_create(args) -> int:
             system.accelerator_type, system
         ),
         machine_label=create_machine_label(system.accelerator_type, system),
-        pathways_rm_args=get_pathways_rm_args(args),
+        pathways_rm_args=get_pathways_rm_args(args, system),
         pathways_worker_args=get_pathways_worker_args(args),
         pathways_proxy_args=get_pathways_proxy_args(args),
         resource_type=resource_type,
@@ -5761,6 +5824,7 @@ def get_workload_list(args) -> tuple[int, str]:
       f' with filter-by-jobs={args.filter_by_job}',
       args,
   )
+
   return return_code, return_value
 
 
@@ -5869,7 +5933,7 @@ def workload_list(args) -> int:
   if return_code != 0:
     xpk_print(f'List Job request returned ERROR {return_code}')
     xpk_exit(return_code)
-  xpk_print(return_value)
+  xpk_print(f'Workload List Output:\n{return_value}')
   xpk_exit(0)
 
 
@@ -6380,10 +6444,10 @@ cluster_create_optional_arguments.add_argument(
 cluster_create_optional_arguments.add_argument(
     '--gke-version',
     type=str,
-    default=default_gke_version,
     help=(
-        'The GKE version of the cluster and respective clusters. The default is'
-        f' "{default_gke_version}".'
+        'The GKE version of the cluster and respective clusters. The'
+        ' default is'
+        ' determined dynamically based on RAPID channel recommended version.'
     ),
 )
 cluster_create_optional_arguments.add_argument(

xpk-0.4.0.dist-info/RECORD

@@ -1,7 +0,0 @@
-xpk.py,sha256=247TxfryVBdgpc6qpwCXOmwXD78nLGi1iHTIfTfflGU,215487
-xpk-0.4.0.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-xpk-0.4.0.dist-info/METADATA,sha256=qIptNh3dJ_OhYAkMYLz3Wltr0H_R2nwSDJQpphPZWiA,44333
-xpk-0.4.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-xpk-0.4.0.dist-info/entry_points.txt,sha256=lhrMqkTA09DLePaqxSMyW2RCLUKs2X1c84baGhMev_k,33
-xpk-0.4.0.dist-info/top_level.txt,sha256=aDe4N0jicmuWExx_6w0TxWQJaEuPSs9BnLU-3aF1GLo,4
-xpk-0.4.0.dist-info/RECORD,,