From 48892ee43b3361578f4708684682f595f040df70 Mon Sep 17 00:00:00 2001
From: Pete Walsh <petew@allenai.org>
Date: Mon, 7 Oct 2024 16:14:51 -0700
Subject: [PATCH] include clone commands in setup steps (#62)

---
 CHANGELOG.md                         |  4 ++++
 src/olmo_core/internal/experiment.py |  4 ++++
 src/olmo_core/launch/beaker.py       | 31 ++++++++++++++--------------
 3 files changed, 24 insertions(+), 15 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..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,10 @@ class BeakerLaunchConfig(Config):
     A description for the experiment.
     """
 
-    setup_steps: List[str] = field(
-        default_factory=lambda: [
-            "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 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
@@ -296,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 = [
@@ -307,9 +311,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,
         ]