Fix Docker deployment

.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/

.github/workflows/docker.yml

@@ -4,9 +4,12 @@ on:
   push:
     branches:
       - main
+      # - fix/docker-deployment
+
   pull_request:
     branches:
       - main
+  workflow_dispatch: 
 
 env:
   IMAGE_NAME: fabnemepfl/mmirage
@@ -26,6 +29,14 @@ jobs:
             path: docker/Dockerfile
             tag_base: arm64
             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.cpu
+            tag_base: arm64
+            name: mmirage-git-cpu
 
     runs-on: ${{ matrix.platform }}
     environment: docker
@@ -62,6 +73,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 +87,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 }}

README.md

@@ -14,14 +14,43 @@ To install the library, you can clone it from GitHub and then use pip to install
 
 ```bash
 git clone git@github.com:EPFLiGHT/MMIRAGE.git
-pip install -e ./MMIRAGE
+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 (API-based LLMs)
+
+For API-based LLMs that do not require a GPU:
+
+```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

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

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", "-c", "import mmirage.core.process, runpy; runpy.run_module('mmirage.shard_process', run_name='__main__')"]

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", "-c", "import mmirage.core.process, runpy; runpy.run_module('mmirage.shard_process', run_name='__main__')"]

pyproject.toml

@@ -12,19 +12,15 @@ 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'",
   "fastapi",
   "openai>=1.0.0",
   "partial_json_parser",
   "sentencepiece",
-  "sgl_kernel",
-  "compressed_tensors",
   "msgspec",
   "nest_asyncio",
-  "xgrammar",
   "PyYAML",
   "json-repair",
   "tqdm",
@@ -38,9 +34,16 @@ dependencies = [
   "jmespath",
   "jinja2>=3.0.0",
   "pillow>=9.0.0",
+  "typing_extensions>=4.5.0; python_version < '3.12'",
 ]
 
 [project.optional-dependencies]
+gpu = [
+  "sglang>=0.5.2",
+  "sgl_kernel",
+  "xgrammar",
+  "compressed_tensors",
+]
 dev = [
   "ruff>=0.5.0",
   "black>=24.3.0",

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,

src/mmirage/core/process/processors/llm/llm_processor.py

@@ -8,7 +8,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
@@ -58,6 +62,8 @@ 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 '.[gpu]'")
         self.llm = sgl.Engine(**asdict(engine_args.server_args))
         self.tokenizer = AutoTokenizer.from_pretrained(
             engine_args.server_args.model_path,