docker size reduced

.dockerignore

@@ -3,7 +3,8 @@ __pycache__/
 .git/
 .pytest_cache/
 tests/
-validation/
+validation/*.py
+validation/requirements.txt
 *.md
 !README.md
 .claude/
@@ -11,3 +12,4 @@ run*.json
 htmlcov/
 .mypy_cache/
 .ruff_cache/
+.env

Dockerfile

@@ -1,22 +1,29 @@
-FROM python:3.12-slim
-
-WORKDIR /app
+# ---- Stage 1: Builder — install + strip aggressively ----
+FROM python:3.12-slim AS builder
 
-# Install system deps (curl for healthcheck)
-RUN apt-get update && apt-get install -y --no-install-recommends curl && \
+RUN apt-get update && apt-get install -y --no-install-recommends binutils && \
     rm -rf /var/lib/apt/lists/*
 
-# Install ALL Python deps + safe cleanup in ONE layer.
-# Docker layers are immutable — cleanup in a separate RUN saves nothing.
-# PyTorch CPU-only (~280MB wheel, ~460MB installed) is the minimum for real
-# torch.nn.Module, torch.autograd, and state_dict() support.
+WORKDIR /build
 COPY requirements.txt .
-RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu && \
-    pip install --no-cache-dir -r requirements.txt && \
-    # Remove non-essential torch components (safe — verified these don't break imports)
+
+RUN pip install --no-cache-dir --no-compile \
+        torch==2.5.1+cpu --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir --no-compile -r requirements.txt && \
+    #
+    # === STRIP DEBUG SYMBOLS FROM ALL .so FILES (saves ~100-200MB) ===
+    find /usr/local/lib/python3.12/site-packages -name "*.so" -exec strip --strip-unneeded {} + 2>/dev/null; \
+    find /usr/local/lib/python3.12/site-packages -name "*.so.*" -exec strip --strip-unneeded {} + 2>/dev/null; \
+    #
+    # === TORCH CLEANUP ===
     rm -rf /usr/local/lib/python3.12/site-packages/torch/test \
            /usr/local/lib/python3.12/site-packages/torch/include \
            /usr/local/lib/python3.12/site-packages/torch/share \
+           /usr/local/lib/python3.12/site-packages/torch/bin/FileStore* \
+           /usr/local/lib/python3.12/site-packages/torch/bin/HashStore* \
+           /usr/local/lib/python3.12/site-packages/torch/bin/TCPStore* \
+           /usr/local/lib/python3.12/site-packages/torch/bin/protoc* \
+           /usr/local/lib/python3.12/site-packages/torch/bin/test_* \
            /usr/local/lib/python3.12/site-packages/torch/utils/benchmark \
            /usr/local/lib/python3.12/site-packages/torch/utils/bottleneck \
            /usr/local/lib/python3.12/site-packages/torch/utils/tensorboard \
@@ -24,18 +31,74 @@ RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/wh
            /usr/local/lib/python3.12/site-packages/torch/lib/libtorchbind_test.so \
            /usr/local/lib/python3.12/site-packages/torch/lib/libjitbackend_test.so \
            /usr/local/lib/python3.12/site-packages/torch/lib/libbackend_with_compiler.so \
-           /usr/local/lib/python3.12/site-packages/caffe2 2>/dev/null; \
-    find /usr/local/lib/python3.12/site-packages/torch -name "*.pyi" -delete 2>/dev/null; \
+           /usr/local/lib/python3.12/site-packages/torch/lib/libaoti_custom_ops.so \
+           /usr/local/lib/python3.12/site-packages/torch/lib/libshm_windows \
+           /usr/local/lib/python3.12/site-packages/caffe2 \
+    #
+    # === BLOATED TRANSITIVE DEPS ===
+           /usr/local/lib/python3.12/site-packages/gradio \
+           /usr/local/lib/python3.12/site-packages/gradio_client \
+           /usr/local/lib/python3.12/site-packages/hf_gradio \
+           /usr/local/lib/python3.12/site-packages/pandas \
+           /usr/local/lib/python3.12/site-packages/PIL \
+           /usr/local/lib/python3.12/site-packages/Pillow* \
+           /usr/local/lib/python3.12/site-packages/pillow* \
+           /usr/local/lib/python3.12/site-packages/networkx \
+           /usr/local/lib/python3.12/site-packages/scipy \
+           /usr/local/lib/python3.12/site-packages/matplotlib \
+           /usr/local/lib/python3.12/site-packages/hf_xet \
+           /usr/local/lib/python3.12/site-packages/ffmpy \
+           /usr/local/lib/python3.12/site-packages/pydub \
+           /usr/local/lib/python3.12/site-packages/groovy \
+           /usr/local/lib/python3.12/site-packages/tomlkit \
+           /usr/local/lib/python3.12/site-packages/semantic_version* \
+           /usr/local/lib/python3.12/site-packages/safehttpx* \
+           /usr/local/lib/python3.12/site-packages/brotli* \
+           /usr/local/lib/python3.12/site-packages/Brotli* \
+           /usr/local/lib/python3.12/site-packages/pip \
+           /usr/local/lib/python3.12/site-packages/setuptools \
+           /usr/local/lib/python3.12/site-packages/docutils \
+           /usr/local/lib/python3.12/site-packages/cryptography \
+           /usr/local/lib/python3.12/site-packages/cryptography* \
+           /usr/local/lib/python3.12/site-packages/pytz 2>/dev/null; \
+    #
+    # === FILE-LEVEL CLEANUP ===
+    find /usr/local/lib/python3.12/site-packages -name "*.pyi" -delete 2>/dev/null; \
+    find /usr/local/lib/python3.12/site-packages -name "*.pyc" -delete 2>/dev/null; \
     find /usr/local/lib/python3.12/site-packages -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null; \
+    find /usr/local/lib/python3.12/site-packages -name "*.egg-info" -type d -exec rm -rf {} + 2>/dev/null; \
+    find /usr/local/lib/python3.12/site-packages -name "tests" -type d -exec rm -rf {} + 2>/dev/null; \
+    find /usr/local/lib/python3.12/site-packages -name "test" -type d -exec rm -rf {} + 2>/dev/null; \
+    # Remove stale dist-info for packages we already deleted
+    rm -rf /usr/local/lib/python3.12/site-packages/gradio*.dist-info \
+           /usr/local/lib/python3.12/site-packages/pandas*.dist-info \
+           /usr/local/lib/python3.12/site-packages/Pillow*.dist-info \
+           /usr/local/lib/python3.12/site-packages/hf_xet*.dist-info \
+           /usr/local/lib/python3.12/site-packages/Brotli*.dist-info \
+           /usr/local/lib/python3.12/site-packages/networkx*.dist-info \
+           /usr/local/lib/python3.12/site-packages/pip \
+           /usr/local/lib/python3.12/site-packages/pip*.dist-info 2>/dev/null; \
     true
 
-# Copy application code
+# ---- Stage 2: Runtime — minimal clean image ----
+FROM python:3.12-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y --no-install-recommends curl && \
+    rm -rf /var/lib/apt/lists/*
+
+# Copy only what's needed from builder
+COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
+COPY --from=builder /usr/local/bin/uvicorn /usr/local/bin/uvicorn
+
 COPY ml_training_debugger/ ml_training_debugger/
 COPY server/ server/
 COPY openenv.yaml .
 COPY baseline_heuristic.py .
 COPY baseline_inference.py .
 COPY README.md .
+COPY validation/reports/ validation/reports/
 
 EXPOSE 7860
 

tests/test_graders.py

@@ -5,10 +5,12 @@ from __future__ import annotations
 import pytest
 
 from ml_training_debugger.graders import (
+    _submitted_diagnosis,
     grade_episode,
     grade_task_001,
     grade_task_003,
     grade_task_005,
+    grade_task_007,
 )
 from ml_training_debugger.models import EpisodeState
 from ml_training_debugger.scenarios import sample_scenario
@@ -166,3 +168,60 @@ class TestGradeEpisode:
         state = EpisodeState()
         score = grade_episode("task_999", state, scenario_001)
         assert score == 0.0
+
+
+class TestGradeTask007:
+    def test_perfect_score(self):
+        scenario = sample_scenario("task_007", seed=42)
+        state = EpisodeState(
+            gradients_inspected=True,
+            data_inspected=True,
+            fix_action_taken=True,
+            restart_after_fix=True,
+            diagnosis_submitted=True,
+            actions_taken=[
+                "inspect_gradients",
+                "inspect_data_batch",
+                "modify_config",
+                "restart_run",
+                "mark_diagnosed:scheduler_misconfigured",
+            ],
+        )
+        score = grade_task_007(state, scenario)
+        assert score == 1.0
+
+    def test_wrong_diagnosis(self):
+        scenario = sample_scenario("task_007", seed=42)
+        state = EpisodeState(
+            diagnosis_submitted=True,
+            actions_taken=["mark_diagnosed:overfitting"],
+        )
+        score = grade_task_007(state, scenario)
+        assert score < 0.5
+
+    def test_score_in_range(self):
+        scenario = sample_scenario("task_007", seed=42)
+        state = EpisodeState()
+        score = grade_task_007(state, scenario)
+        assert 0.0 <= score <= 1.0
+
+
+class TestSubmittedDiagnosis:
+    def test_finds_diagnosis(self):
+        state = EpisodeState(
+            actions_taken=["inspect_gradients", "mark_diagnosed:lr_too_high"],
+        )
+        assert _submitted_diagnosis(state) == "lr_too_high"
+
+    def test_no_diagnosis(self):
+        state = EpisodeState(actions_taken=["inspect_gradients"])
+        assert _submitted_diagnosis(state) is None
+
+    def test_latest_diagnosis(self):
+        state = EpisodeState(
+            actions_taken=[
+                "mark_diagnosed:overfitting",
+                "mark_diagnosed:lr_too_high",
+            ],
+        )
+        assert _submitted_diagnosis(state) == "lr_too_high"

tests/test_new_endpoints.py

@@ -0,0 +1,80 @@
+"""Tests for new endpoints: curriculum, leaderboard, replay, validation-report."""
+
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+
+from server.app import app
+
+
+@pytest.fixture
+def client():
+    return TestClient(app)
+
+
+class TestCurriculumEndpoint:
+    def test_returns_curriculum(self, client) -> None:
+        resp = client.get("/curriculum")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert "curriculum" in data
+        assert "total_episodes" in data
+        assert data["total_episodes"] > 0
+
+    def test_curriculum_has_difficulty_levels(self, client) -> None:
+        resp = client.get("/curriculum")
+        curriculum = resp.json()["curriculum"]
+        levels = {entry["difficulty_level"] for entry in curriculum}
+        assert 1 in levels
+        assert 3 in levels
+        assert 5 in levels
+
+    def test_curriculum_covers_all_tasks(self, client) -> None:
+        resp = client.get("/curriculum")
+        curriculum = resp.json()["curriculum"]
+        task_ids = {entry["task_id"] for entry in curriculum}
+        assert "task_001" in task_ids
+        assert "task_007" in task_ids
+
+
+class TestLeaderboardEndpoint:
+    def test_returns_leaderboard(self, client) -> None:
+        resp = client.get("/leaderboard")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert "entries" in data
+        assert "total" in data
+
+    def test_leaderboard_after_baseline(self, client) -> None:
+        # Run baseline to populate scores
+        client.post("/baseline")
+        resp = client.get("/leaderboard")
+        data = resp.json()
+        assert data["total"] > 0
+
+
+class TestReplayEndpoint:
+    def test_missing_episode(self, client) -> None:
+        resp = client.get("/replay/nonexistent_episode_123")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert "error" in data
+
+    def test_replay_after_baseline(self, client) -> None:
+        # Run baseline to create episodes
+        client.post("/baseline")
+        resp = client.get("/replay/baseline_task_001")
+        data = resp.json()
+        # Should have episode data or error
+        assert "episode_id" in data or "error" in data
+
+
+class TestValidationReportEndpoint:
+    def test_returns_real_report(self, client) -> None:
+        resp = client.get("/validation-report")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert "results" in data
+        assert "summary" in data
+        assert data["summary"]["passed"] > 0

tests/test_real_training.py

@@ -0,0 +1,135 @@
+"""Tests for real mini-training in pytorch_engine.py."""
+
+from __future__ import annotations
+
+import torch
+
+from ml_training_debugger.pytorch_engine import (
+    SimpleCNN,
+    SimpleMLP,
+    _TRAINING_CACHE,
+    run_real_training,
+)
+from ml_training_debugger.scenarios import sample_scenario
+
+
+class TestRunRealTraining:
+    def test_returns_20_epoch_curves(self) -> None:
+        s = sample_scenario("task_001", seed=42)
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+        assert len(curves["val_loss_history"]) == 20
+        assert len(curves["val_acc_history"]) == 20
+
+    def test_all_values_are_floats(self) -> None:
+        s = sample_scenario("task_003", seed=42)
+        curves = run_real_training(s)
+        for key in ["loss_history", "val_loss_history", "val_acc_history"]:
+            for v in curves[key]:
+                assert isinstance(v, float), f"{key} has non-float: {type(v)}"
+
+    def test_caching_works(self) -> None:
+        _TRAINING_CACHE.clear()
+        s = sample_scenario("task_001", seed=42)
+        c1 = run_real_training(s)
+        c2 = run_real_training(s)
+        assert c1 is c2  # Same object reference = cached
+
+    def test_reproducible_across_calls(self) -> None:
+        _TRAINING_CACHE.clear()
+        s = sample_scenario("task_002", seed=42)
+        c1 = run_real_training(s)
+        _TRAINING_CACHE.clear()
+        c2 = run_real_training(s)
+        assert c1["loss_history"] == c2["loss_history"]
+        assert c1["val_acc_history"] == c2["val_acc_history"]
+
+    def test_different_seeds_different_curves(self) -> None:
+        s1 = sample_scenario("task_001", seed=42)
+        s2 = sample_scenario("task_001", seed=99)
+        c1 = run_real_training(s1)
+        c2 = run_real_training(s2)
+        assert c1["loss_history"] != c2["loss_history"]
+
+    def test_task_001_high_lr_instability(self) -> None:
+        s = sample_scenario("task_001", seed=42)
+        curves = run_real_training(s)
+        max_loss = max(v for v in curves["loss_history"] if v != float("inf"))
+        assert max_loss > 3.0  # High LR causes loss spikes
+
+    def test_task_002_vanishing_slow_learning(self) -> None:
+        s = sample_scenario("task_002", seed=42)
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+
+    def test_task_003_data_leakage(self) -> None:
+        s = sample_scenario("task_003", seed=42)
+        curves = run_real_training(s)
+        # With leakage, val accuracy may be elevated
+        assert len(curves["val_acc_history"]) == 20
+
+    def test_task_004_overfitting(self) -> None:
+        s = sample_scenario("task_004", seed=42)
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+
+    def test_task_005_batchnorm_eval(self) -> None:
+        s = sample_scenario("task_005", seed=42)
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+
+    def test_task_006_code_bug(self) -> None:
+        s = sample_scenario("task_006", seed=42)
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+
+    def test_task_007_scheduler(self) -> None:
+        s = sample_scenario("task_007", seed=42)
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+
+    def test_mlp_architecture(self) -> None:
+        """Find a scenario that uses MLP and verify training works."""
+        for seed in range(1, 20):
+            s = sample_scenario("task_001", seed=seed)
+            if s.model_type == "mlp":
+                curves = run_real_training(s)
+                assert len(curves["loss_history"]) == 20
+                return
+        # If no MLP found in 20 seeds, test directly
+        from ml_training_debugger.scenarios import ScenarioParams
+        from ml_training_debugger.models import RootCauseDiagnosis
+        s = ScenarioParams(
+            task_id="task_001",
+            root_cause=RootCauseDiagnosis.LR_TOO_HIGH,
+            seed=999,
+            learning_rate=0.1,
+            model_type="mlp",
+        )
+        curves = run_real_training(s)
+        assert len(curves["loss_history"]) == 20
+
+
+class TestSimpleMLP:
+    def test_is_nn_module(self) -> None:
+        model = SimpleMLP()
+        assert isinstance(model, torch.nn.Module)
+
+    def test_param_count(self) -> None:
+        model = SimpleMLP()
+        count = sum(p.numel() for p in model.parameters())
+        assert 10_000 < count < 500_000
+
+    def test_forward_pass(self) -> None:
+        model = SimpleMLP()
+        x = torch.randn(4, 3, 32, 32)
+        out = model(x)
+        assert out.shape == (4, 10)
+
+    def test_has_batchnorm(self) -> None:
+        model = SimpleMLP()
+        has_bn = any(
+            isinstance(m, torch.nn.BatchNorm1d)
+            for m in model.modules()
+        )
+        assert has_bn

tests/test_simulation_fallback.py

@@ -0,0 +1,149 @@
+"""Tests for parametric fallback in simulation.py.
+
+These test the fallback paths that run when real training is unavailable.
+We force fallback by monkeypatching _get_real_curves to return None.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import patch
+
+from ml_training_debugger.scenarios import sample_scenario
+from ml_training_debugger.simulation import (
+    gen_loss_history,
+    gen_val_accuracy_history,
+    gen_val_loss_history,
+)
+
+
+def _force_fallback(*args, **kwargs):
+    return None
+
+
+class TestParametricFallbackLoss:
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_001_fallback(self) -> None:
+        s = sample_scenario("task_001", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_002_fallback(self) -> None:
+        s = sample_scenario("task_002", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_003_fallback(self) -> None:
+        s = sample_scenario("task_003", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_004_fallback(self) -> None:
+        s = sample_scenario("task_004", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_005_fallback(self) -> None:
+        s = sample_scenario("task_005", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_006_fallback(self) -> None:
+        s = sample_scenario("task_006", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_007_fallback(self) -> None:
+        s = sample_scenario("task_007", seed=42)
+        hist = gen_loss_history(s)
+        assert len(hist) == 20
+
+
+class TestParametricFallbackValAcc:
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_001_fallback(self) -> None:
+        s = sample_scenario("task_001", seed=42)
+        hist = gen_val_accuracy_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_003_fallback(self) -> None:
+        s = sample_scenario("task_003", seed=42)
+        hist = gen_val_accuracy_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_004_fallback(self) -> None:
+        s = sample_scenario("task_004", seed=42)
+        hist = gen_val_accuracy_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_005_fallback(self) -> None:
+        s = sample_scenario("task_005", seed=42)
+        hist = gen_val_accuracy_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_006_fallback(self) -> None:
+        s = sample_scenario("task_006", seed=42)
+        hist = gen_val_accuracy_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_007_fallback(self) -> None:
+        s = sample_scenario("task_007", seed=42)
+        hist = gen_val_accuracy_history(s)
+        assert len(hist) == 20
+
+
+class TestParametricFallbackValLoss:
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_001_fallback(self) -> None:
+        s = sample_scenario("task_001", seed=42)
+        hist = gen_val_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_004_fallback(self) -> None:
+        s = sample_scenario("task_004", seed=42)
+        hist = gen_val_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_005_fallback(self) -> None:
+        s = sample_scenario("task_005", seed=42)
+        hist = gen_val_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_006_fallback(self) -> None:
+        s = sample_scenario("task_006", seed=42)
+        hist = gen_val_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_task_007_fallback(self) -> None:
+        s = sample_scenario("task_007", seed=42)
+        hist = gen_val_loss_history(s)
+        assert len(hist) == 20
+
+    @patch("ml_training_debugger.simulation._get_real_curves", _force_fallback)
+    def test_fallback_default(self) -> None:
+        """Test the final fallback path for unknown root cause."""
+        from ml_training_debugger.models import RootCauseDiagnosis
+        from ml_training_debugger.scenarios import ScenarioParams
+
+        # Use scheduler root cause but force fallback
+        s = ScenarioParams(
+            task_id="task_999",
+            root_cause=RootCauseDiagnosis.SCHEDULER_MISCONFIGURED,
+            seed=42,
+        )
+        hist = gen_val_loss_history(s)
+        assert len(hist) == 20