feat: Allow using custom containers with pre-built dependencies

.gitattributes

@@ -0,0 +1 @@
+uv.lock -diff

.github/workflows/ci_mocked_api.yaml

@@ -7,81 +7,23 @@ on:
   pull_request:
 
 jobs:
-  formatting:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out the code
-        uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.11' 
-
-      - name: Install poetry
-        run: pip install poetry
-
-      - name: Determine dependencies
-        run: poetry lock
-
-      - name: Install Dependencies using Poetry
-        run: poetry install
-
-      - name: Check formatting
-        run: poetry run black --check .
-
-  linting:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out the code
-        uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v4
-        with:
-          python-version: '3.11' 
-
-      - name: Install poetry
-        run: pip install poetry
-
-      - name: Determine dependencies
-        run: poetry lock
-
-      - name: Install Dependencies using Poetry
-        run: poetry install
-
-      - name: Check code
-        run: poetry run flake8
-
   testing-mocked-api:
     runs-on: ubuntu-latest
     permissions:
-      contents: 'read'
-      id-token: 'write'
+      contents: "read"
+      id-token: "write"
     steps:
       - uses: actions/checkout@v4
 
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
-          python-version: '3.11' 
-
-      - name: Install poetry
-        run: pip install poetry
+          python-version: "3.11"
 
-      - name: Determine dependencies
-        run: poetry lock
+      - name: Install uv based on the version defined in pyproject.toml
+        uses: astral-sh/setup-uv@v6
 
       - name: Install dependencies
-        run: poetry install
-
-      - id: 'auth'
-        uses: 'google-github-actions/auth@v1'
-        with:
-          credentials_json: '${{ secrets.GCP_SA_KEY }}'
-
-      - name: 'Set up Cloud SDK'
-        uses: 'google-github-actions/setup-gcloud@v1'
-
-      - name: 'Use gcloud CLI'
-        run: 'gcloud info'
+        run: uv sync --all-extras
 
       - name: Setup minio
         uses: comfuture/minio-action@v1
@@ -92,7 +34,7 @@ jobs:
 
       - name: Run pytest
         run: |
-          poetry run coverage run -m pytest tests/tests_mocked_api.py -v
+          uv run coverage run -m pytest tests/ -v
 
       - name: Run Coverage
-        run: poetry run coverage report -m
+        run: uv run coverage report -m

.github/workflows/release-please.yml

@@ -22,30 +22,22 @@ jobs:
     needs: release-please
     if: ${{ needs.release-please.outputs.release_created }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
 
-      - uses: actions/setup-python@v2
+      - uses: actions/setup-python@v5
         with:
           python-version: "3.11"
 
-      - name: Install poetry
-        run: pip install poetry
+      - name: Install uv based on the version defined in pyproject.toml
+        uses: astral-sh/setup-uv@v6
 
-      - name: Determine dependencies
-        run: poetry lock
-
-      - uses: actions/setup-python@v4
-        with:
-          python-version: "3.11"
-          cache: poetry
-
-      - name: Install Dependencies using Poetry
-        run: |
-          pip install connection-pool # because it is incompatible with poetry
-          poetry install
+      - name: Install dependencies
+        run: uv sync --all-extras
 
       - name: Publish to PyPi
         env:
           PYPI_USERNAME: __token__
           PYPI_PASSWORD: ${{ secrets.PYPI_TOKEN }}
-        run: poetry publish --build --username $PYPI_USERNAME --password $PYPI_PASSWORD
+        run: |
+          uv build
+          uv publish --build --username "$PYPI_USERNAME" --password "$PYPI_PASSWORD"

.pre-commit-config.yaml

@@ -0,0 +1,30 @@
+repos:
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: v0.12.7
+    hooks:
+      - args: [--fix, --exit-non-zero-on-fix]
+        id: ruff-check
+      - id: ruff-format
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.17.1
+    hooks:
+      - id: mypy
+        additional_dependencies:
+          - types-requests
+          - types-protobuf>=5,<6
+          - types-cachetools
+          - numpy
+  - repo: https://github.com/rhysd/actionlint
+    rev: v1.7.7
+    hooks:
+      - id: actionlint
+        args:
+          - "-shellcheck="
+          - "-ignore"
+          - 'input "attestations" is not defined in action'
+          - "-ignore"
+          - '"github.head_ref" is potentially untrusted'
+  - repo: https://github.com/hadialqattan/pycln
+    rev: v2.5.0
+    hooks:
+      - id: pycln

example/hello-world-cos/Dockerfile

@@ -0,0 +1,43 @@
+# syntax=docker/dockerfile:1.4@sha256:9ba7531bd80fb0a858632727cf7a112fbfd19b17e94c4e84ced81e24ef1a0dbc
+
+FROM mirror.gcr.io/mambaorg/micromamba:2.3-debian13-slim@sha256:3f92db74f6780f60c9a2c6d983c76360a1d4cfdcb597395ab9e64747f067b15b
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+
+# Revert to being root temporarily
+USER 0
+# Install base dependencies
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
+    rm -f /etc/apt/apt.conf.d/docker-clean \
+    && echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | tee /etc/apt/apt.conf.d/keep-cache \
+    && apt-get -qq update -y \
+    && apt-get -qq install --no-install-recommends -y curl git ca-certificates git python3-pip build-essential gnupg
+
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg
+RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
+    rm -f /etc/apt/apt.conf.d/docker-clean \
+    && apt-get -qq update -y \
+    && apt-get -qq install --no-install-recommends -y google-cloud-cli
+
+# Create mount path used by gcsfuse to mount data
+RUN mkdir -p /mnt/ && chown mambauser:mambauser /mnt/ && chmod g+w /mnt/
+
+# Create and set permissions for the share directory
+RUN mkdir -p /mnt/share && chown mambauser:mambauser /mnt/share && chmod g+w /mnt/share
+VOLUME /mnt/share
+
+# Back to non-root to install dependencies
+USER mambauser
+
+
+COPY environment.yml .
+RUN micromamba env create --download-only --yes -f ./environment.yml
+RUN micromamba env create --yes -f ./environment.yml
+
+
+COPY docker-entrypoint.sh /
+ENTRYPOINT ["/docker-entrypoint.sh"]

example/hello-world-cos/README.md

@@ -1,20 +1,58 @@
 # Hello World COS
 
 The example is for the container operating system (COS), which basically means running inside a container
-instead of on a native virtual machine. This also means using the default Snakemake image. Here is how to see images in the project:
+instead of on a native virtual machine. This requires building a custom image. Here is how to see images in
+the project:
 
 ```bash
 gcloud compute images list \
-    --project=batch-custom-image \
-    --no-standard-images
+      --project=batch-custom-image \
+      --no-standard-images
+```
+
+An example rule in snakemake:
+
+```snakemake
+rule interpolated_flat_fluvial:
+    input:
+        interpolated_layers="input.path"
+    output:
+        touch("output.path")
+    resources:
+        googlebatch_service_account="[email protected]",
+        googlebatch_image_family="batch-cos-stable-official",
+        googlebatch_image_project="batch-custom-image",
+        googlebatch_entrypoint="/docker-entrypoint.sh",
+        googlebatch_machine_type="e2-standard-4",
+        googlebatch_container="europe-docker.pkg.dev/myproject/my_image_name:latest",
+        googlebatch_container_dependencies_installed=True,
+        googlebatch_work_tasks=1,
+        googlebatch_work_tasks_per_node=2,
+    script:
+        "relative/path/to/script.py"
 ```
 
 Here is an example command:
 
 ```bash
 GOOGLE_PROJECT=myproject
-snakemake --jobs 1 --executor googlebatch --googlebatch-image-family batch-cos-stable-official --googlebatch-region us-central1 --googlebatch-image-project batch-custom-image --googlebatch-project ${GOOGLE_PROJECT} --default-storage-provider s3 --default-storage-prefix s3://my-snakemake-testing
+snakemake \
+  --jobs 1 \
+  --executor googlebatch \
+  --googlebatch-image-family batch-cos-stable-official \
+  --googlebatch-region us-central1 \
+  --googlebatch-image-project batch-custom-image \
+  --googlebatch-project ${GOOGLE_PROJECT} \
+  --container-image europe-docker.pkg.dev/${GOOGLE_PROJECT}/my_image_name:latest \
+  --preemptible-rules name_of_preemptible_rule \
+  --default-storage-provider gcs \
+  --default-storage-prefix gs://my_bucket_name \
+  --verbose \
+  my_rule_name
 ```
 
-See [this link](https://cloud.google.com/batch/docs/vm-os-environment-overview#supported_vm_os_images) for how to find a compatible COS image project and family.
-You can also see information [here](https://cloud.google.com/batch/docs/view-os-images),
+See the [Dockerfile](./Dockerfile) and [entrypoint script](./docker-entrypoint.sh) for an example of using micromamba to
+build an image with all dependencies installed.
+
+See [this link](https://cloud.google.com/batch/docs/vm-os-environment-overview#supported_vm_os_images) for how to find a
+compatible COS image project and family. You can also see [OS images documentation](https://cloud.google.com/batch/docs/view-os-images),

example/hello-world-cos/docker-entrypoint.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env sh
+
+set -ex
+
+micromamba run -n my_environment_name bash -ex -c "$@"