Add customizable binaries support

src/xpk/core/config.py

@@ -54,6 +54,7 @@ def _get_version() -> str:
 CLIENT_ID_KEY = 'client-id'
 SEND_TELEMETRY_KEY = 'send-telemetry'
 ZONE_KEY = 'zone'
+CUSTOM_BINARIES_PATH_KEY = 'custom-binaries-path'
 
 DEFAULT_KEYS = [
     CFG_BUCKET_KEY,
@@ -62,6 +63,7 @@ def _get_version() -> str:
     CLIENT_ID_KEY,
     SEND_TELEMETRY_KEY,
     ZONE_KEY,
+    CUSTOM_BINARIES_PATH_KEY,
 ]
 VERTEX_TENSORBOARD_FEATURE_FLAG = XPK_CURRENT_VERSION >= '0.4.0'
 

src/xpk/main.py

@@ -39,10 +39,11 @@
 from .parser.core import set_parser
 from .parser.common import extract_command_path, enable_flags_usage_tracking, retrieve_flags
 from .core.updates import print_xpk_hello
-from .core.config import set_config, FileSystemConfig
+from .core.config import set_config, get_config, FileSystemConfig, CUSTOM_BINARIES_PATH_KEY
 from .core.telemetry import MetricsCollector, send_clearcut_payload, should_send_telemetry
 from .utils.console import xpk_print, exit_code_to_int
 from .utils.execution_context import set_context
+from .utils.environment import custom_binaries_path_env
 from .utils.kubectl import sandbox_kubeconfig
 ################### Compatibility Check ###################
 # Check that the user runs the below version or greater.
@@ -92,7 +93,10 @@ def main() -> None:
         and main_args.sandbox_kubeconfig
     )
     opt_sandbox = sandbox_kubeconfig if is_sandbox else contextlib.nullcontext
-    with opt_sandbox():
+    with (
+        opt_sandbox(),
+        custom_binaries_path_env(get_config().get(CUSTOM_BINARIES_PATH_KEY)),
+    ):
       main_args.func(main_args)
     xpk_print('XPK Done.', flush=True)
     MetricsCollector.log_complete(0)

src/xpk/utils/environment.py

@@ -0,0 +1,43 @@
+"""
+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 contextlib
+import os
+from typing import Iterator
+
+
+@contextlib.contextmanager
+def custom_binaries_path_env(
+    custom_binaries_path: str | None = None,
+) -> Iterator[None]:
+  """Context manager to use a custom binaries path."""
+  if custom_binaries_path:
+    environ = os.environ
+    backup = environ.get('PATH')
+    environ['PATH'] = (
+        f'{custom_binaries_path}{os.pathsep}{backup}'
+        if backup is not None
+        else custom_binaries_path
+    )
+    try:
+      yield
+    finally:
+      if backup is None:
+        del environ['PATH']
+      else:
+        environ['PATH'] = backup
+  else:
+    yield