Fix: --no-enable-autoupgrade when gke_version specified (#783)

* fix: --no-enable-autoupgrade when gke_version specified and no rapid channel set

* make goldens

* AcceleratorType.GPU

* Completely remove setting Rapid channel

* Reverse test case names

Author: jamOne-

goldens/Basic_cluster_create.txt

@@ -9,7 +9,7 @@ gcloud container get-server-config --project=golden-project --region=us-central1
 [XPK] Task: `Find if Cluster Exists` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=location~"us-central1.*" --format="csv[no-heading](name)"
 [XPK] Task: `GKE Cluster Create` is implemented by the following command not running since it is a dry run. 
-gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6   --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --location-policy=BALANCED --scopes=storage-full,gke-default
+gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6 --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --location-policy=BALANCED --scopes=storage-full,gke-default
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)"
 [XPK] Task: `Check if Private Nodes is enabled in cluster.` is implemented by the following command not running since it is a dry run. 

goldens/Cluster_create_private.txt

@@ -9,7 +9,7 @@ gcloud container get-server-config --project=golden-project --region=us-central1
 [XPK] Task: `Find if Cluster Exists` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=location~"us-central1.*" --format="csv[no-heading](name)"
 [XPK] Task: `GKE Cluster Create` is implemented by the following command not running since it is a dry run. 
-gcloud beta container clusters create golden-cluster-private --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=n1-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 4   --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --enable-master-authorized-networks --enable-private-nodes --location-policy=BALANCED --scopes=storage-full,gke-default --enable-ip-alias
+gcloud beta container clusters create golden-cluster-private --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=n1-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 4 --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --enable-master-authorized-networks --enable-private-nodes --location-policy=BALANCED --scopes=storage-full,gke-default --enable-ip-alias
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=name=golden-cluster-private --format="value(location)"
 [XPK] Task: `Check if Private Nodes is enabled in cluster.` is implemented by the following command not running since it is a dry run. 

goldens/Cluster_create_with_gb200-4.txt

@@ -9,7 +9,7 @@ gcloud container get-server-config --project=golden-project --region=us-central1
 [XPK] Task: `Find if Cluster Exists` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=location~"us-central1.*" --format="csv[no-heading](name)"
 [XPK] Task: `GKE Cluster Create` is implemented by the following command not running since it is a dry run. 
-gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6   --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --enable-dataplane-v2 --enable-multi-networking --no-enable-autoupgrade --enable-ip-alias
+gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6 --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --no-enable-autoupgrade --enable-dataplane-v2 --enable-multi-networking --enable-ip-alias
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)"
 [XPK] Task: `Check if Private Nodes is enabled in cluster.` is implemented by the following command not running since it is a dry run. 

goldens/NAP_cluster-create.txt

@@ -9,7 +9,7 @@ gcloud container get-server-config --project=golden-project --region=us-central1
 [XPK] Task: `Find if Cluster Exists` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=location~"us-central1.*" --format="csv[no-heading](name)"
 [XPK] Task: `GKE Cluster Create` is implemented by the following command not running since it is a dry run. 
-gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6   --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --location-policy=BALANCED --scopes=storage-full,gke-default
+gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6 --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --location-policy=BALANCED --scopes=storage-full,gke-default
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)"
 [XPK] Task: `Check if Private Nodes is enabled in cluster.` is implemented by the following command not running since it is a dry run. 

goldens/NAP_cluster-create_with_pathways.txt

@@ -9,7 +9,7 @@ gcloud container get-server-config --project=golden-project --region=us-central1
 [XPK] Task: `Find if Cluster Exists` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=location~"us-central1.*" --format="csv[no-heading](name)"
 [XPK] Task: `GKE Cluster Create` is implemented by the following command not running since it is a dry run. 
-gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6   --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --location-policy=BALANCED --scopes=storage-full,gke-default --enable-ip-alias
+gcloud beta container clusters create golden-cluster --project=golden-project --region=us-central1 --node-locations=us-central1-a --cluster-version=0 --machine-type=e2-standard-16 --enable-autoscaling --total-min-nodes 1 --total-max-nodes 1000 --num-nodes 6 --enable-dns-access --autoscaling-profile=optimize-utilization --labels=gke_product_type=xpk --location-policy=BALANCED --scopes=storage-full,gke-default --enable-ip-alias
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
 gcloud container clusters list --project=golden-project --filter=name=golden-cluster --format="value(location)"
 [XPK] Task: `Check if Private Nodes is enabled in cluster.` is implemented by the following command not running since it is a dry run. 

src/xpk/commands/cluster.py

@@ -1142,12 +1142,6 @@ def run_gke_cluster_create_command(
   # benefit from a larger initial `--num-nodes`. After the cluster is created,
   # the auto-scaler can reduce/increase the nodes based on the load.
 
-  # If the user passes in the gke version then we use that directly instead of the rapid release.
-  # This allows users to directly pass a specified gke version without release channel constraints.
-  rapid_release_cmd = ''
-  if args.gke_version is not None:
-    rapid_release_cmd = ' --release-channel rapid'
-
   command = (
       'gcloud beta container clusters create'
       f' {args.cluster} --project={args.project}'
@@ -1158,13 +1152,14 @@ def run_gke_cluster_create_command(
       ' --enable-autoscaling'
       ' --total-min-nodes 1 --total-max-nodes 1000'
       f' --num-nodes {args.default_pool_cpu_num_nodes}'
-      f' {args.custom_cluster_arguments}'
-      f' {rapid_release_cmd}'
       ' --enable-dns-access'
       ' --autoscaling-profile=optimize-utilization'
       ' --labels=gke_product_type=xpk'
   )
 
+  if args.gke_version or system.accelerator_type == AcceleratorType.GPU:
+    command += ' --no-enable-autoupgrade'
+
   enable_ip_alias = False
 
   if args.private or args.authorized_networks is not None:
@@ -1173,10 +1168,7 @@ def run_gke_cluster_create_command(
 
   if system.accelerator_type == AcceleratorType.GPU:
     enable_ip_alias = True
-    command += (
-        ' --enable-dataplane-v2'
-        ' --enable-multi-networking --no-enable-autoupgrade'
-    )
+    command += ' --enable-dataplane-v2 --enable-multi-networking'
   else:
     command += ' --location-policy=BALANCED --scopes=storage-full,gke-default'
 
@@ -1216,6 +1208,9 @@ def run_gke_cluster_create_command(
     addons_str = ','.join(addons)
     command += f' --addons={addons_str}'
 
+  if args.custom_cluster_arguments:
+    command += f' {args.custom_cluster_arguments}'
+
   return_code = run_command_with_updates(command, 'GKE Cluster Create')
   if return_code != 0:
     xpk_print(f'GKE Cluster Create request returned ERROR {return_code}')

src/xpk/commands/cluster_test.py

@@ -16,11 +16,13 @@
 
 from argparse import Namespace
 from dataclasses import dataclass
+from typing import Any
 from unittest.mock import MagicMock, patch
 import pytest
 
-from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args
+from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command
 from xpk.core.system_characteristics import SystemCharacteristics, UserFacingNameToSystemCharacteristics
+from xpk.core.testing.commands_tester import CommandsTester
 from xpk.utils.feature_flags import FeatureFlags
 
 
@@ -29,10 +31,11 @@ class _Mocks:
   common_print_mock: MagicMock
   commands_print_mock: MagicMock
   commands_get_reservation_deployment_type: MagicMock
+  commands_tester: CommandsTester
 
 
 @pytest.fixture
-def mock_common_print_and_exit(mocker):
+def mocks(mocker) -> _Mocks:
   common_print_mock = mocker.patch(
       'xpk.commands.common.xpk_print',
       return_value=None,
@@ -48,103 +51,137 @@ def mock_common_print_and_exit(mocker):
       common_print_mock=common_print_mock,
       commands_get_reservation_deployment_type=commands_get_reservation_deployment_type,
       commands_print_mock=commands_print_mock,
+      commands_tester=CommandsTester(
+          mocker,
+          run_command_with_updates_path=(
+              'xpk.commands.cluster.run_command_with_updates'
+          ),
+      ),
   )
 
 
-DEFAULT_TEST_SYSTEM: SystemCharacteristics = (
-    UserFacingNameToSystemCharacteristics['l4-1']
-)
+def construct_args(**kwargs: Any) -> Namespace:
+  args_dict = dict(
+      project='project',
+      zone='us-central1-a',
+      reservation='',
+      default_pool_cpu_machine_type='test-machine-type',
+      cluster='test-cluster',
+      default_pool_cpu_num_nodes='100',
+      sub_slicing=False,
+      gke_version='',
+      private=False,
+      authorized_networks=None,
+      enable_pathways=False,
+      enable_ray_cluster=False,
+      enable_workload_identity=False,
+      enable_gcsfuse_csi_driver=False,
+      enable_gcpfilestore_csi_driver=False,
+      enable_parallelstore_csi_driver=False,
+      enable_pd_csi_driver=False,
+      enable_lustre_csi_driver=False,
+      custom_cluster_arguments='',
+      num_slices=1,
+      num_nodes=1,
+      flex=False,
+      memory_limit='100Gi',
+      cpu_limit=100,
+      cluster_cpu_machine_type='',
+  )
+  args_dict.update(kwargs)
+  return Namespace(**args_dict)
+
+
+GPU_TEST_SYSTEM: SystemCharacteristics = UserFacingNameToSystemCharacteristics[
+    'l4-1'
+]
 SUB_SLICING_SYSTEM: SystemCharacteristics = (
     UserFacingNameToSystemCharacteristics['v6e-4x4']
 )
+TPU_TEST_SYSTEM: SystemCharacteristics = UserFacingNameToSystemCharacteristics[
+    'v6e-4x4'
+]
 
 
 def test_validate_cluster_create_args_for_correct_args_pass(
-    mock_common_print_and_exit: _Mocks,
+    mocks: _Mocks,
 ):
   args = Namespace()
 
-  _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
+  _validate_cluster_create_args(args, GPU_TEST_SYSTEM)
 
-  assert mock_common_print_and_exit.common_print_mock.call_count == 0
+  assert mocks.common_print_mock.call_count == 0
 
 
 def test_validate_cluster_create_args_for_correct_sub_slicing_args_pass(
-    mock_common_print_and_exit: _Mocks,
+    mocks: _Mocks,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = True
-  args = Namespace(
+  args = construct_args(
       sub_slicing=True,
       reservation='test-reservation',
-      project='project',
-      zone='zone',
   )
 
   _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
 
-  assert mock_common_print_and_exit.common_print_mock.call_count == 0
+  assert mocks.common_print_mock.call_count == 0
 
 
 def test_validate_cluster_create_args_for_not_supported_system_throws(
-    mock_common_print_and_exit: _Mocks,
+    mocks: _Mocks,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = True
-  args = Namespace(
+  args = construct_args(
       sub_slicing=True,
       reservation='test-reservation',
-      project='project',
-      zone='zone',
   )
 
   with pytest.raises(SystemExit):
-    _validate_cluster_create_args(args, DEFAULT_TEST_SYSTEM)
+    _validate_cluster_create_args(args, GPU_TEST_SYSTEM)
 
-  assert mock_common_print_and_exit.common_print_mock.call_count == 1
+  assert mocks.common_print_mock.call_count == 1
   assert (
-      mock_common_print_and_exit.common_print_mock.call_args[0][0]
+      mocks.common_print_mock.call_args[0][0]
       == 'Error: l4-1 does not support Sub-slicing.'
   )
 
 
 def test_validate_cluster_create_args_for_missing_reservation(
-    mock_common_print_and_exit: _Mocks,
+    mocks: _Mocks,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = True
-  args = Namespace(
-      sub_slicing=True, project='project', zone='zone', reservation=None
+  args = construct_args(
+      sub_slicing=True,
+      reservation=None,
   )
 
   with pytest.raises(SystemExit):
     _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
 
-  assert mock_common_print_and_exit.commands_print_mock.call_count == 1
+  assert mocks.commands_print_mock.call_count == 1
   assert (
       'Validation failed: Sub-slicing cluster creation requires'
-      in mock_common_print_and_exit.commands_print_mock.call_args[0][0]
+      in mocks.commands_print_mock.call_args[0][0]
   )
 
 
 def test_validate_cluster_create_args_for_invalid_reservation(
-    mock_common_print_and_exit: _Mocks,
+    mocks: _Mocks,
 ):
   FeatureFlags.SUB_SLICING_ENABLED = True
-  args = Namespace(
+  args = construct_args(
       sub_slicing=True,
-      project='project',
-      zone='zone',
       reservation='test-reservation',
   )
-  mock_common_print_and_exit.commands_get_reservation_deployment_type.return_value = (
-      'SPARSE'
-  )
+  mocks.commands_get_reservation_deployment_type.return_value = 'SPARSE'
 
   with pytest.raises(SystemExit):
     _validate_cluster_create_args(args, SUB_SLICING_SYSTEM)
 
-  assert mock_common_print_and_exit.commands_print_mock.call_count == 5
+  assert mocks.commands_print_mock.call_count == 5
   assert (
       'Refer to the documentation for more information on creating Cluster'
-      in mock_common_print_and_exit.commands_print_mock.call_args[0][0]
+      in mocks.commands_print_mock.call_args[0][0]
   )
 
 
@@ -155,17 +192,73 @@ def test_install_kueue_returns_kueue_installation_code(
   mock_kueue_manager_install.return_value = 17
 
   code = _install_kueue(
-      args=Namespace(
-          num_slices=1,
-          num_nodes=1,
-          flex=False,
-          memory_limit='100Gi',
-          cpu_limit=100,
-          enable_pathways=False,
-          sub_slicing=False,
-      ),
-      system=DEFAULT_TEST_SYSTEM,
+      args=construct_args(),
+      system=GPU_TEST_SYSTEM,
       autoprovisioning_config=None,
   )
 
   assert code == 17
+
+
+def test_run_gke_cluster_create_command_specifies_custom_cluster_arguments_last(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(
+          custom_cluster_arguments='--enable-autoscaling=False --foo=baz'
+      ),
+      gke_control_plane_version='1.2.3',
+      system=TPU_TEST_SYSTEM,
+  )
+
+  assert result == 0
+  mocks.commands_tester.assert_command_run(
+      'clusters create',
+      ' --enable-autoscaling',
+      ' --enable-autoscaling=False --foo=baz',
+  )
+
+
+def test_run_gke_cluster_create_command_without_gke_version_does_not_have_no_autoupgrade_flag(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(gke_version=''),
+      gke_control_plane_version='1.2.3',
+      system=TPU_TEST_SYSTEM,
+  )
+
+  assert result == 0
+  mocks.commands_tester.assert_command_not_run(
+      'clusters create', ' --no-enable-autoupgrade'
+  )
+
+
+def test_run_gke_cluster_create_command_with_gke_version_has_no_autoupgrade_flag(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(gke_version='1.2.3'),
+      gke_control_plane_version='1.2.3',
+      system=TPU_TEST_SYSTEM,
+  )
+
+  assert result == 0
+  mocks.commands_tester.assert_command_run(
+      'clusters create', ' --no-enable-autoupgrade'
+  )
+
+
+def test_run_gke_cluster_create_command_with_gpu_system_has_no_enable_autoupgrade(
+    mocks: _Mocks,
+):
+  result = run_gke_cluster_create_command(
+      args=construct_args(gke_version=''),
+      gke_control_plane_version='1.2.3',
+      system=GPU_TEST_SYSTEM,
+  )
+
+  assert result == 0
+  mocks.commands_tester.assert_command_run(
+      'clusters create', ' --no-enable-autoupgrade'
+  )