From 080bf31c0dbb207a9b52c5ddcb51d08867e78442 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Mon, 7 Oct 2024 15:52:22 -0700
Subject: [PATCH 1/2] include clone commands in setup steps

---
 CHANGELOG.md                         |  4 ++++
 src/olmo_core/internal/experiment.py |  4 ++++
 src/olmo_core/launch/beaker.py       | 10 +++++-----
 3 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0d58ed86..3b161d55 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- `BeakerLaunchConfig.setup_steps` should now include steps to clone your repo (which it will by default). This change allows support for private repos.
+
 ## [v1.4.0](https://github.com/allenai/OLMo-core/releases/tag/v1.4.0) - 2024-10-02
 
 ### Changed
diff --git a/src/olmo_core/internal/experiment.py b/src/olmo_core/internal/experiment.py
index b45357eb..45b798a1 100644
--- a/src/olmo_core/internal/experiment.py
+++ b/src/olmo_core/internal/experiment.py
@@ -150,6 +150,10 @@ def build_common_components(
             BeakerEnvSecret(name="WEKA_ENDPOINT_URL", secret="WEKA_ENDPOINT_URL"),
         ],
         setup_steps=[
+            # Clone repo.
+            'git clone "${REPO_URL}" .',
+            'git checkout "${GIT_REF}"',
+            "git submodule update --init --recursive",
             # Setup python environment.
             "conda shell.bash activate base",
             "pip install -e '.[all]'",
diff --git a/src/olmo_core/launch/beaker.py b/src/olmo_core/launch/beaker.py
index f3f7dd83..2a3c27d7 100644
--- a/src/olmo_core/launch/beaker.py
+++ b/src/olmo_core/launch/beaker.py
@@ -122,14 +122,17 @@ class BeakerLaunchConfig(Config):
 
     setup_steps: List[str] = field(
         default_factory=lambda: [
+            'git clone "${REPO_URL}" .',
+            'git checkout "${GIT_REF}"',
+            "git submodule update --init --recursive",
             "conda shell.bash activate base",
             "pip install -e '.[all]'",
             "pip freeze",
         ]
     )
     """
-    A list of shell commands to run for installing dependencies, running arbitrary scripts,
-    and other setup steps.
+    A list of shell commands to run for cloning your repo, installing dependencies,
+    and other arbitrary setup steps.
     """
 
     beaker_image: str = OLMoCoreBeakerImage.stable
@@ -307,9 +310,6 @@ def build_experiment_spec(self, torchrun: bool = True) -> ExperimentSpec:
             "set -exuo pipefail",
             "mkdir -p /olmo-core-runtime",
             "cd /olmo-core-runtime",
-            'git clone "${REPO_URL}" .',
-            'git checkout "${GIT_REF}"',
-            "git submodule update --init --recursive",
             *self.setup_steps,
         ]
 

From f1ca4617fc37864305b9d11589e87472b2b29d1e Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Mon, 7 Oct 2024 16:00:34 -0700
Subject: [PATCH 2/2] allow private repos, but raise error if not handled by
 user

---
 src/olmo_core/launch/beaker.py | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/src/olmo_core/launch/beaker.py b/src/olmo_core/launch/beaker.py
index 2a3c27d7..d060dfa3 100644
--- a/src/olmo_core/launch/beaker.py
+++ b/src/olmo_core/launch/beaker.py
@@ -84,6 +84,16 @@ class BeakerWekaBucket(Config):
     mount: str
 
 
+DEFAULT_SETUP_STEPS = (
+    'git clone "${REPO_URL}" .',
+    'git checkout "${GIT_REF}"',
+    "git submodule update --init --recursive",
+    "conda shell.bash activate base",
+    "pip install -e '.[all]'",
+    "pip freeze",
+)
+
+
 @dataclass
 class BeakerLaunchConfig(Config):
     """
@@ -120,16 +130,7 @@ class BeakerLaunchConfig(Config):
     A description for the experiment.
     """
 
-    setup_steps: List[str] = field(
-        default_factory=lambda: [
-            'git clone "${REPO_URL}" .',
-            'git checkout "${GIT_REF}"',
-            "git submodule update --init --recursive",
-            "conda shell.bash activate base",
-            "pip install -e '.[all]'",
-            "pip freeze",
-        ]
-    )
+    setup_steps: List[str] = field(default_factory=lambda: list(DEFAULT_SETUP_STEPS))
     """
     A list of shell commands to run for cloning your repo, installing dependencies,
     and other arbitrary setup steps.
@@ -299,10 +300,10 @@ def build_experiment_spec(self, torchrun: bool = True) -> ExperimentSpec:
         # Get repository account, name, and current ref.
         github_account, github_repo, git_ref, is_public = ensure_repo(self.allow_dirty)
 
-        if not is_public:
+        if not is_public and self.setup_steps == DEFAULT_SETUP_STEPS:
             raise OLMoConfigurationError(
-                "Only public repositories are supported at the moment. "
-                "Please use beaker-gantry to launch jobs with private repos."
+                "It looks like your repository is private and private repositories will require "
+                "custom 'setup_steps' in order to clone the repo."
             )
 
         entrypoint_script = [