sientia-dataops-model-manager/CLAUDE.md

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

# 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 (boot fails if that changes the version of an already-imported package, so the requirements.txt pins must match the factory runtimes; app_up goes to 1 only after this check), 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 RetryPolicys 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).