Merge pull request #1 from epoch-research/better-layering-2

Improved use of layering to reduce image size by 6-10x

Author: tadamcz

.github/workflows/build.yaml

@@ -1,15 +1,22 @@
-
 name: Build images
 
 env:
   GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
 
 on:
   workflow_dispatch:
-
+    inputs:
+      architecture:
+        description: 'CPU architecture (x64 or arm64)'
+        required: true
+        type: choice
+        options:
+          - x64
+          - arm64
 jobs:
   build:
-    runs-on: ubuntu-latest
+    # The script will detect the architecture and build the images for that architecture
+    runs-on: epoch-research-${{ inputs.architecture }}-32core-128GB-1200GB
     defaults:
       run:
         shell: bash -l {0}
@@ -18,15 +25,23 @@ jobs:
         uses: actions/checkout@v2
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.9'
+          python-version: '3.11'
       - name: Install uv
         run: |
           curl -LsSf https://astral.sh/uv/install.sh | sh
       - name: Install dependencies
         run: |
             uv pip install --python ${Python_ROOT_DIR} '.'
-      - name: Install dev dependencies
-        run: |
-            uv pip install --python ${Python_ROOT_DIR} pytest pytest-cov
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
       - name: Run build script
-        run: python build.py
+        run: python build.py
+      - name: Store Docker build logs as artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: build-logs-${{ inputs.architecture }}
+          path: logs

.github/workflows/pytest.yaml

@@ -34,7 +34,7 @@ jobs:
         uses: actions/checkout@v2
       - uses: actions/setup-python@v5
         with:
-          python-version: '3.9'
+          python-version: '3.11'
       - name: Install uv
         run: |
           curl -LsSf https://astral.sh/uv/install.sh | sh

build.py

@@ -0,0 +1,33 @@
+from swebench.harness import prepare_images
+import datasets
+
+
+def build(dataset_name, repo, limit=None):
+    dataset = datasets.load_dataset(dataset_name)
+    split = 'test'
+    dataset = dataset[split]
+
+    if repo:
+        dataset = dataset.filter(lambda x: x['repo'] == repo)
+
+    instance_ids = dataset['instance_id']
+
+    if limit:
+        instance_ids = instance_ids[:limit]
+
+    print(f"Building images for {len(instance_ids)} instances")
+
+    prepare_images.main(
+        dataset_name=dataset_name,
+        split=split,
+        instance_ids=instance_ids,
+        max_workers=32,
+        force_rebuild=False,
+        open_file_limit=8192,
+    )
+
+
+if __name__ == '__main__':
+    dataset_name = 'princeton-nlp/SWE-bench'
+    repo: str | None = None
+    build(dataset_name, repo)