Pin compatible versions: torch 2.6.0 + torchao <0.9 + transformers <5.0

Dockerfile

@@ -2,7 +2,7 @@
 # Serves both the UI dashboard AND GRPO training.
 # Set env OPENGRID_MODE=training for training mode.
 
-FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04
+FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04
 
 LABEL org.opencontainers.image.title="OpenGrid"
 LABEL org.opencontainers.image.description="Renewable energy grid load-balancing environment"
@@ -12,7 +12,7 @@ LABEL openenv="true"
 USER root
 RUN apt-get update && apt-get install -y --no-install-recommends \
     python3.10 python3-pip python3.10-venv python3-dev \
-    build-essential gcc g++ && \
+    build-essential gcc g++ git && \
     ln -sf /usr/bin/python3.10 /usr/bin/python && \
     ln -sf /usr/bin/pip3 /usr/bin/pip && \
     rm -rf /var/lib/apt/lists/*
@@ -27,39 +27,36 @@ ENV PATH="/home/user/.local/bin:$PATH"
 WORKDIR /app
 
 # --- Dependencies ---
-# Install server deps first (cached across builds)
+# 1. Install server deps (cached across builds)
 COPY --chown=user requirements.txt .
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
-# Install PyTorch (latest version to support torchao/torch.int1)
-RUN pip install --no-cache-dir torch
+# 2. Install PyTorch 2.6.0 with CUDA 12.1 (supports torch.int1 for torchao)
+RUN pip install --no-cache-dir \
+    torch==2.6.0 \
+    --extra-index-url https://download.pytorch.org/whl/cu121
 
-# Install training deps (only re-runs if training reqs change)
+# 3. Install training deps with pinned compatible versions
 COPY --chown=user requirements-training.txt .
 RUN pip install --no-cache-dir --upgrade -r requirements-training.txt
 
-# --- Application code (selective COPY for lean images) ---
-# Core Python modules
+# --- Application code ---
 COPY --chown=user src/ /app/src/
 COPY --chown=user training/ /app/training/
 
-# Entry points
 COPY --chown=user app.py /app/
 COPY --chown=user run_training.py /app/
 COPY --chown=user inference.py /app/
 COPY --chown=user entrypoint.sh /app/
 
-# Frontend (small, needed for server mode)
 COPY --chown=user static/ /app/static/
 
-# Config
 COPY --chown=user pyproject.toml /app/
 COPY --chown=user openenv.yaml /app/
 
 RUN chmod +x entrypoint.sh
 
-# Default mode controlled by OPENGRID_MODE env var (set via HF secrets)
-# server = FastAPI UI, training = GRPO pipeline
+# OPENGRID_MODE=training → GRPO pipeline, OPENGRID_MODE=server → UI (default)
 EXPOSE 7860
 
 HEALTHCHECK --interval=60s --timeout=10s --start-period=600s \

requirements-training.txt

@@ -1,5 +1,10 @@
-# Training dependencies (torch installed separately in Dockerfile with CUDA)
-transformers>=4.51.3
+# Training dependencies (torch 2.6.0 installed separately in Dockerfile)
+# -------------------------------------------------------------------
+# Pin versions that are mutually compatible with:
+#   torch==2.6.0, unsloth==2025.11.1, CUDA 12.1
+# -------------------------------------------------------------------
+
+transformers>=4.51.3,<5.0
 trl>=0.12.0,<1.0
 peft>=0.13.0
 accelerate>=1.0.0
@@ -7,6 +12,7 @@ datasets>=3.0.0
 bitsandbytes
 unsloth==2025.11.1
 unsloth_zoo
+torchao>=0.5.0,<0.9
 
 # Shared with environment
 fastapi