feat(storage): Week 4 — add MLflow experiment & model management (Week 4)

New storage/ component (storage -> shared only; api -> storage; inference
untouched). MLflowTracker does fire-and-forget inference/prompt logging on the
request path (async via to_thread, best-effort log-and-drop, bounded pending
tasks), plus experiment-tracking and model/adapter registry scaffolding. Wired
through api.main lifespan + generation service; ATLAS_MLFLOW_* config, off by
default. Uncommented + wired the compose MLflow service (self-contained sqlite
store; host port 5001 to avoid macOS AirPlay). Both Dockerfiles copy storage/
and gain an optional INCLUDE_MLFLOW arg. Hot-path logging uses MlflowClient with
explicit run ids (not the fluent global-active-run API) for concurrency safety.

Refs: Week 4

Tests: tests/storage/mlflow/ + tests/api/test_inference_logging.py

Author: Julina Maharjan

.env.example

@@ -38,3 +38,21 @@ ATLAS_TRUST_REMOTE_CODE=false
 ATLAS_ENABLE_BATCHING=true
 ATLAS_MAX_BATCH_SIZE=16
 ATLAS_BATCH_WAIT_MS=5
+
+# --- MLflow: experiment & model tracking (Week 4) ---
+# Off by default: mock/CI and hosts without `mlflow` installed run unchanged.
+# Enable to log inference (params/latency/tokens/output) + track runs/registry.
+# Needs `pip install -r requirements-mlflow.txt` and a reachable tracking store.
+ATLAS_MLFLOW_ENABLED=false
+# Tracking server or store URI. Empty + enabled => local file store "file:./mlruns".
+# With docker compose, point at the bundled service: http://mlflow:5000
+ATLAS_MLFLOW_TRACKING_URI=
+# ATLAS_MLFLOW_REGISTRY_URI=          # defaults to the tracking URI
+ATLAS_MLFLOW_EXPERIMENT=atlasai-inference
+# Per-request inference logging (fire-and-forget, never blocks the request path)
+ATLAS_MLFLOW_LOG_INFERENCE=true
+# Persist prompt + completion text (set false to log only params/metrics)
+ATLAS_MLFLOW_LOG_PROMPTS=true
+ATLAS_MLFLOW_MAX_TEXT_CHARS=2000
+# Backpressure: cap concurrent in-flight log tasks (excess dropped, not blocked)
+ATLAS_MLFLOW_MAX_PENDING_LOGS=100

api/main.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from contextlib import asynccontextmanager
 from typing import AsyncIterator
 
@@ -22,6 +23,7 @@
 from shared.errors import register_exception_handlers
 from shared.logging import configure_logging, get_logger
 from shared.rate_limit import RateLimiter
+from storage.mlflow import MLflowTracker
 
 logger = get_logger(__name__)
 
@@ -36,14 +38,20 @@ async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
     """Load the inference engine on startup; release it on shutdown."""
     settings: Settings = app.state.settings
     engine: InferenceEngine = app.state.engine
+    tracker: MLflowTracker = app.state.mlflow
     logger.info(
         "gateway.startup",
         extra={"env": settings.env, "version": settings.app_version},
     )
     await engine.startup()
+    # Configure MLflow off the event loop; no-op + never fails when disabled.
+    await asyncio.to_thread(tracker.start)
+    if tracker.enabled:
+        logger.info("mlflow.ready", extra=tracker.health())
     try:
         yield
     finally:
+        await tracker.shutdown()
         await engine.shutdown()
         logger.info("gateway.shutdown")
 
@@ -52,13 +60,17 @@ def create_app(
     settings: Settings | None = None,
     *,
     engine: InferenceEngine | None = None,
+    tracker: MLflowTracker | None = None,
 ) -> FastAPI:
     """Build and return a configured :class:`FastAPI` application.
 
     Args:
         settings: Optional settings override (primarily for tests).
         engine: Optional pre-built inference engine (primarily for tests, to
             inject a mock-backed engine). Defaults to ``InferenceEngine.build()``.
+        tracker: Optional MLflow tracker (primarily for tests, to inject a spy or
+            disabled tracker). Defaults to :meth:`MLflowTracker.from_settings`,
+            which is a no-op unless ``ATLAS_MLFLOW_ENABLED`` is set.
 
     Returns:
         A fully wired FastAPI application instance.
@@ -77,6 +89,7 @@ def create_app(
     )
     app.state.settings = settings
     app.state.engine = engine or InferenceEngine.build()
+    app.state.mlflow = tracker or MLflowTracker.from_settings()
 
     # Bind the resolved settings to the `get_settings` dependency so every
     # `Depends(get_settings)` call site (auth, health, ...) uses *this* app's

api/services/generation.py

@@ -1,20 +1,27 @@
 """Generation service — gateway-side adapter over the Inference Engine.
 
-Week 1 shipped a stub here. Week 2 makes this a thin async adapter that
-delegates to :class:`inference.engine.InferenceEngine`, preserving the
-``generate(request) -> GenerationResponse`` seam the routers depend on.
+Week 1 shipped a stub here. Week 2 made this a thin async adapter that delegates
+to :class:`inference.engine.InferenceEngine`. Week 4 adds fire-and-forget
+inference logging: after a response is produced (or a stream completes) the
+service hands an :class:`InferenceRecord` to the :class:`MLflowTracker`, which
+schedules the write on a background task. The tracker call never blocks or
+raises into the request path, so logging is invisible to latency and to
+clients — and a tracking outage cannot fail a generation.
 """
 
 from __future__ import annotations
 
+import time
 from collections.abc import AsyncIterator
 
 from fastapi import Request
 
 from inference.engine import InferenceEngine
 from inference.types import StreamChunk
+from shared.context import get_request_id
 from shared.logging import get_logger
 from shared.schemas.generation import GenerationRequest, GenerationResponse
+from storage.mlflow import InferenceRecord, MLflowTracker
 
 logger = get_logger(__name__)
 
@@ -24,23 +31,90 @@ class GenerationService:
 
     Args:
         engine: The process-wide inference engine (from ``app.state.engine``).
+        tracker: The process-wide MLflow tracker (from ``app.state.mlflow``).
+            Defaults to a disabled no-op tracker so the service stays usable
+            without tracking configured.
     """
 
-    def __init__(self, engine: InferenceEngine) -> None:
+    def __init__(
+        self, engine: InferenceEngine, tracker: MLflowTracker | None = None
+    ) -> None:
         self._engine = engine
+        self._tracker = tracker or MLflowTracker.from_settings()
 
     async def generate(self, request: GenerationRequest) -> GenerationResponse:
-        """Return a full generation for ``request``."""
-        return await self._engine.generate(request)
+        """Return a full generation for ``request`` and log it (best-effort)."""
+        start = time.perf_counter()
+        response = await self._engine.generate(request)
+        latency_ms = (time.perf_counter() - start) * 1000.0
+
+        # Fire-and-forget: returns immediately, never raises.
+        self._tracker.log_inference(
+            InferenceRecord(
+                request_id=get_request_id(),
+                model=request.model,
+                backend=self._engine.backend_name,
+                kind="generate",
+                max_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                adapter=request.adapter,
+                finish_reason=response.finish_reason,
+                prompt_tokens=response.usage.prompt_tokens,
+                completion_tokens=response.usage.completion_tokens,
+                total_tokens=response.usage.total_tokens,
+                latency_ms=latency_ms,
+                prompt=request.prompt,
+                output=response.text,
+            )
+        )
+        return response
 
     async def generate_stream(
         self, request: GenerationRequest
     ) -> AsyncIterator[StreamChunk]:
-        """Yield streaming chunks for ``request``."""
+        """Yield streaming chunks for ``request``; log once the stream ends.
+
+        Accumulates the streamed text and timing (including time-to-first-token)
+        and emits a single :class:`InferenceRecord` after the terminal chunk.
+        """
+        start = time.perf_counter()
+        ttft_ms: float | None = None
+        parts: list[str] = []
+        finish_reason: str | None = None
+
         async for chunk in self._engine.generate_stream(request):
+            if chunk.delta:
+                if ttft_ms is None:
+                    ttft_ms = (time.perf_counter() - start) * 1000.0
+                parts.append(chunk.delta)
+            if chunk.finish_reason is not None:
+                finish_reason = chunk.finish_reason
             yield chunk
 
+        latency_ms = (time.perf_counter() - start) * 1000.0
+        self._tracker.log_inference(
+            InferenceRecord(
+                request_id=get_request_id(),
+                model=request.model,
+                backend=self._engine.backend_name,
+                kind="stream",
+                max_tokens=request.max_tokens,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                adapter=request.adapter,
+                finish_reason=finish_reason,
+                latency_ms=latency_ms,
+                ttft_ms=ttft_ms,
+                prompt=request.prompt,
+                output="".join(parts),
+            )
+        )
+
 
 def get_generation_service(request: Request) -> GenerationService:
-    """FastAPI provider that wires the request's app engine into the service."""
-    return GenerationService(engine=request.app.state.engine)
+    """FastAPI provider wiring the app's engine + tracker into the service."""
+    return GenerationService(
+        engine=request.app.state.engine,
+        tracker=getattr(request.app.state, "mlflow", None),
+    )

deployment/docker/Dockerfile.cpu

@@ -28,20 +28,24 @@ ENV PYTHONDONTWRITEBYTECODE=1 \
 # When 1 (default), also install the hf backend stack (torch + transformers).
 # Set to 0 for a mock-only image with no heavy ML dependencies.
 ARG INCLUDE_HF=1
+# When 1, install the MLflow tracking client (Week 4). Needed only if you run
+# with ATLAS_MLFLOW_ENABLED=true; the tracker imports mlflow lazily otherwise.
+ARG INCLUDE_MLFLOW=1
 
 WORKDIR /build
 
 # Copy only dependency manifests first so this layer is cached across source
 # edits (the expensive pip install only re-runs when requirements change).
-COPY requirements.txt requirements-hf.txt ./
+COPY requirements.txt requirements-hf.txt requirements-mlflow.txt ./
 
 # Build an isolated venv that the runtime stage copies wholesale.
 RUN python -m venv /opt/venv
 ENV PATH="/opt/venv/bin:$PATH"
 
 RUN pip install --upgrade pip \
  && pip install -r requirements.txt \
- && if [ "$INCLUDE_HF" = "1" ]; then pip install -r requirements-hf.txt; fi
+ && if [ "$INCLUDE_HF" = "1" ]; then pip install -r requirements-hf.txt; fi \
+ && if [ "$INCLUDE_MLFLOW" = "1" ]; then pip install -r requirements-mlflow.txt; fi
 
 # --- Stage 2: runtime - slim, non-root ---------------------------------------
 FROM ${PYTHON_IMAGE} AS runtime
@@ -76,6 +80,7 @@ COPY --from=builder /opt/venv /opt/venv
 COPY --chown=atlas:atlas api/ ./api/
 COPY --chown=atlas:atlas shared/ ./shared/
 COPY --chown=atlas:atlas inference/ ./inference/
+COPY --chown=atlas:atlas storage/ ./storage/
 
 USER atlas
 EXPOSE 8000

deployment/docker/Dockerfile.gpu

@@ -51,18 +51,24 @@ ENV PATH="/opt/venv/bin:$PATH"
 
 WORKDIR /app
 
+# When 1, install the MLflow tracking client (Week 4). Needed only if you run
+# with ATLAS_MLFLOW_ENABLED=true; the tracker imports mlflow lazily otherwise.
+ARG INCLUDE_MLFLOW=1
+
 # Dependency manifests first for layer caching (vLLM install is expensive).
-COPY requirements.txt requirements-inference.txt ./
+COPY requirements.txt requirements-inference.txt requirements-mlflow.txt ./
 RUN pip install --upgrade pip \
  && pip install -r requirements.txt \
- && pip install -r requirements-inference.txt
+ && pip install -r requirements-inference.txt \
+ && if [ "$INCLUDE_MLFLOW" = "1" ]; then pip install -r requirements-mlflow.txt; fi
 
 # Run as an unprivileged user.
 RUN useradd --create-home --uid 10001 atlas
 
 COPY --chown=atlas:atlas api/ ./api/
 COPY --chown=atlas:atlas shared/ ./shared/
 COPY --chown=atlas:atlas inference/ ./inference/
+COPY --chown=atlas:atlas storage/ ./storage/
 
 USER atlas
 EXPOSE 8000

deployment/docker/docker-compose.yml

@@ -1,14 +1,20 @@
 # AtlasAI - local orchestration (Week 3, Docker Platform).
 #
-# Brings up the API Gateway + Inference Engine using the CPU image. Datastores
-# that arrive in later weeks (Postgres, Redis, MLflow) are scaffolded as
-# commented stubs below; uncomment each when its week lands.
+# Brings up the API Gateway + Inference Engine using the CPU image, plus the
+# MLflow tracking service (Week 4). Datastores that arrive in later weeks
+# (Postgres, Redis) remain scaffolded as commented stubs below; uncomment each
+# when its week lands.
 #
 # Usage (from this directory):
 #   docker compose up --build                              # mock backend
 #   ATLAS_INFERENCE_BACKEND=hf docker compose up --build   # real generation
+#   ATLAS_MLFLOW_ENABLED=true docker compose up --build    # + inference tracking
 #
 # The API is then on http://localhost:8000  (/docs, /tester, /health, /ready).
+# The MLflow UI is on http://localhost:5001 by default (host port). Port 5000 is
+# avoided because macOS AirPlay Receiver binds it; override with MLFLOW_HOST_PORT.
+# Note: this only remaps the *host* port — inside the compose network the api
+# still reaches MLflow at http://mlflow:5000 (the container port is unchanged).
 
 name: atlasai
 
@@ -19,6 +25,7 @@ services:
       dockerfile: deployment/docker/Dockerfile.cpu
       args:
         INCLUDE_HF: "1"
+        INCLUDE_MLFLOW: "1"
     image: atlasai:cpu
     ports:
       - "8000:8000"
@@ -29,6 +36,15 @@ services:
       ATLAS_API_KEYS: ${ATLAS_API_KEYS:-dev-key-local}
       ATLAS_LOG_LEVEL: ${ATLAS_LOG_LEVEL:-INFO}
       ATLAS_ENV: ${ATLAS_ENV:-development}
+      # MLflow tracking (Week 4). Off by default; flip ATLAS_MLFLOW_ENABLED=true
+      # to log inference to the bundled server below. TRACKING_URI targets the
+      # in-network `mlflow` service (override to use an external tracking store).
+      ATLAS_MLFLOW_ENABLED: ${ATLAS_MLFLOW_ENABLED:-false}
+      ATLAS_MLFLOW_TRACKING_URI: ${ATLAS_MLFLOW_TRACKING_URI:-http://mlflow:5000}
+      ATLAS_MLFLOW_EXPERIMENT: ${ATLAS_MLFLOW_EXPERIMENT:-atlasai-inference}
+    depends_on:
+      mlflow:
+        condition: service_healthy
     # To load a full local .env instead (copy from .env.example), uncomment:
     # env_file:
     #   - ../../.env
@@ -76,21 +92,33 @@ services:
   #     timeout: 5s
   #     retries: 5
 
-  # mlflow:              # Week 9 - experiment / model tracking
-  #   image: ghcr.io/mlflow/mlflow:v2.16.2
-  #   command: >
-  #     mlflow server --host 0.0.0.0 --port 5000
-  #     --backend-store-uri postgresql://atlas:atlas@postgres/atlas
-  #     --artifacts-destination /mlflow/artifacts
-  #   ports:
-  #     - "5000:5000"
-  #   depends_on:
-  #     - postgres
-  #   volumes:
-  #     - mlflow-artifacts:/mlflow/artifacts
+  mlflow:               # Week 4 - experiment / model tracking
+    image: ghcr.io/mlflow/mlflow:v2.16.2
+    # Self-contained for Week 4: SQLite backend store (required for the model
+    # registry) + a local artifacts dir, both on named volumes. When Postgres
+    # lands (Week 6) point --backend-store-uri at it and drop the mlflow-db volume.
+    command: >
+      mlflow server --host 0.0.0.0 --port 5000
+      --backend-store-uri sqlite:////mlflow/db/mlflow.db
+      --artifacts-destination /mlflow/artifacts
+    ports:
+      # Host 5001 -> container 5000 (5000 clashes with macOS AirPlay Receiver).
+      - "${MLFLOW_HOST_PORT:-5001}:5000"
+    volumes:
+      - mlflow-db:/mlflow/db
+      - mlflow-artifacts:/mlflow/artifacts
+    healthcheck:
+      test: ["CMD", "python", "-c",
+             "import urllib.request,sys; sys.exit(0) if urllib.request.urlopen('http://localhost:5000/health').status==200 else sys.exit(1)"]
+      interval: 30s
+      timeout: 5s
+      retries: 3
+      start_period: 20s
+    restart: unless-stopped
 
 volumes:
   hf-cache:
+  mlflow-db:
+  mlflow-artifacts:
   # pg-data:
   # redis-data:
-  # mlflow-artifacts:

inference/engine.py

@@ -192,6 +192,11 @@ def ready(self) -> bool:
         """Whether the engine has completed startup."""
         return self._ready
 
+    @property
+    def backend_name(self) -> str:
+        """Short identifier of the active generation backend (mock|hf|vllm)."""
+        return self._gen.name
+
     def health(self) -> dict[str, object]:
         """Return a JSON-serializable engine health snapshot."""
         return {