include clone commands in setup steps (#62)

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

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]'",

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,20 +300,17 @@ 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 = [
             "#!/usr/bin/env bash",
             "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,
         ]