diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..c308611 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,17 @@ +.git +__pycache__ +*.pyc +*.pyo +*.egg-info +build/ +dist/ +.env +.venv +env/ +venv/ +tests/output/ +tests/merged/ +logs/ +.mypy_cache/ +.pytest_cache/ +.ruff_cache/ \ No newline at end of file diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index af7a344..42d0bd1 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -4,12 +4,14 @@ on: push: branches: - main + pull_request: branches: - main + workflow_dispatch: env: - IMAGE_NAME: fabnemepfl/mmirage + IMAGE_NAMESPACE: fabnemepfl REGISTRY: docker.io jobs: @@ -22,13 +24,17 @@ jobs: path: docker/Dockerfile tag_base: amd64 name: mmirage-git + - platform: ubuntu-latest + path: docker/Dockerfile.cpu + tag_base: amd64 + name: mmirage-git-cpu - platform: ubuntu-24.04-arm - path: docker/Dockerfile + path: docker/Dockerfile.cpu tag_base: arm64 - name: mmirage-git + name: mmirage-git-cpu runs-on: ${{ matrix.platform }} - environment: docker + environment: ${{ github.event_name != 'pull_request' && 'docker' || 'docker-ci' }} steps: - name: Free space (ARM) @@ -62,6 +68,7 @@ jobs: uses: actions/checkout@v4 - name: Log in to DockerHub + if: github.event_name != 'pull_request' uses: docker/login-action@v3 with: username: ${{ secrets.DOCKERHUB_USERNAME }} @@ -75,7 +82,7 @@ jobs: with: context: . file: ${{ matrix.path }} - push: true + push: ${{ github.event_name != 'pull_request' }} tags: | - ${{ secrets.DOCKERHUB_USERNAME }}/${{ matrix.name }}:latest-${{ matrix.tag_base }} - ${{ secrets.DOCKERHUB_USERNAME }}/${{ matrix.name }}:${{ github.sha }}-${{ matrix.tag_base }} + ${{ env.REGISTRY }}/${{ env.IMAGE_NAMESPACE }}/${{ matrix.name }}:latest-${{ matrix.tag_base }} + ${{ env.REGISTRY }}/${{ env.IMAGE_NAMESPACE }}/${{ matrix.name }}:${{ github.sha }}-${{ matrix.tag_base }} diff --git a/README.md b/README.md index d35da70..0828347 100644 --- a/README.md +++ b/README.md @@ -10,18 +10,59 @@ MMIRAGE, which stands for **M**odular **M**ultimodal **I**ntelligent **R**eforma ## How to install -To install the library, you can clone it from GitHub and then use pip to install it directly. It is recommended to have already installed `torch` and `sglang` to take advantage of GPU acceleration. +To install the library, clone it from GitHub and install it with pip. The base +install does not include the local SGLang runtime: ```bash git clone git@github.com:EPFLiGHT/MMIRAGE.git pip install -e ./MMIRAGE ``` +Install the GPU extra when using the SGLang-backed `llm` processor for local +GPU inference: + +```bash +pip install -e './MMIRAGE[gpu]' +``` + For testing and scripts that make use of the library, it is advised to create a .env file: ```bash ./scripts/generate_env.sh ``` +## Docker + +### Build + +```bash +docker compose build +``` + +### Run + +```bash +docker compose run --rm mmirage --config configs/your_config.yaml +``` + +The container requires an NVIDIA GPU. The `docker-compose.yml` is configured to request GPU access, but the host must have: +- NVIDIA GPU drivers installed +- NVIDIA Container Toolkit / `nvidia-container-runtime` configured for Docker +- A recent Docker Engine and Docker Compose version with GPU support enabled + +Without these host-side prerequisites, `docker compose run` may fail to detect or use the GPU. + +### CPU-only + +The CPU image installs MMIRAGE without the GPU extra. It is suitable for +workflows that do not instantiate the SGLang-backed `llm` processor, and is +intended to support API-backed processors once they are available. Current +configs that use `type: llm` require the GPU image or an install with the +`[gpu]` extra. + +```bash +docker compose run --rm mmirage-cpu --config configs/your_config.yaml +``` + ## Key features - **Multimodal Support**: Process both text and images with vision-language models diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..d49759a --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,17 @@ +services: + mmirage: + build: + context: . + dockerfile: docker/Dockerfile + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + + mmirage-cpu: + build: + context: . + dockerfile: docker/Dockerfile.cpu diff --git a/docker/Dockerfile b/docker/Dockerfile index b7371ff..b2a4e2a 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,5 +1,11 @@ -FROM docker.io/lmsysorg/sglang:latest +FROM lmsysorg/sglang:latest + +ENV NVIDIA_VISIBLE_DEVICES=all +ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility COPY . /workspace/MMIRAGE WORKDIR /workspace/MMIRAGE -RUN pip install --no-cache-dir -e . + +RUN pip install --no-cache-dir .[gpu] + +ENTRYPOINT ["python3", "-m", "mmirage.shard_process"] diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu new file mode 100644 index 0000000..2529d00 --- /dev/null +++ b/docker/Dockerfile.cpu @@ -0,0 +1,8 @@ +FROM python:3.11-slim + +COPY . /workspace/MMIRAGE +WORKDIR /workspace/MMIRAGE + +RUN pip install --no-cache-dir . + +ENTRYPOINT ["python3", "-m", "mmirage.shard_process"] diff --git a/pyproject.toml b/pyproject.toml index 6ccc56a..76ba4f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,6 @@ authors = [{ name = "Meditron team" }] # Core runtime deps for your scripts dependencies = [ - "sglang>=0.5.2", "transformers>=4.46.0", "pyzmq", "uvloop<0.22; platform_system != 'Windows'", @@ -20,11 +19,8 @@ dependencies = [ "openai>=1.0.0", "partial_json_parser", "sentencepiece", - "sgl_kernel", - "compressed_tensors", "msgspec", "nest_asyncio", - "xgrammar", "PyYAML", "json-repair", "tqdm", @@ -38,10 +34,17 @@ dependencies = [ "jmespath", "jinja2>=3.0.0", "pillow>=9.0.0", + "typing_extensions>=4.5.0; python_version < '3.12'", "humanize>=4.0.0", ] [project.optional-dependencies] +gpu = [ + "sglang>=0.5.2", + "sgl_kernel", + "xgrammar", + "compressed_tensors", +] dev = [ "ruff>=0.5.0", "black>=24.3.0", diff --git a/src/mmirage/core/loader/jsonl.py b/src/mmirage/core/loader/jsonl.py index 6677b7c..93879a1 100644 --- a/src/mmirage/core/loader/jsonl.py +++ b/src/mmirage/core/loader/jsonl.py @@ -3,7 +3,13 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Dict, Optional, Union, override +from typing import Dict, Optional, Union + +try: + from typing import override +except ImportError: + from typing_extensions import override # type: ignore + from datasets import ( Dataset, DatasetDict, diff --git a/src/mmirage/core/process/processors/llm/llm_processor.py b/src/mmirage/core/process/processors/llm/llm_processor.py index 5107582..19cae67 100644 --- a/src/mmirage/core/process/processors/llm/llm_processor.py +++ b/src/mmirage/core/process/processors/llm/llm_processor.py @@ -9,7 +9,11 @@ from typing import Any, List, Tuple import jinja2 -import sglang as sgl +try: + import sglang as sgl + SGLANG_AVAILABLE = True +except ImportError: + SGLANG_AVAILABLE = False from transformers import AutoTokenizer from mmirage.core.process.base import BaseProcessor, ProcessorRegistry, TokenCounts @@ -59,6 +63,12 @@ def __init__(self, engine_args: SGLangLLMConfig, **kwargs) -> None: **kwargs: Additional arguments passed to base class. """ super().__init__(engine_args, **kwargs) + if not SGLANG_AVAILABLE: + raise RuntimeError( + "SGLang is not installed. Install with: pip install 'mmirage[gpu]' " + "or, from a source checkout, pip install -e '.[gpu]'" + ) + server_kwargs = asdict(engine_args.server_args) extra = server_kwargs.pop("extra_engine_args", {}) or {} server_kwargs.update(extra)