Merge branch 'develop' into chzheng/docker_image_flag

Author: ycchenzheng

src/xpk/commands/cluster.py

@@ -141,8 +141,6 @@ def cluster_adapt(args) -> None:
     if not tensorboard_config:
       xpk_exit(1)
 
-  # Provision node pools dynamically based on incoming workloads:
-  # Currently autoprovisioning is not supported with Pathways.
   autoprovisioning_config = None
   if args.enable_autoprovisioning:
     xpk_print('Enabling Autoprovisioning')
@@ -294,7 +292,7 @@ def cluster_create(args) -> None:
   # Provision node pools dynamically based on incoming workloads:
   # Currently autoprovisioning is not supported with Pathways.
   autoprovisioning_config = None
-  if not args.enable_pathways and args.enable_autoprovisioning:
+  if args.enable_autoprovisioning:
     xpk_print('Enabling Autoprovisioning')
     autoprovisioning_config, return_code = enable_autoprovisioning_on_cluster(
         args, system
@@ -819,7 +817,7 @@ def check_deployment_exists(args, deployment_name: str, namespace: str) -> bool:
 
 
 def verify_coredns_readiness(
-    args, timeout: int = 120, namespace: str = 'kube-system'
+    args, timeout: int = 240, namespace: str = 'kube-system'
 ):
   """Verifies CoreDNS readiness using kubectl wait commands."""
   xpk_print('Now verifying CoreDNS readiness...')

src/xpk/commands/workload.py

@@ -84,6 +84,7 @@
 from ..core.vertex import create_vertex_experiment
 from ..core.workload import (
     check_if_workload_exists,
+    get_jobsets_list_gcp_link,
     get_workload_list,
     wait_for_job_completion,
     zone_to_region,
@@ -761,4 +762,8 @@ def workload_list(args) -> None:
     xpk_print(f'List Job request returned ERROR {return_code}')
     xpk_exit(return_code)
   xpk_print(f'Workload List Output:\n{return_value}')
+
+  workload_list_gcp_link = get_jobsets_list_gcp_link(project=args.project)
+  xpk_print(f'See your workloads in Cloud Console: {workload_list_gcp_link}')
+
   xpk_exit(0)

src/xpk/core/cluster.py

@@ -62,8 +62,8 @@ def set_jobset_on_cluster(args) -> int:
     0 if successful and 1 otherwise.
   """
   command = (
-      'kubectl apply --server-side -f'
-      f' https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml'
+      'kubectl apply --server-side --force-conflicts'
+      f' -f https://github.com/kubernetes-sigs/jobset/releases/download/{JOBSET_VERSION}/manifests.yaml'
   )
   task = f'Install Jobset on {args.cluster}'
   return_code = run_command_with_updates_retry(command, task, args)

src/xpk/core/kueue.py

@@ -43,7 +43,7 @@
 KUEUE_VERSION = 'v0.12.2'
 CLUSTER_QUEUE_NAME = 'cluster-queue'
 LOCAL_QUEUE_NAME = 'multislice-queue'
-WAIT_FOR_KUEUE_TIMEOUT = '5m'
+WAIT_FOR_KUEUE_TIMEOUT = '10m'
 MEMORY_SIZE_PER_VM = 1.2
 MIN_MEMORY_LIMIT_SIZE = 4096
 

src/xpk/core/nap.py

@@ -269,9 +269,6 @@ def is_autoprovisioning_enabled(
     bool is true if autoprovisioning is enabled, false otherwise.
     int of 0 if successful and 1 otherwise.
   """
-  # Currently autoprovisioning is not enabled for Pathways workloads. b/360898087
-  if args.use_pathways:
-    return False, 0
 
   resources_configmap_name = f'{args.cluster}-{CLUSTER_RESOURCES_CONFIGMAP}'
   cluster_config_map = get_cluster_configmap(args, resources_configmap_name)

src/xpk/core/resources.py

@@ -108,13 +108,7 @@ def create_cluster_configmaps(
   device_type = system.device_type
   if system.accelerator_type == AcceleratorType['GPU']:
     resources_data = f'{device_type}: "{int(args.num_nodes)}"'
-  elif (
-      not args.enable_pathways
-      and args.enable_autoprovisioning
-      and autoprovisioning_config
-  ):
-    # Currently autoprovisioning is not supported with Pathways.
-    # Auto provisioning will have variable topologies for a gke accelerator type.
+  elif args.enable_autoprovisioning and autoprovisioning_config:
     resources_data = (
         f'{system.gke_accelerator}: {AUTOPROVISIONING_CONFIG_VALUE}'
     )

src/xpk/core/tests/unit/test_workload.py

@@ -0,0 +1,28 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from xpk.core.workload import get_jobsets_list_gcp_link
+
+
+def test_get_jobsets_list_gcp_link():
+  result = get_jobsets_list_gcp_link(
+      project='test-project',
+  )
+
+  assert (
+      result
+      == 'https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project=test-project'
+  )

src/xpk/core/workload.py

@@ -14,6 +14,7 @@
 limitations under the License.
 """
 
+import re
 from ..utils.console import xpk_exit, xpk_print
 from .commands import run_command_for_value
 from .gcloud_context import zone_to_region
@@ -240,3 +241,13 @@ def wait_for_job_completion(args) -> int:
     xpk_print('Your workload did not complete successfully')
     return 125
   return 0
+
+
+GCP_NAME_FILTER_VALUE_REGEX = re.compile(r'[a-z0-9\-]+')
+"""Defines correct name prefix value (contains only letters, numbers and dashes) that can be used in GCP filter chips."""
+
+
+def get_jobsets_list_gcp_link(project: str) -> str:
+  """Returns a link to Cloud Console JobSets list"""
+
+  return f'https://console.cloud.google.com/kubernetes/aiml/deployments/jobs?project={project}'

src/xpk/core/workload_decorators/tcpx_decorator.py

@@ -131,6 +131,7 @@ def add_volumes(job_manifest: dict):
   })
   volumes.append({'name': 'sys', 'hostPath': {'path': '/sys'}})
   volumes.append({'name': 'proc-sys', 'hostPath': {'path': '/proc/sys'}})
+  volumes.append({'name': 'tcpx-socket', 'hostPath': {'path': '/run/tcpx'}})
   volumes.append(
       {'name': 'dshm', 'emptyDir': {'medium': 'Memory', 'sizeLimit': '128Gi'}}
   )

src/xpk/core/workload_decorators/tests/__init__.py

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