Fix dockerfile issues

pytorch · Jun 12, 2024 · aeaaa0d · aeaaa0d
1 parent bacbff4
commit aeaaa0d
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/docker/build-torchbench-nightly-docker.sh b/docker/build-torchbench-nightly-docker.sh
@@ -1 +1,4 @@
-docker build . -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:latest
+TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH:-main}
+
+docker build . -f torchbench-nightly.dockerfile -t ghcr.io/pytorch/torchbench:latest \
+    --build-arg TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH}
diff --git a/docker/torchbench-nightly.dockerfile b/docker/torchbench-nightly.dockerfile
@@ -8,10 +8,13 @@ ENV SETUP_SCRIPT=/workspace/setup_instance.sh
 ARG TORCHBENCH_BRANCH=${TORCHBENCH_BRANCH:-main}
 ARG FORCE_DATE=${FORCE_DATE}
 
-# Setup Conda env and CUDA
+# Checkout Torchbench and submodules
 RUN git clone -b "${TORCHBENCH_BRANCH}" --single-branch \
- https://github.com/pytorch/benchmark /workspace/benchmark
+    https://github.com/pytorch/benchmark /workspace/benchmark
+RUN cd /workspace/benchmark \
+    git submodule update --init --recursive
 
+# Setup conda env and CUDA
 RUN cd /workspace/benchmark && \
     . ${SETUP_SCRIPT} && \
     python ./utils/python_utils.py --create-conda-env ${CONDA_ENV} && \

diff --git a/userbenchmark/triton/install.py b/userbenchmark/triton/install.py
@@ -9,7 +9,7 @@
 
 def install_fbgemm():
     cmd = [sys.executable, "setup.py", "bdist_wheel", "--package_variant=genai"]
-    subprocess.check_call(cmd, cwd=FBGEMM_PATH)
+    subprocess.check_call(cmd, cwd=str(FBGEMM_PATH.resolve()))
 
 def test_fbgemm():
     cmd = [sys.executable, "-c", '"import fbgemm_gpu.experimental.gen_ai"']