Skip to content

Commit 6360856

Browse files
authored
Add Docker development environment (#50)
* Add Dockerfile * try speeding up Docker build * Free up some disk-space on runner first * Add smoke test for Docker run * Run smoke-test quickly * Fix CUDA env variable settings * Run part of the smoke test inside to simplify across nodes * Remove tool cache too for space * Refactoring and hopefully slimming down
1 parent cf1ce78 commit 6360856

File tree

6 files changed

+283
-0
lines changed

6 files changed

+283
-0
lines changed

.dockerignore

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# Build artifacts
2+
build/
3+
dist/
4+
*.egg-info
5+
**/__pycache__
6+
.venv/
7+
*.pyc
8+
*.so
9+
10+
# Version control
11+
.git/
12+
.github/
13+
14+
# Documentation
15+
docs/_build/
16+
docs/generated/
17+
18+
# Testing output
19+
tests/output/
20+
.benchmarks/
21+
.pytest_cache/
22+
.coverage
23+
.coverage.*
24+
htmlcov/
25+
26+
# IDE files
27+
.vscode/
28+
.idea/
29+
*.code-workspace
30+
31+
# Environment and logs
32+
.env
33+
*.log
34+
.python-version
35+
36+
# Other artifacts
37+
*.h5
38+
!ATS*
39+
wheelhouse/

.github/workflows/docker-build.yml

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
name: Build and push Docker image
2+
3+
on:
4+
workflow_dispatch:
5+
pull_request:
6+
paths:
7+
- "Dockerfile"
8+
- "docker-compose.yml"
9+
- ".dockerignore"
10+
- ".github/workflows/docker-build.yml"
11+
push:
12+
branches:
13+
- main
14+
release:
15+
types:
16+
- published
17+
18+
env:
19+
REGISTRY: ghcr.io
20+
IMAGE_NAME: ${{ github.repository }}-dev
21+
22+
jobs:
23+
build-and-push:
24+
name: Build and push Docker image
25+
runs-on: ubuntu-22.04
26+
permissions:
27+
contents: read
28+
packages: write
29+
outputs:
30+
image-tag: ${{ steps.export-tag.outputs.tag }}
31+
32+
steps:
33+
- name: Free Disk Space (Ubuntu)
34+
uses: jlumbroso/free-disk-space@main
35+
with:
36+
# aggressively free disk space because
37+
# the container is so large
38+
tool-cache: true
39+
40+
- name: Check out repository
41+
uses: actions/checkout@v4
42+
43+
- name: Set up Docker Buildx
44+
uses: docker/setup-buildx-action@v3
45+
46+
- name: Log in to Container Registry
47+
uses: docker/login-action@v3
48+
with:
49+
registry: ${{ env.REGISTRY }}
50+
username: ${{ github.actor }}
51+
password: ${{ secrets.GITHUB_TOKEN }}
52+
53+
- name: Extract metadata for Docker
54+
id: meta
55+
uses: docker/metadata-action@v5
56+
with:
57+
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
58+
tags: |
59+
type=ref,event=branch
60+
type=ref,event=pr
61+
type=semver,pattern={{version}}
62+
type=semver,pattern={{major}}.{{minor}}
63+
type=semver,pattern={{major}}
64+
type=sha
65+
type=raw,value=latest,enable={{is_default_branch}}
66+
67+
- name: Build and push Docker image
68+
id: build
69+
uses: docker/build-push-action@v6
70+
with:
71+
context: .
72+
push: ${{ github.event_name != 'pull_request' }}
73+
load: ${{ github.event_name == 'pull_request' }}
74+
tags: ${{ steps.meta.outputs.tags }}
75+
labels: ${{ steps.meta.outputs.labels }}
76+
cache-from: type=gha
77+
cache-to: type=gha,mode=max
78+
79+
- name: Export image tag
80+
id: export-tag
81+
run: |
82+
# Get first tag from metadata output
83+
TAG=$(echo "${{ steps.meta.outputs.tags }}" | head -n1)
84+
echo "tag=$TAG" >> $GITHUB_OUTPUT
85+
86+
- name: Smoke test image (no GPU)
87+
run: |
88+
docker run --rm "${{ steps.export-tag.outputs.tag }}" \
89+
python -c "import mach; print(f'✓ mach version: {mach.__version__}')"
90+
91+
smoke-test-gpu:
92+
name: Smoke test Docker image on GPU
93+
needs: build-and-push
94+
if: github.event_name != 'pull_request'
95+
runs-on: linux-x64-nvidia-gpu-t4
96+
permissions:
97+
packages: read
98+
timeout-minutes: 10
99+
100+
container:
101+
image: ${{ needs.build-and-push.outputs.image-tag }}
102+
credentials:
103+
username: ${{ github.actor }}
104+
password: ${{ secrets.GITHUB_TOKEN }}
105+
options: --gpus all
106+
107+
steps:
108+
- name: Test mach import and version
109+
run: |
110+
python -c "import mach; print(f'✓ mach version: {mach.__version__}')"
111+
112+
- name: Test CUDA backend availability
113+
run: |
114+
python -c "import mach._cuda_impl; print('✓ CUDA backend loaded successfully')"

Dockerfile

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# syntax=docker/dockerfile:1
2+
3+
# Development environment for mach-beamform
4+
# Provides CUDA compilation without requiring local CUDA installation
5+
6+
ARG CUDA_VERSION=12.6.3
7+
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
8+
9+
# Avoid interactive prompts during package installation
10+
ENV DEBIAN_FRONTEND=noninteractive
11+
12+
# Install build dependencies
13+
RUN apt-get update && apt-get install -y --no-install-recommends \
14+
gcc \
15+
g++ \
16+
make \
17+
cmake \
18+
ninja-build \
19+
curl \
20+
ca-certificates \
21+
&& rm -rf /var/lib/apt/lists/*
22+
23+
# Install uv (will automatically install Python when needed)
24+
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
25+
26+
# Set CUDA environment variables
27+
ENV CUDA_HOME=/usr/local/cuda
28+
ENV PATH="${CUDA_HOME}/bin:${PATH}"
29+
ENV LD_LIBRARY_PATH="${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}"
30+
31+
# Silence warning about not being able to use hard links with cache mount
32+
ENV UV_LINK_MODE=copy
33+
34+
# Shared flags for uv sync: install all optional extras but exclude heavy dev-dependencies
35+
ENV UV_SYNC_FLAGS="--frozen --all-extras --no-dev --no-group test --no-group profile --no-group build --no-group array --no-group compare --no-group docs"
36+
37+
# Set working directory
38+
WORKDIR /workspace
39+
40+
# Copy dependency files first for better layer caching
41+
# Dependencies only rebuild when these files change
42+
COPY pyproject.toml uv.lock ./
43+
44+
# Install dependencies with cache mount
45+
# This layer is cached and reused when only source code changes
46+
# Install all optional dependencies but exclude heavy dev-dependencies
47+
RUN --mount=type=cache,target=/root/.cache/uv \
48+
uv sync $UV_SYNC_FLAGS --no-install-project
49+
50+
# Copy the rest of the project
51+
# Source code changes won't trigger dependency reinstall
52+
COPY . .
53+
54+
# Install the project itself (fast, no dependency downloads)
55+
RUN --mount=type=cache,target=/root/.cache/uv \
56+
uv sync $UV_SYNC_FLAGS
57+
58+
# Add virtual environment to PATH so Python and installed packages are available
59+
ENV PATH="/workspace/.venv/bin:${PATH}"
60+
61+
# OCI labels
62+
LABEL org.opencontainers.image.title="mach-beamform-dev" \
63+
org.opencontainers.image.description="Development environment for ultrafast GPU-accelerated beamforming" \
64+
org.opencontainers.image.source="https://github.com/Forest-Neurotech/mach" \
65+
org.opencontainers.image.vendor="Forest Neurotech"
66+
67+
# Default command: interactive bash shell
68+
CMD ["/bin/bash"]

Makefile

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,4 +109,12 @@ clean: ## Cleans build artifacts
109109
@echo "Cleaning build artifacts..."
110110
rm -rf dist/ build/ mach.*.so
111111

112+
.PHONY: docker-build
113+
docker-build: ## Builds the Docker development image
114+
docker compose build
115+
116+
.PHONY: docker-dev
117+
docker-dev: ## Runs the development container
118+
docker compose run --rm dev
119+
112120
.DEFAULT_GOAL := help

README.md

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,34 @@ Build prerequisites:
4949
* `gcc >= 8`
5050
* `nvcc >= 11.0`
5151

52+
### Docker Development
53+
54+
Compile and test without installing the CUDA *toolkit* using our Docker development environment.
55+
56+
**Prerequisites:**
57+
* Docker Engine with [nvidia-container-toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)
58+
* CUDA-capable GPU with driver >= 12.3
59+
60+
**Quick start:**
61+
62+
```bash
63+
# Build and start development container
64+
docker compose run --rm dev
65+
66+
# Or use make shortcuts
67+
make docker-build # Build image (first time: ~2-3 min, rebuilds: ~30s)
68+
make docker-dev # Run container
69+
```
70+
71+
**Inside the container:**
72+
73+
```bash
74+
make compile # Compile CUDA extension
75+
make test # Run tests
76+
```
77+
78+
Your source code is mounted from the host, so you can edit files locally and compile in the container. Build artifacts (`.venv/` and `build/`) are stored in anonymous volumes to avoid permission issues. Dependencies are pre-installed in the image and cached, so rebuilds are fast when only source code changes.
79+
5280
## Examples
5381

5482
Try our [examples](https://forest-neurotech.github.io/mach/examples/):

docker-compose.yml

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
services:
2+
# Default development service
3+
# Mount project as read-write, mask build artifacts with volumes
4+
dev:
5+
build:
6+
context: .
7+
dockerfile: Dockerfile
8+
image: mach-beamform-dev:latest
9+
stdin_open: true
10+
tty: true
11+
working_dir: /workspace
12+
13+
# Mount project directory, but use volumes for artifacts to avoid host conflicts
14+
volumes:
15+
- .:/workspace
16+
- /workspace/.venv # Anonymous volume masks .venv from host
17+
- /workspace/build # Anonymous volume masks build from host
18+
19+
# GPU access configuration
20+
deploy:
21+
resources:
22+
reservations:
23+
devices:
24+
- driver: nvidia
25+
count: all
26+
capabilities: [gpu]

0 commit comments

Comments
 (0)