train_act/Dockerfile at main · ioai-tech/train_act · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
# syntax=docker/dockerfile:1.7

ARG PYTORCH_IMAGE=pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime
ARG ACT_GIT_URL=https://github.com/tonyzhaozh/act.git
ARG ACT_GIT_REF=742c753c0d4a5d87076c8f69e5628c79a8cc5488

FROM debian:bookworm-slim AS act-source
ARG ACT_GIT_URL
ARG ACT_GIT_REF
RUN apt-get update \
    && apt-get install -y --no-install-recommends ca-certificates git \
    && rm -rf /var/lib/apt/lists/* \
    && git clone --filter=blob:none "${ACT_GIT_URL}" /opt/act \
    && git -C /opt/act checkout --detach "${ACT_GIT_REF}" \
    && test "$(git -C /opt/act rev-parse HEAD)" = "${ACT_GIT_REF}" \
    && rm -rf /opt/act/.git

FROM ${PYTORCH_IMAGE} AS runtime

ARG ACT_GIT_REF
ARG BUILD_VERSION=dev
ARG VCS_REF=unknown

LABEL org.opencontainers.image.title="IOAI ACT Trainer" \
      org.opencontainers.image.description="ACT training for LeRobot v2/v3 datasets" \
      org.opencontainers.image.source="https://github.com/ioai-tech/train_act" \
      org.opencontainers.image.documentation="https://github.com/ioai-tech/train_act#readme" \
      org.opencontainers.image.licenses="MIT" \
      org.opencontainers.image.version="${BUILD_VERSION}" \
      org.opencontainers.image.revision="${VCS_REF}" \
      io.ioai.train-act.upstream-commit="${ACT_GIT_REF}"

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    WANDB_MODE=disabled \
    TORCH_HOME=/opt/torch-cache \
    PYTHONPATH=/opt/act \
    PIP_BREAK_SYSTEM_PACKAGES=1

WORKDIR /app/train_act

COPY requirements.txt docker/check_dependencies.py /tmp/
RUN --mount=type=cache,target=/root/.cache/pip \
    python -m pip install -r /tmp/requirements.txt \
    && python /tmp/check_dependencies.py \
    && rm /tmp/requirements.txt

# This layer changes only when the PyTorch/torchvision base changes. torchvision
# validates the f37072fd hash prefix while downloading the ImageNet checkpoint.
RUN mkdir -p "${TORCH_HOME}/hub/checkpoints" \
    && python - <<'PY'
from pathlib import Path
from torchvision.models import ResNet18_Weights, resnet18

resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
path = Path("/opt/torch-cache/hub/checkpoints/resnet18-f37072fd.pth")
if not path.is_file() or path.stat().st_size < 40_000_000:
    raise RuntimeError(f"Invalid ResNet18 cache: {path}")
print(f"Cached {path} ({path.stat().st_size} bytes)")
PY
RUN chmod -R a+rX "${TORCH_HOME}"

COPY --from=act-source /opt/act /opt/act
COPY src/act_overrides/ /opt/act/
RUN python -m pip install --no-deps /opt/act/detr

COPY pyproject.toml README.md LICENSE THIRD_PARTY_NOTICES.md ./
COPY src/train_act ./src/train_act
RUN --mount=type=cache,target=/root/.cache/pip \
    python -m pip install --no-deps . \
    && python /tmp/check_dependencies.py \
    && python -c "import train_act; from train_act.cli import build_parser; build_parser()" \
    && test -f "${TORCH_HOME}/hub/checkpoints/resnet18-f37072fd.pth" \
    && rm /tmp/check_dependencies.py

COPY entrypoint.sh /usr/local/bin/train-act-entrypoint
RUN chmod 0755 /usr/local/bin/train-act-entrypoint \
    && mkdir -p /data/input /data/output

ENTRYPOINT ["/usr/local/bin/train-act-entrypoint"]