From a91ef813f943e7d80fe53998d2ea311ce0a17bae Mon Sep 17 00:00:00 2001
From: Jason Wang <blahblahj.wsy@gmail.com>
Date: Tue, 10 Sep 2024 18:30:50 -0700
Subject: [PATCH] fix extra context docker build bug (#724)

---
 docker/tpu/Dockerfile.incremental | 2 +-
 src/levanter/infra/docker.py      | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/docker/tpu/Dockerfile.incremental b/docker/tpu/Dockerfile.incremental
index f0369736c..64c14b4c9 100644
--- a/docker/tpu/Dockerfile.incremental
+++ b/docker/tpu/Dockerfile.incremental
@@ -22,4 +22,4 @@ ADD . /opt/levanter
 
 # Add $EXTRA_CTX to the same location as in local machine.
 # it's already in the image, so we don't need to copy it. just move it if we set EXTRA_CTX
-RUN if [ -f ".mnt" ]; then mkdir -p $(dirname $EXTRA_CTX) && mv .mnt $EXTRA_CTX; fi
+RUN if [ -f ".mnt" ] || [ -d ".mnt" ]; then mkdir -p $(dirname $EXTRA_CTX) && mv .mnt $EXTRA_CTX; fi
diff --git a/src/levanter/infra/docker.py b/src/levanter/infra/docker.py
index 2f8052f87..63a51ae2f 100644
--- a/src/levanter/infra/docker.py
+++ b/src/levanter/infra/docker.py
@@ -159,7 +159,7 @@ def copy_extra_ctx(extra_ctx):
         mount_dst = Path(".mnt")
         _cp(extra_ctx, mount_dst)
         try:
-            yield mount_dst
+            yield extra_ctx
         finally:
             _rm(mount_dst)
     else: