xpk is set to a specific version

GoogleCloudPlatform · Jan 9, 2025 · 5d0c2c6 · 5d0c2c6
1 parent 738dc43
commit 5d0c2c6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/xlml/utils/xpk.py b/xlml/utils/xpk.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utilities to run workloads with xpk (https://github.com/google/xpk)."""
+"""Utilities to run workloads with xpk (https://github.com/AI-Hypercomputer/xpk)."""
 
 import os
 import tempfile
@@ -83,7 +83,7 @@ def run_workload(
     )
     cmds = [
         "set -xue",
-        f"git clone https://github.com/google/xpk {tmpdir}/xpk",
+        f"git clone --branch v0.4.1 https://github.com/AI-Hypercomputer/xpk {tmpdir}/xpk",
         "pip install ruamel.yaml docker",
     ]
     if accelerator_type == GpuVersion.XPK_H100_MEGA.value: