# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## What this is

Enterprise ML model training orchestration platform built on **Temporal**. Orchestrates training workflows (validation → train → save to MLFlow → cleanup), backed by PostgreSQL (experiment tracking), MinIO (data/artifacts), MongoDB (notifications), and a Git-based "Plugin Store" that supplies model runtimes/wrappers dynamically. Python 3.11+.

## Commands

```bash
# Local run (loads .env, starts worker)
./run_local.sh
# or manually:
python -m model_manager.worker.worker

# Full validation (format check, lint, mypy, bandit, tests w/ coverage) — run before committing
./validate.sh
./validate.sh --fix          # auto-fix format/lint, no tests
./validate.sh --no-tests     # skip tests
./validate.sh --only-tests   # tests only

# Individual tools
ruff format --check model_manager/ tests/
ruff format model_manager/ tests/          # auto-fix
ruff check model_manager/ tests/
ruff check --fix model_manager/ tests/
mypy model_manager/
bandit -r model_manager/ -ll

# Tests
pytest                                          # runs tests/ and e2e/ (see pyproject testpaths)
pytest tests/                                   # unit tests only
pytest tests/activities/test_training.py -v     # single file
pytest tests/activities/test_training.py::test_name -v   # single test
pytest tests/ --cov=model_manager --cov-report=term-missing --cov-report=html

# Manual integration/E2E scripts (require live Temporal/Postgres/MinIO stack; run via IDE "Run Cell" or directly)
python scripts/run_training_test.py     # supports --all for batch scenario execution
python scripts/run_cleanup_test.py
```

Coverage gate in CI is 80% minimum; this repo targets 100% in practice — don't add uncovered code paths without tests.

## Architecture

### Temporal layering
- **Worker** (`model_manager/worker/worker.py`): entrypoint. Builds config from env, connects Postgres/MLFlow/MinIO/Plugin Store, installs the `RUNTIME` model runtime via `PluginStore.install_runtime`, starts two Temporal workers (one per task queue) plus Prometheus metrics server and the cleanup schedule. All heavy imports are wrapped in `workflow.unsafe.imports_passed_through()` — required because Temporal workflow sandboxing forbids non-deterministic imports at module scope outside that block.
- **Workflows** (`model_manager/workflows/`): pure orchestration, no I/O beyond `execute_activity_method` calls.
  - `train_model.py` (`TrainModel`, workflow name `train_model`): validate `experiment_run_id` → `load_model_metadata` → `validate_train_params` → update DB status `ORCHESTRATOR_WAITING_PROC` → `train_model` activity → update DB status `TRAINING_SUCCESS`/`TRAINING_ERROR` → `cleanup_resources` (always runs in `finally` if a `run_dir` exists, even on failure). Uses distinct `RetryPolicy`s per operation class: `network_retry_policy` (transient I/O), `no_retry_policy` (training/validation — data errors are permanent, and retrying would duplicate MLflow writes), `database_retry_policy` (Postgres). Timeouts are env-driven (`TIMEOUT_*`), extend `TIMEOUT_TRAIN_MODEL` instead of adding retries for slow training.
  - `cleanup_files.py` (`CleanupFiles`): scheduled (`cleanup_schedule.py`, default daily at midnight UTC via `CLEANUP_CRON`) deletion of stale local temp dirs matching `name_YYYYMMDD_HHMMSS_microseconds`, older than `CLEANUP_RETENTION_HOURS`.
  - Task queues are derived from `RUNTIME` env var (default `single`) via `prepare_worker.build_queue_name`: `train_model-<runtime>-queue`, `cleanup_files-<runtime>-queue`. A Temporal client starting a workflow must target the matching queue.
- **Activities** (`model_manager/activities/`): `Activities` (`activities.py`) combines `ExperimentTracking` + `Training` + `Cleanup` via multiple inheritance into one class registered with the worker — mind MRO and `__init__`/`__del__` ordering when touching this class (see the defensive `hasattr(self, 'engine')` check in `__del__`).
  - `experiment_tracking.py`: single `update_experiment_run()` entrypoint handling three update kinds via `UpdateType`: `STATUS`, `STATUS_WITH_ERROR`, `MODEL_SAVED`. Error messages truncated to 1024 chars before persisting.
  - `training.py`: single `train_model()` entrypoint. Receives pre-downloaded data as `BytesIO` (not re-downloaded per-activity) to avoid memory leaks; never raises on training failure — returns a result/error so the workflow can update DB status without an unhandled exception mid-activity.
  - `cleanup.py`: local-filesystem-only cleanup (`cleanup_temp_directories`), no MinIO involved — MinIO artifacts are managed externally.

### Data/config layer
- `model_manager/utils/connectors_config.py`: single source of truth for env-var → config-dict builders (`build_postgres_config`, `build_mlflow_config`, `build_minio_config`, `build_mongodb_config`, `build_plugin_store_config`). Read this file (and `worker.py`) before trusting any env var default — the README table mirrors it but this file is authoritative.
- `model_manager/utils/models/train_model_params.py`: `TrainModelParams` — the workflow's single input contract. `from_dict` does type coercion; `validate_business_rules()` runs after, enforcing 7 rules (train_size range, non-empty variable_columns, required model_metadata, JSON-Schema validation of `data_model_kwargs`/`model_kwargs`/`opt_params` against schemas in `model_metadata` via `Draft202012Validator`, required non-blank strings, allowed `date_format` values, numeric `experiment_run_id`). Extending training params means updating this file's validation, not just the workflow.
- `model_manager/utils/repository/data_manager_repository.py`: data loading, feature prep, metrics, report generation — the actual ML/data logic activities delegate into.
- `model_manager/sientia/`: RCE drift metrics (`metrics.py`), Evidently-based HTML report generation (`reports.py`), custom exceptions.
- `model_manager/metrics.py`: Prometheus metric definitions (`app_up`, `workflow_execution_total`, `activity_execution_total`, training/cleanup/RCE-drift gauges). New metrics go here, not inline in activities.

### External dependencies
Two private Aignosi packages carry non-trivial logic: `sientia_do` (observability/logger, notification handler, metrics controller, Postgres/MinIO repositories) and `sientia_model` (Plugin Store, MLflow repository, model wrappers). Both are pinned in `requirements.txt`/`requirements-local.txt`; when behavior looks like it lives outside this repo, it's likely in one of these.

### Test bootstrap quirk
`tests/conftest.py` stubs `evidently` and parts of `sientia_do` at `pytest_configure` time so unit tests can import `model_manager.sientia.reports` without the full (heavy/optional) Evidently install. If you add new imports from these packages in production code paths that unit tests exercise, you likely need to extend the stubs here too.

### Test layout
- `tests/` — unit tests, mirrors `model_manager/` package structure (`tests/activities/`, `tests/schedules/`, `tests/sientia/`, etc.).
- `e2e/` — integration tests against a real Temporal/Postgres/MinIO/MongoDB stack (via `testcontainers`), driven by JSON scenario files in `docs/test-scenarios/` (see `e2e/scenarios.md` for the scenario→test mapping). Scenario payloads are snake_case, shaped like `TrainModelParams`/`input-sample.json`.
- New scenario: copy an existing `docs/test-scenarios/*.json`, add/extend a test in `e2e/test_train_model_workflow.py`, update `e2e/scenarios.md`.

## Conventions
- Ruff: line-length 100, single quotes, `select = [E,W,F,I,B,C4,UP,N,YTT,S,BLE,A,C90]`, max complexity 15. `N802`/`N806` ignored (Temporal decorators use non-lowercase names). `S101`/`S105`/`S106` ignored in tests.
- mypy: `disallow_untyped_defs = false` but `check_untyped_defs = true` — existing untyped defs are tolerated, but code inside them is still checked. Several external packages (`temporalio`, `mlflow`, `sientia_model`, etc.) have `ignore_missing_imports`; `model_repository.py`/`data_manager_repository.py` have `ignore_errors = true` (legacy — don't extend that list without reason).