Support torch/torch_xla wheel locally (#142)

* Support torch/torch_xla wheel locally

* nit

* nit

Author: zpcore

.gitignore

@@ -35,6 +35,7 @@ pip-wheel-metadata/
 # Ignore build and distribution directories
 build/
 dist/
+local_dist/
 *.egg-info/
 
 # Ignore environment configuration files

README.md

@@ -149,6 +149,15 @@ add flag `--use-hf` to `tp run` command:
 tp run --use-hf torchprime/hf_models/train.py
 ```
 
+## Run with local torch/torch_xla wheel
+Torchprime supports run with user specified torch and torch_xla wheels placed
+under `local_dist/` directory. The wheel will be automatically installed in the
+docker image when use `tp run` command. To use the wheel, add flag
+`--use-local-wheel` to `tp run` command:
+```
+tp run --use-local-wheel torchprime/hf_models/train.py
+```
+
 ## Contributing
 
 Contributions are welcome! Please feel free to submit a pull request.

torchprime/launcher/Dockerfile

@@ -3,6 +3,8 @@
 FROM us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:nightly_3.10_tpuvm_cxx11_20250227
 
 ARG USE_TRANSFORMERS=false
+ARG USE_LOCAL_WHEEL=false
+
 # Install system dependencies
 RUN apt-get update && apt-get install -y curl gnupg
 
@@ -39,6 +41,23 @@ RUN pip install -e .
 
 # Now copy the rest of the repo
 COPY . /workspaces/torchprime
+
+# Install torch and torch_xla from local wheels if USE_LOCAL_WHEEL and exists
+# under local_dist directory. Note that you need to build the torch and
+# torch_xla using the 
+RUN if [ "$USE_LOCAL_WHEEL" = "true" ]; then \
+        if [ -d "local_dist" ] && [ "$(find local_dist -name 'torch-*.whl' | wc -l)" -gt 0 ]; then \
+            pip install local_dist/torch-*.whl; \
+        else \
+            echo "torch wheel not found in local_dist directory"; \
+        fi; \
+        if [ -d "local_dist" ] && [ "$(find local_dist -name 'torch_xla-*.whl' | wc -l)" -gt 0 ]; then \
+            pip install local_dist/torch_xla-*.whl; \
+        else \
+            echo "torch_xla wheel not found in local_dist directory"; \
+        fi; \
+    fi
+
 # This should not install any packages. Only symlink the source code.
 RUN pip install --no-deps -e .
 

torchprime/launcher/buildpush.py

@@ -48,7 +48,8 @@ def buildpush(
 
   build_cmd = f"{sudo_cmd} docker build"
   if build_arg:
-    build_cmd += f" --build-arg {build_arg}"
+    for _arg in build_arg:
+      build_cmd += f" --build-arg {_arg}"
   build_cmd += (
     f" --network=host --progress=auto -t {docker_tag} {context_dir} -f {docker_file}"
   )

torchprime/launcher/cli.py

@@ -208,8 +208,13 @@ def create_and_activate_gcloud(gcloud_config_name, config: Config):
   default=None,
 )
 @click.option("--use-hf", is_flag=True, help="Use HuggingFace transformer")
+@click.option(
+  "--use-local-wheel",
+  is_flag=True,
+  help="Use local torch and torch_xla wheels under folder local_dist/",
+)
 @interactive
-def run(args, name: str | None, use_hf: bool):
+def run(args, name: str | None, use_hf: bool, use_local_wheel: bool):
   """
   Runs the provided SPMD training command as an xpk job on a GKE cluster.
   """
@@ -218,7 +223,11 @@ def run(args, name: str | None, use_hf: bool):
   click.echo(get_project_dir().absolute())
 
   # Build docker image.
-  build_arg = "USE_TRANSFORMERS=true" if use_hf else None
+  build_arg = []
+  if use_hf:
+    build_arg.append("USE_TRANSFORMERS=true")
+  if use_local_wheel:
+    build_arg.append("USE_LOCAL_WHEEL=true")
   docker_project = config.docker_project
   if docker_project is None:
     docker_project = config.project