Show prompt to update xpk (#771)

scaliby · web-flow · commit bdabf90caa44 · 2025-10-31T07:28:50.000Z
* feat: show prompt to update xpk

* style: lint

* fix: xpk print

* update goldens

* style: lint
diff --git a/goldens/Basic_cluster_create.txt b/goldens/Basic_cluster_create.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --tpu-type=tpu7x-8 --spot --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster create for cluster golden-cluster:
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Determine server supported GKE versions for default rapid gke version` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Batch.txt b/goldens/Batch.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py batch --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run batch-read.sh
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Working on golden-project and us-central1-a
 [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)"
diff --git a/goldens/Cluster_create_private.txt b/goldens/Cluster_create_private.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster create-pathways --project=golden-project --zone=us-central1-a --cluster=golden-cluster-private --private --tpu-type=v5p-8  --num-slices=1 --default-pool-cpu-machine-type=n1-standard-16 --default-pool-cpu-num-nodes=4 --reservation=golden-reservation --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster create for cluster golden-cluster-private:
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Determine server supported GKE versions for default rapid gke version` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Cluster_create_with_gb200-4.txt b/goldens/Cluster_create_with_gb200-4.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --device-type=gb200-4 --reservation=golden-reservation --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster create for cluster golden-cluster:
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Determine server supported GKE versions for default rapid gke version` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Cluster_delete.txt b/goldens/Cluster_delete.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster delete --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster delete for cluster: golden-cluster
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Cluster_delete_force.txt b/goldens/Cluster_delete_force.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster delete --project=golden-project --zone=us-central1-a --cluster=golden-cluster --force --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster delete for cluster: golden-cluster
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Job_cancel.txt b/goldens/Job_cancel.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py job cancel golden-job --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting job cancel for job: ['golden-job']
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Job_info.txt b/goldens/Job_info.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py job info golden-job --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Task: `Getting job data` is implemented by the following command not running since it is a dry run. 
 kubectl-kjob describe slurm golden-job
 [XPK] Task: `Getting job info` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Job_list.txt b/goldens/Job_list.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py job ls --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Working on golden-project and us-central1-a
 [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)"
diff --git a/goldens/NAP_cluster-create.txt b/goldens/NAP_cluster-create.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster create --project=golden-project --zone=us-central1-a --enable-autoprovisioning --cluster=golden-cluster --tpu-type=tpu7x-8 --on-demand --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster create for cluster golden-cluster:
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Determine server supported GKE versions for default rapid gke version` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/NAP_cluster-create_with_pathways.txt b/goldens/NAP_cluster-create_with_pathways.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py cluster create-pathways --project=golden-project --zone=us-central1-a --enable-autoprovisioning --cluster=golden-cluster --tpu-type=tpu7x-8 --on-demand --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting cluster create for cluster golden-cluster:
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Determine server supported GKE versions for default rapid gke version` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Storage_list.txt b/goldens/Storage_list.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py storage list --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 NAME    TYPE    AUTO MOUNT    MOUNT POINT    READONLY    MANIFEST
 ------  ------  ------------  -------------  ----------  ----------
 [XPK] XPK Done.
diff --git a/goldens/Workload_create.txt b/goldens/Workload_create.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py workload create --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v5p-8 --num-slices=1 --script-dir=/tmp --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
 kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Workload_create_pathways.txt b/goldens/Workload_create_pathways.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py workload create-pathways --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --command "bash hello" --tpu-type=v5p-8 --num-slices=1 --script-dir=/tmp --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Task: `Check if Workload Already Exists` is implemented by the following command not running since it is a dry run. 
 kubectl get workloads -o=custom-columns='Jobset:.metadata.ownerReferences[0].name'
 [XPK] Task: `GKE Cluster Get ConfigMap` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Workload_delete.txt b/goldens/Workload_delete.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py workload delete --project=golden-project --zone=us-central1-a --cluster=golden-cluster --workload=golden-workload --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting Workload delete
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
diff --git a/goldens/Workload_list.txt b/goldens/Workload_list.txt
@@ -1,5 +1,5 @@
 $ python3 xpk.py workload list --project=golden-project --zone=us-central1-a --cluster=golden-cluster --dry-run
-[XPK] Starting xpk
+[XPK] Starting xpk v0.14.2
 [XPK] Starting workload list
 [XPK] Working on golden-project and us-central1-a
 [XPK] Task: `Find cluster region or zone` is implemented by the following command not running since it is a dry run. 
diff --git a/src/xpk/core/updates.py b/src/xpk/core/updates.py
@@ -0,0 +1,57 @@
+"""
+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 json
+from json.decoder import JSONDecodeError
+from .commands import run_command_for_value
+from ..utils.console import xpk_print
+from ..utils.execution_context import is_dry_run
+from packaging.version import Version
+from .config import __version__
+
+
+def get_latest_xpk_version() -> tuple[int, Version | None]:
+  if is_dry_run():
+    return 0, Version(__version__)
+
+  return_code, result = run_command_for_value(
+      command="pip index versions xpk --json",
+      task="Retrieve latest XPK version",
+      quiet=True,
+  )
+
+  if return_code != 0:
+    return return_code, None
+
+  try:
+    parsed = json.loads(result.strip())
+    return 0, Version(parsed["latest"])
+  except JSONDecodeError:
+    return 1, None
+
+
+def print_xpk_hello() -> None:
+  current_version = Version(__version__)
+  xpk_print(f"Starting xpk v{current_version}", flush=True)
+  return_code, latest_version = get_latest_xpk_version()
+  if return_code != 0 or latest_version is None:
+    return
+  if current_version < latest_version:
+    xpk_print(
+        f"XPK version v{current_version} is outdated. Please consider upgrading"
+        f" to v{latest_version}",
+        flush=True,
+    )
diff --git a/src/xpk/core/updates_test.py b/src/xpk/core/updates_test.py
@@ -0,0 +1,80 @@
+"""
+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 ..utils.execution_context import set_dry_run
+from .updates import get_latest_xpk_version, print_xpk_hello
+from packaging.version import Version
+from .config import __version__
+from unittest.mock import MagicMock, patch
+
+
+def test_get_latest_xpk_version_returns_current_version_for_dry_run():
+  set_dry_run(True)
+  return_code, version = get_latest_xpk_version()
+  assert return_code == 0
+  assert version == Version(__version__)
+
+
+@patch('xpk.core.updates.run_command_for_value')
+def test_get_latest_xpk_version_returns_error_when_underlaying_command_errors(
+    run_command_for_value: MagicMock,
+):
+  run_command_for_value.return_value = (1, None)
+  set_dry_run(False)
+  return_code, version = get_latest_xpk_version()
+  assert return_code == 1
+  assert version is None
+
+
+@patch('xpk.core.updates.run_command_for_value')
+def test_get_latest_xpk_version_returns_version_returned_from_command(
+    run_command_for_value: MagicMock,
+):
+  run_command_for_value.return_value = (0, '{"latest": "1.0.0"}')
+  set_dry_run(False)
+  return_code, version = get_latest_xpk_version()
+  assert return_code == 0
+  assert version == Version('1.0.0')
+
+
+@patch('xpk.core.updates.xpk_print')
+@patch('xpk.core.updates.get_latest_xpk_version')
+def test_print_xpk_hello_does_not_print_update_when_version_check_fails(
+    get_latest_xpk_version: MagicMock, xpk_print: MagicMock
+):
+  get_latest_xpk_version.return_value = (1, None)
+  print_xpk_hello()
+  xpk_print.assert_called_once()
+
+
+@patch('xpk.core.updates.xpk_print')
+@patch('xpk.core.updates.get_latest_xpk_version')
+def test_print_xpk_hello_does_not_print_update_when_xpk_is_up_to_date(
+    get_latest_xpk_version: MagicMock, xpk_print: MagicMock
+):
+  get_latest_xpk_version.return_value = (0, Version(__version__))
+  print_xpk_hello()
+  xpk_print.assert_called_once()
+
+
+@patch('xpk.core.updates.xpk_print')
+@patch('xpk.core.updates.get_latest_xpk_version')
+def test_print_xpk_hello_prints_update_when_xpk_is_outdated(
+    get_latest_xpk_version: MagicMock, xpk_print: MagicMock
+):
+  get_latest_xpk_version.return_value = (0, Version('99.99.99'))
+  print_xpk_hello()
+  assert xpk_print.call_count == 2
diff --git a/src/xpk/main.py b/src/xpk/main.py
@@ -35,6 +35,7 @@
 import sys
 
 from .parser.core import set_parser
+from .core.updates import print_xpk_hello
 from .utils.console import xpk_print
 from .utils.execution_context import set_context
 ################### Compatibility Check ###################
@@ -62,13 +63,13 @@ def main() -> None:
   parser = argparse.ArgumentParser(description='xpk command', prog='xpk')
   set_parser(parser=parser)
 
-  xpk_print('Starting xpk', flush=True)
   main_args = parser.parse_args()
   main_args.enable_ray_cluster = False
   set_context(
       dry_run_value='dry_run' in main_args and main_args.dry_run,
       quiet_value='quiet' in main_args and main_args.quiet,
   )
+  print_xpk_hello()
   main_args.func(main_args)
   xpk_print('XPK Done.', flush=True)
 
diff --git a/src/xpk/utils/execution_context.py b/src/xpk/utils/execution_context.py
@@ -20,9 +20,19 @@
 
 def set_context(dry_run_value: bool, quiet_value: bool) -> None:
   """Sets the dry_run and quiet flags."""
+  set_dry_run(dry_run_value)
+  set_quiet(quiet_value)
+
+
+def set_dry_run(dry_run_value: bool) -> None:
+  """Sets the dry_run flag."""
   global dry_run
-  global quiet
   dry_run = dry_run_value
+
+
+def set_quiet(quiet_value: bool) -> None:
+  """Sets the quiet flag."""
+  global quiet
   quiet = quiet_value
 
 
