From 0dc0013e72495f6977dad10731c0639218061ea6 Mon Sep 17 00:00:00 2001
From: epwalsh <epwalsh10@gmail.com>
Date: Mon, 20 Mar 2023 15:49:08 -0700
Subject: [PATCH] update docker images

---
 .github/workflows/main.yml | 2 +-
 Makefile                   | 4 +++-
 README.md                  | 4 ++--
 docker/Dockerfile.gantry   | 2 +-
 docker/Dockerfile.test     | 2 +-
 5 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 95990c57f..42067de73 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -108,7 +108,7 @@ jobs:
     timeout-minutes: 15
     env:
       BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
-      BEAKER_IMAGE: dolma-test
+      BEAKER_IMAGE: dolma-torch2-test
       BEAKER_WORKSPACE: ai2/llm-testing
     steps:
       - name: Determine current commit SHA (pull request)
diff --git a/Makefile b/Makefile
index c299a43b7..9e7be2dcc 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,5 @@
 # If you update this, also update BEAKER_IMAGE in .github/workflows/main.yml
-IMAGE_NAME_BASE = dolma
+IMAGE_NAME_BASE = dolma-torch2
 # If you update this, also update BEAKER_WORKSPACE in .github/workflows/main.yml
 BEAKER_WORKSPACE = "ai2/llm-testing"
 
@@ -32,6 +32,7 @@ base-image :
 gantry-image : base-image
 	docker build -f docker/Dockerfile.gantry -t $(IMAGE_NAME_BASE)-gantry .
 	beaker image create $(IMAGE_NAME_BASE)-gantry --name $(IMAGE_NAME_BASE)-gantry-tmp --workspace $(BEAKER_WORKSPACE)
+	echo "Image uploaded successfully"
 	beaker image delete $(GANTRY_IMAGE) || true
 	beaker image rename $(BEAKER_USER)/$(IMAGE_NAME_BASE)-gantry-tmp $(IMAGE_NAME_BASE)-gantry
 
@@ -39,6 +40,7 @@ gantry-image : base-image
 test-image : base-image
 	docker build -f docker/Dockerfile.test -t $(IMAGE_NAME_BASE)-test .
 	beaker image create $(IMAGE_NAME_BASE)-test --name $(IMAGE_NAME_BASE)-test-tmp --workspace $(BEAKER_WORKSPACE)
+	echo "Image uploaded successfully"
 	beaker image delete $(TEST_IMAGE) || true
 	beaker image rename $(BEAKER_USER)/$(IMAGE_NAME_BASE)-test-tmp $(IMAGE_NAME_BASE)-test
 
diff --git a/README.md b/README.md
index 461bec08b..33e44b371 100644
--- a/README.md
+++ b/README.md
@@ -24,7 +24,7 @@ gantry run \
   --nfs \
   --priority preemptible \
   --gpus 8 \
-  --beaker-image dolma-gantry \
+  --beaker-image dolma-torch2-gantry \
   --cluster 'ai2/*-cirrascale' \
   --allow-dirty \
   -- composer scripts/train.py configs/1.2b-c4.yaml
@@ -36,7 +36,7 @@ Train the 70B model on c4 with gantry across multiple nodes:
 gantry run \
   --workspace ai2/llm-testing \
   --priority "high" \
-  --beaker-image dolma-gantry \
+  --beaker-image dolma-torch2-gantry \
   --cluster ai2/general-cirrascale-a100-80g-ib \
   --gpus 8 \
   --nfs \
diff --git a/docker/Dockerfile.gantry b/docker/Dockerfile.gantry
index 1387ebe66..14a902870 100644
--- a/docker/Dockerfile.gantry
+++ b/docker/Dockerfile.gantry
@@ -4,7 +4,7 @@
 # To build and push the image to Beaker, run 'make gantry-image'.
 # To test the image after pushing to Beaker, run 'make gantry-test'.
 
-FROM dolma-base
+FROM dolma-torch2-base
 
 WORKDIR /stage
 
diff --git a/docker/Dockerfile.test b/docker/Dockerfile.test
index eb301a845..35614df8b 100644
--- a/docker/Dockerfile.test
+++ b/docker/Dockerfile.test
@@ -4,7 +4,7 @@
 #
 # To build and push the image to Beaker, run 'make test-image'.
 
-FROM dolma-base
+FROM dolma-torch2-base
 
 COPY scripts/test_entrypoint.sh /entrypoint.sh
 RUN chmod +x /entrypoint.sh