FastGen v0.2.0

.env.example

@@ -0,0 +1,42 @@
+# FastGen environment configuration template.
+#
+# Copy to `.env` at the repository root and adjust. `fastgen/__init__.py` loads `.env`
+# with python-dotenv when the package is imported, before the config modules read the
+# environment, so these values become the defaults for the roots below. Variables already
+# exported in the shell (e.g. by the job scheduler or CI) always take precedence. `${VAR}` references
+# expand from earlier entries and from the shell. Keep secrets (W&B, Hugging Face, S3 keys)
+# out of this file: export them or use the config's credentials_dir (./credentials/ by
+# default), as described in fastgen/configs/README.md. Non-empty credential files loaded
+# during setup override their matching environment variables.
+
+# Root for run outputs (checkpoints, logs); also the base of the DATA/MODEL/.cache roots.
+FASTGEN_OUTPUT_ROOT=./FASTGEN_OUTPUT
+
+# Dataset root and root for local pretrained checkpoints (mirrored S3/HTTP files land here).
+DATA_ROOT_DIR=${FASTGEN_OUTPUT_ROOT}/DATA
+CKPT_ROOT_DIR=${FASTGEN_OUTPUT_ROOT}/MODEL
+
+# Remote prepared-dataset prefix (S3 or HTTP(S)). S3 credentials use the standard AWS
+# environment or s3.json in the config's credentials_dir (default ./credentials/).
+# REMOTE_DATA_ROOT_DIR=s3://data
+
+# Hugging Face cache. Defaults to $FASTGEN_OUTPUT_ROOT/.cache; point it at a shared,
+# pre-populated cache to reuse downloads. Set LOCAL_FILES_ONLY=true to run offline.
+# HF_HOME=/path/to/shared/hf_cache
+LOCAL_FILES_ONLY=false
+
+# Reduce allocator fragmentation in long and multi-forward GPU workloads. Keep both equal:
+# older PyTorch reads the CUDA name, newer PyTorch reads the generic name.
+PYTORCH_ALLOC_CONF=expandable_segments:True
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+# One intra-op CPU thread per process. FastGen's main process is GPU-bound, so extra
+# OpenMP workers only spin against the DataLoader workers (no gain measured at 16 threads),
+# and one thread keeps CPU reductions deterministic. torchrun sets this itself for
+# nproc_per_node > 1; the entry covers single-process launches. Raise to 2-4 only when a
+# profile shows the main process CPU-bound on tensor work (e.g. num_workers=0 preprocessing).
+OMP_NUM_THREADS=1
+
+# OpenAI-compatible endpoint for the VLM log module (VLMLogger); the bearer token comes from
+# OPENAI_API_KEY or vlm_api_token.txt in the config's credentials_dir.
+# VLM_BASE_URL=https://api.openai.com/v1

.github/workflows/ci.yml

@@ -8,13 +8,15 @@ on:
 
 env:
   WANDB_MODE: disabled
+  # Keep pip's cache inside the workspace so the actions/cache step below can save and restore it.
+  PIP_CACHE_DIR: ${{ github.workspace }}/.cache/pip
 
 jobs:
   lint:
     runs-on: [FG_Runner]
 
     container:
-      image: ghcr.io/juliusberner/fastgen:0.0.3c
+      image: ghcr.io/juliusberner/fastgen:0.0.6
       options: --entrypoint "" --gpus all --shm-size=8g
       credentials:
         username: ${{ github.actor }}
@@ -48,7 +50,7 @@ jobs:
     runs-on: [FG_Runner]
 
     container:
-      image: ghcr.io/juliusberner/fastgen:0.0.3c
+      image: ghcr.io/juliusberner/fastgen:0.0.6
       options: --entrypoint "" --gpus all --shm-size=8g
       credentials:
         username: ${{ github.actor }}
@@ -68,9 +70,6 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-pip-mypy-
 
-      - name: Increase file descriptor limit
-        run: ulimit -n 4096 || echo "Could not increase file descriptor limit"
-
       - name: Run Pytest
         run: make pytest
 
@@ -79,7 +78,7 @@ jobs:
     runs-on: [FG_Runner] 
 
     container:
-      image: ghcr.io/juliusberner/fastgen:0.0.3c
+      image: ghcr.io/juliusberner/fastgen:0.0.6
       options: --entrypoint "" --gpus all --shm-size=8g
       credentials:
         username: ${{ github.actor }}
@@ -100,4 +99,4 @@ jobs:
             ${{ runner.os }}-pip-mypy-
 
       - name: Install fastgen
-        run: make install-fastgen
+        run: make install-fastgen

.gitignore

@@ -17,8 +17,8 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
-lib64/
+/lib/
+/lib64/
 parts/
 sdist/
 var/
@@ -167,7 +167,8 @@ cython_debug/
 
 # large files
 *.zip
-core
+/core
+core.[0-9]*
 
 # temporal folder
 tmp/
@@ -179,6 +180,11 @@ credentials
 FASTGEN_OUTPUT/
 wandb/
 *.err
+*.out
 
 # vscode
-*.vscode/
+*.vscode/
+
+# home-directory files when the container runs with HOME=/workspace
+.local/
+.bash_history

CONTRIBUTING.md

@@ -41,7 +41,7 @@ make lint    # Check compliance
 
 ## Continuous Integration
 
-The GitHub CI pipeline (currently disabled) consists of three stages:
+The CI pipeline runs automatically on every push to `main` and every pull request targeting `main`:
 
 | Stage | Commands |
 |-------|----------|
@@ -51,6 +51,10 @@ The GitHub CI pipeline (currently disabled) consists of three stages:
 
 **Note:** Before submitting a pull request, ensure all checks pass locally.
 
+The test suite also checks that the built wheel includes configs and runtime
+data and works outside the checkout. Run only this check with
+`python -m pytest tests/test_packaging.py`.
+
 ## Pull Request Process
 
 1. Fork the repository (for external contributors) or create a feature branch.

Dockerfile

@@ -1,4 +1,5 @@
-FROM nvcr.io/nvidia/pytorch:25.06-py3
+# Base image 25.11-py3 ships torch 2.10.0a0 / CUDA 13.0 / Python 3.12
+FROM nvcr.io/nvidia/pytorch:25.11-py3
 
 # Defines the architecture (amd64 or arm64) automatically during build
 ARG TARGETARCH
@@ -31,32 +32,43 @@ RUN apt-get update --fix-missing && \
 RUN S5CMD_VERSION="2.1.0-beta.1" && \
     if [ "$TARGETARCH" = "arm64" ]; then \
         S5CMD_ARCH="Linux-arm64"; \
+        S5CMD_SHA256="f64dcaabbb837ce223d716c9e64b4829e5cd780f151a42f2d2cb7edccedb7b06"; \
     else \
         S5CMD_ARCH="Linux-64bit"; \
+        S5CMD_SHA256="de5efa37941ef83e892662aefc936603a1a8acbc5a6a9e003a71721a7140b610"; \
     fi && \
-    wget "https://github.com/peak/s5cmd/releases/download/v${S5CMD_VERSION}/s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz" && \
-    tar -xf "s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz" && \
+    S5CMD_ARCHIVE="s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz" && \
+    wget "https://github.com/peak/s5cmd/releases/download/v${S5CMD_VERSION}/${S5CMD_ARCHIVE}" && \
+    echo "${S5CMD_SHA256}  ${S5CMD_ARCHIVE}" | sha256sum --check --strict - && \
+    tar -xf "${S5CMD_ARCHIVE}" && \
     install s5cmd /usr/local/bin/ && \
-    rm s5cmd "s5cmd_${S5CMD_VERSION}_${S5CMD_ARCH}.tar.gz"
+    rm s5cmd "${S5CMD_ARCHIVE}"
 
 # Set build optimization flags
 ENV MAX_JOBS=4
-ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0" 
+ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0;10.0;12.0f+PTX"
 
 # Upgrade pip
-RUN pip install --upgrade pip setuptools wheel
-# numpy<2.0.0 is required for the preinstalled numba and thinc
-RUN pip install wandb[media] \
-    diffusers==0.35.1 \
-    transformers==4.49.0 \
+RUN pip install --no-cache-dir --root-user-action=ignore --upgrade pip setuptools wheel
+# numpy<2.0.0: required for the preinstalled numba and thinc
+# Pin the PEFT/TorchAO pair tested with diffusers 0.40 and torch 2.10.
+RUN pip install --no-cache-dir --root-user-action=ignore wandb[media] \
+    diffusers==0.40.0 \
+    transformers==5.15.0 \
+    peft==0.20.0 \
+    torchao==0.16.0 \
     accelerate \
     safetensors \
     scipy \
+    attrs \
+    matplotlib \
+    psutil \
     einops \
-    jupyter \
     hydra-core \
-    imageio \
+    pyyaml \
+    "imageio[ffmpeg]" \
     tqdm \
+    click \
     webdataset \
     av \
     loguru \
@@ -65,12 +77,15 @@ RUN pip install wandb[media] \
     ftfy \
     opencv-python-headless \
     sentencepiece \
-    "numpy<2.0.0" 
-RUN pip uninstall -y apex 
+    openai \
+    python-dotenv \
+    "numpy<2.0.0"
+RUN pip uninstall -y apex
+RUN pip uninstall -y pynvml && pip install --no-cache-dir --root-user-action=ignore nvidia-ml-py
 
 WORKDIR /workspace
 
-# Entry point
-RUN (printf '#!/bin/bash\nexec \"$@\"\n' >> /entry.sh) && chmod a+x /entry.sh
+# Entry point; without a command, the container opens a shell
+RUN printf '%s\n' '#!/bin/bash' 'exec "$@"' > /entry.sh && chmod a+x /entry.sh
 ENTRYPOINT ["/entry.sh"]
-
+CMD ["/bin/bash"]

MANIFEST.in

@@ -0,0 +1,5 @@
+# setup.py reads requirements.txt at metadata time, so a source distribution must ship it.
+include requirements.txt
+
+# Keep repository tests out of the source distribution.
+prune tests

Makefile

@@ -1,28 +1,42 @@
-sinclude .env
-
 all: help
 
 install-linters: ## install the linters
-	pip install black==23.10.0 ruff==0.6.9 mypy==1.9.0 types-psutil
+	pip install ruff==0.6.9 mypy==2.1.0
 
-install: install-linters
+install: install-linters ## install the development dependencies (alias for install-linters)
 
-mypy:
-	python3 -m mypy --check-untyped-defs --follow-imports=silent --exclude third_party/ train.py
+mypy: ## smoke-check train.py with mypy (fastgen/ modules are not analysed)
+	python3 -m mypy --check-untyped-defs --follow-imports=skip train.py
 
-lint:
-	python3 -m ruff format --check --exclude fastgen/third_party/
-	python3 -m ruff check --exclude fastgen/third_party/
+lint: ## check formatting and lint rules with ruff
+	python3 -m ruff format --check
+	python3 -m ruff check
 
-format:
-	python3 -m ruff format --exclude fastgen/third_party/
-	python3 -m ruff check --fix --exclude fastgen/third_party/
+format: ## format the code and auto-fix lint rules with ruff
+	python3 -m ruff format
+	python3 -m ruff check --fix
 
-install-fastgen:
+install-fastgen: ## install FastGen in editable mode
 	python3 -m pip install -e .
 
-pytest:
-	ulimit -n 4096 && python3 -m pytest --ignore=FASTGEN_OUTPUT --ignore=runs --ignore=tmp --ignore third_party
+# Run the suite in three sequential pytest invocations so memory accumulated by
+# tests that load large pretrained models (HF caches, CUDA allocator state,
+# diffusers/transformers globals) is reclaimed between batches. Each batch
+# always runs (so CI gets a complete picture) and the final exit code reflects
+# whether any batch failed.
+pytest: ## run the test suite in three batches
+	@rc=0; \
+	ulimit -n 4096 || echo "Could not raise file descriptor limit"; \
+	echo "===== batch 1/3: all tests except tests/test_network.py and tests/test_network_fsdp.py ====="; \
+	python3 -m pytest \
+		--ignore=FASTGEN_OUTPUT --ignore=runs --ignore=tmp \
+		--ignore=tests/test_network.py --ignore=tests/test_network_fsdp.py \
+		|| rc=$$?; \
+	echo "===== batch 2/3: tests/test_network.py ====="; \
+	python3 -m pytest tests/test_network.py || rc=$$?; \
+	echo "===== batch 3/3: tests/test_network_fsdp.py ====="; \
+	python3 -m pytest tests/test_network_fsdp.py || rc=$$?; \
+	exit $$rc
 
 .PHONY: help
 help: