kailash-ml 2.2.0

Machine learning lifecycle for the Kailash ecosystem

kailash-ml

Machine learning lifecycle for the Kailash ecosystem. Train models, store features, serve predictions, monitor drift, and optionally augment decisions with AI agents -- all built on polars DataFrames and backed by a single database.

Part of the Kailash Python SDK by the Terrene Foundation.

Version: 2.2.0 | License: Apache-2.0 | Python: 3.11+

---

Why kailash-ml

Polars-native from the ground up. Every engine accepts and returns polars.DataFrame. Conversion to numpy, pandas, or framework-specific formats happens at the boundary, never inside your pipeline logic. This means you work with a single, fast, memory-efficient data format throughout.

Zero-config persistence. All engines store their state -- features, model metadata, drift reports, experiment logs -- through the same database your application already uses. No separate tracking server, no object store to configure, no infrastructure to maintain.

Agent-augmented, not agent-dependent. Six Kaizen agents can augment your ML workflow with LLM-powered recommendations. They require double opt-in (install the extras AND pass agent=True), enforce cost budgets, and always run alongside a pure algorithmic baseline. You never need agents -- they are a power-user option.

Production lifecycle built in. Models move through a governed lifecycle (staging, shadow, production, archived). Drift monitoring runs on a schedule. ONNX export happens automatically. The inference server caches models in memory and exposes predictions via Nexus.

Security by default. Model class imports are restricted to an allowlist. All SQL is encapsulated in a single auditable file. Financial fields are validated against NaN and Infinity. Agent cost budgets are hard-capped.

---

Installation

# Base install (~195MB): polars, sklearn, LightGBM, XGBoost, scipy, plotly, ONNX
#   XGBoost>=2.0 ships with CUDA built in and auto-detects GPU at runtime
#   (CPU fallback otherwise) — no separate GPU wheel is required.
pip install kailash-ml

# Optional extras
pip install kailash-ml[dl]        # + PyTorch, Lightning, transformers
pip install kailash-ml[dl-gpu]    # + onnxruntime-gpu (CUDA inference)
pip install kailash-ml[rl]        # + Stable-Baselines3, Gymnasium
pip install kailash-ml[agents]    # + Kaizen (LLM-augmented ML)
# [xgb] is a no-op alias retained for backward-compat; XGBoost is now a base dep.
pip install kailash-ml[catboost]  # + CatBoost
pip install kailash-ml[stats]     # + statsmodels
pip install kailash-ml[hpo]       # + Optuna (Bayesian HPO, successive halving)
pip install kailash-ml[imbalance] # + imbalanced-learn (SMOTE, ADASYN)
pip install kailash-ml[explain]   # + SHAP (model explainability)
pip install kailash-ml[all]       # Everything above
pip install kailash-ml[all-gpu]   # Everything + GPU runtime

Development

pip install kailash-ml[dev]       # pytest, hypothesis, mypy, ruff

---

Quick Start

df below is any polars.DataFrame whose columns include a target you want to predict (supply the column name via target="...").

Install with pip install kailash-ml. The following block is executed by CI against a real DataFrame and is the canonical Quick Start for the 1.0.0 release. Every line is load-bearing — do NOT abbreviate it in copy-paste examples on external sites.

import kailash_ml as km
async with km.track("demo") as run:
    result = await km.train(df, target="y")
    registered = await km.register(result, name="demo")
server = await km.serve("demo@production")
# $ kailash-ml-dashboard  (separate shell)

What this does

Each line maps to a single deliberate effect:

• import kailash_ml as km — imports the package. Every public entry point is a km.* function; there are no constructors to wire in the Quick Start (see specs/ml-engines-v2.md §2.1 MUST 1, zero-arg construction).
• async with km.track("demo") as run: — opens an ambient ExperimentRun context named "demo". Every km.train(...) / km.register(...) / checkpoint / metric emitted inside the block is auto-attached to this run (specs/ml-tracking.md §2.4). The run finalises (status, end time, artifact manifest) on context exit.
• result = await km.train(df, target="y") — picks the best model family for the target's dtype and row count, trains it on the chosen backend (CPU / CUDA / MPS / ROCm / XPU / TPU auto-detected per specs/ml-backends.md), and returns a TrainingResult with populated metrics and device: DeviceReport (specs/ml-engines-v2.md §4.2 MUST 1).
• registered = await km.register(result, name="demo") — creates a new ModelVersion under the logical name "demo" (auto-bumped version number), serialises the trained artefact as ONNX by default (specs/ml-engines-v2.md §2.1 MUST 9), and returns a RegisterResult with artifact_uris pointing at the framework-agnostic ONNX + native format pair.
• server = await km.serve("demo@production") — resolves the "demo" model's production alias to its current version, spins up a local in-process inference server exposing REST (default) and MCP channels, and returns a ServeResult with .uris["rest"], .channels, .model_uri, and .model_version (specs/ml-serving.md §2.2). (The low-level InferenceServer.start() returns a separate ServeHandle carrying .stop() / .status.)
• # $ kailash-ml-dashboard (separate shell) — invites the reader to launch the optional ML dashboard CLI from a second shell. The dashboard is a non-blocking visualisation surface for runs, models, and serving handles (specs/ml-dashboard.md). The comment is retained verbatim because the Quick Start's 5-to-10-line budget (specs/ml-engines-v2.md §16.2 MUST 1) accommodates exactly this one-line operational pointer, and newbies otherwise never discover the dashboard exists.

For the full engine surface, multi-tenancy, checkpoint / resume, RL, AutoML, drift monitoring, and production deployment patterns, see the full spec set under specs/ml-*.md. The canonical Quick Start is intentionally narrow — it demonstrates the zero-ceremony promise of the 1.0.0 Engine and nothing beyond that.

---

Architecture Overview

                              kailash-ml
    +---------------------------------------------------------+
    |                                                         |
    |   types.py (FeatureSchema, ModelSignature, Protocols)   |
    |                          |                              |
    |   +-------- engines/ --------+                          |
    |   |                          |                          |
    |   |  FeatureStore            |  interop.py              |
    |   |  TrainingPipeline  <---->|  (polars <-> sklearn     |
    |   |  ModelRegistry           |   polars <-> lightgbm    |
    |   |  InferenceServer         |   polars <-> pandas      |
    |   |  DriftMonitor            |   polars <-> arrow       |
    |   |  ExperimentTracker       |   polars <-> HuggingFace)|
    |   |  HyperparameterSearch    |                          |
    |   |  AutoMLEngine            |                          |
    |   |  EnsembleEngine          |                          |
    |   |  PreprocessingPipeline   |                          |
    |   |  DataExplorer            |                          |
    |   |  FeatureEngineer         |                          |
    |   |  ModelVisualizer         |                          |
    |   |  ModelExplainer          |                          |
    |   +------------|-------------+                          |
    |                |                                        |
    |   bridge/      |     compat/       dashboard/           |
    |   OnnxBridge   |     MlflowFormat  MLDashboard          |
    |                |                                        |
    |   agents/ (optional, double opt-in)                     |
    |   DataScientist, FeatureEngineer, ModelSelector,        |
    |   ExperimentInterpreter, DriftAnalyst, RetrainingDecision|
    |                |                                        |
    |   rl/ (optional)                                        |
    |   RLTrainer, EnvironmentRegistry, PolicyRegistry        |
    |                |                                        |
    +----------------|----------------------------------------+
                     |
           ConnectionManager (kailash.db)
                     |
               SQLite / PostgreSQL

All engines persist through ConnectionManager from the core Kailash SDK. The interop.py module is the sole conversion point between polars and external frameworks -- no engine file contains direct pandas or numpy conversion logic.

---

Engines Reference

kailash-ml provides 14 engines plus a bridge and compatibility layer, organized by purpose and stability.

Core Engines (P0 -- stable API, full test coverage)

FeatureStore

# Canonical 1.0+ read surface (the top-level default since 2.0.0)
from kailash_ml import FeatureStore  # → kailash_ml.features.FeatureStore

The canonical FeatureStore is a thin, polars-native DataFlow bridge: it wraps a live DataFlow instance and serves point-in-time-correct feature reads through dataflow.ml_feature_source. Constructor: FeatureStore(dataflow, *, default_tenant_id=None). It owns no DDL — you materialise features as ordinary @db.model rows (db.express.create(...)) and read them back through the store.

Key operation: get_features(schema, timestamp=None, *, tenant_id=None, entity_ids=None) — returns a polars.DataFrame (latest value per entity, or the as-of-timestamp value). See specs/ml-feature-store.md § 4.1 and MIGRATION.md for the write-via-model recipe.

# Legacy 0.x write-capable surface (explicit import; ConnectionManager-based)
from kailash_ml.engines.feature_store import FeatureStore

The legacy engine computes, versions, and stores features itself via ConnectionManager; all raw SQL is encapsulated in _feature_sql.py. Constructor: FeatureStore(conn, *, table_prefix="kml_feat_"). Key operations: initialize(), register_features(), compute(), store(), get_features(), get_training_set(), get_features_lazy(), list_schemas(). Retained for callers needing the self-contained write path; see MIGRATION.md for the canonical migration.

ModelRegistry

from kailash_ml.engines.model_registry import ModelRegistry, LocalFileArtifactStore

Manages the complete model lifecycle through four stages: staging (newly trained), shadow (running alongside production), production (serving traffic), and archived (retired). Stores artifacts on the local filesystem with optional ONNX export. Reads and writes MLflow MLmodel format v1 for interoperability.

Key operations: register_model(), promote_model(), get_model() (pass stage="production" for the live version), get_model_versions(), list_models(), load_artifact().

TrainingPipeline

from kailash_ml.engines.training_pipeline import TrainingPipeline, ModelSpec, EvalSpec

Orchestrates the full training lifecycle: load features, train a model, evaluate metrics, and register the result. Supports scikit-learn and LightGBM model classes out of the box. Model class imports are restricted to a security allowlist (sklearn.*, lightgbm.*, xgboost.*, catboost.*, torch.*, pytorch_lightning.*, lightning.*, kailash_ml.*) to prevent arbitrary code execution. Includes model calibration (Platt scaling, isotonic regression) and optional auto-logging to ExperimentTracker.

Key operations: train(), calibrate(), retrain().

InferenceServer

from kailash_ml.serving.server import InferenceServer, InferenceServerConfig

Loads a registered model and serves predictions over the requested channels (REST, MCP, gRPC). Construction takes an InferenceServerConfig envelope (tenant_id, model_name, model_version, channels, runtime, batch_size); from_registry() is the convenience classmethod that resolves a models://name@alias URI through the registry's alias layer. Single-record and batch prediction flow through predict(features, *, tenant_id=None); the model signature gates input shape (W25 invariant 1).

Key operations: from_registry(), start(), predict(), stop(), health().

DriftMonitor

from kailash_ml.engines.drift_monitor import DriftMonitor, DriftSpec

Detects distribution shifts in production data using PSI (Population Stability Index) and the Kolmogorov-Smirnov test. Stores reference distributions and drift reports in the database. Reports classify drift as none, moderate, or severe per feature and overall. Optionally augments reports with agent-powered interpretation (double opt-in).

Key operations: set_reference_data(), check_drift(), get_drift_history(), check_performance(), schedule_monitoring().

ExperimentTracker

from kailash_ml.engines.experiment_tracker import ExperimentTracker

Provides MLflow-compatible experiment tracking: experiments, runs, parameters, step-based metrics (for training curves), and artifact metadata. Supports nested runs for hierarchical tracking (parent sweep with child trials). Artifacts are stored on the local filesystem; the database holds metadata only -- no binary blobs in SQL.

Key operations: create_experiment(), run() (context manager), start_run(parent_run_id=...), log_param(), log_metric(), log_artifact(), get_run(), list_runs(), list_child_runs(). Standalone usage: await ExperimentTracker.create("sqlite:///ml.db").

Search and Optimization Engines (P1 -- tested, API may evolve)

HyperparameterSearch

from kailash_ml.engines.hyperparameter_search import HyperparameterSearch, SearchSpace, SearchConfig

Supports four search strategies for hyperparameter optimization: grid (exhaustive), random (sampling), bayesian (surrogate-model guided via Optuna), and successive halving (progressive resource allocation with Optuna pruning). Auto-logs trials to ExperimentTracker as nested child runs when a tracker is provided.

Key operations: search().

AutoMLEngine

from kailash_ml.automl import AutoMLEngine, AutoMLConfig
# Equivalent: from kailash_ml import AutoMLEngine

Orchestrates search-strategy sweeps over a user-supplied ParamSpec list, enforces a microdollar cost budget via CostTracker, consults PACT admission, and persists a tenant-scoped audit row per trial. Requires double opt-in for agent augmentation: set agent=True in AutoMLConfig AND have kailash-ml[agents] installed. Cost budgets cap LLM spending via max_llm_cost_usd.

Key operations: run().

EnsembleEngine

from kailash_ml.engines.ensemble import EnsembleEngine

Creates ensemble models through four strategies: blending (weighted average of predictions), stacking (meta-learner on base model outputs), bagging (bootstrap aggregation), and boosting (sequential error correction). All data handling uses polars internally; conversion to sklearn happens at the boundary via interop.py.

Key operations: blend(), stack(), bag(), boost().

PreprocessingPipeline

from kailash_ml.engines.preprocessing import PreprocessingPipeline

Automatic data preprocessing for ML workflows. Auto-detects task type, encodes categoricals, scales numerics (zscore, minmax, robust, maxabs), imputes missing values (mean, median, mode, KNN, iterative), removes multicollinear features, handles class imbalance (SMOTE, ADASYN, class_weight), and splits train/test with stratification support. Returns a SetupResult with transformed data ready for training.

Key operations: setup(), transform(), inverse_transform().

Experimental Engines (P2 -- functional, API may change)

DataExplorer

from kailash_ml.engines.data_explorer import DataExplorer

Computes summary statistics, distributions, correlations, and missing value analysis using polars. Generates interactive plotly visualizations. Optionally augments profiling output with agent-generated narrative (double opt-in).

Key operations: profile(), visualize(), correlation_matrix().

FeatureEngineer

from kailash_ml.engines.feature_engineer import FeatureEngineer

Generates candidate features (interactions, polynomial, binning, temporal) from source data, evaluates their predictive power, and selects the best subset. Returns ranked features with importance scores.

Key operations: generate(), select(), rank().

ModelVisualizer

from kailash_ml.engines.model_visualizer import ModelVisualizer

Produces interactive plotly visualizations for ML model analysis: confusion matrix, ROC curve, precision-recall curve, feature importance, learning curves, residual plots, calibration curves, metric comparison, and training history. All methods return plotly Figure objects.

Key operations: confusion_matrix(), roc_curve(), precision_recall(), feature_importance(), learning_curve().

ModelExplainer

from kailash_ml.engines.model_explainer import ModelExplainer

SHAP-based model explainability for global and local feature importance. Requires pip install kailash-ml[explain]. Provides per-prediction explanations (why did the model make this specific prediction?) alongside global feature importance (which features matter most overall). All visualizations use plotly for consistency with ModelVisualizer.

Key operations: explain_global(), explain_local(), explain_dependence(), to_plotly().

Bridge and Compatibility

OnnxBridge

from kailash_ml.bridge.onnx_bridge import OnnxBridge

Exports trained models to ONNX format for cross-runtime serving ("train in Python, serve in Rust"). Performs a pre-flight compatibility check, exports the model, and validates the ONNX artifact. Export failure is never fatal -- the model falls back to native Python inference. Approximate success rates: sklearn ~90%, LightGBM ~95%, PyTorch feedforward ~70-85%.

Key operations: check_compatibility(), export(), validate().

MlflowFormatReader / MlflowFormatWriter

from kailash_ml.compat.mlflow_format import MlflowFormatReader, MlflowFormatWriter

Reads and writes the MLflow MLmodel YAML format v1. This is format interoperability -- metadata round-trips through MLflow without data loss. kailash-ml does not run an MLflow tracking server, replace experiment tracking, or integrate with the MLflow UI.

Key operations: MlflowFormatReader.read(), MlflowFormatWriter.write().

Quality Tier Promotion

• P2 to P1: 3 integration tests, 2 real-world user validations, no open bugs above LOW
• P1 to P0: No API changes for 3+ minor releases, complete documentation, performance benchmarks

---

Type Contracts (kailash_ml.types)

All cross-engine and cross-framework type contracts live in kailash_ml.types. This module replaced the former kailash-ml-protocols package. Every type supports to_dict() / from_dict() round-trip serialization.

FeatureField

Defines a single feature column.

from kailash_ml.types import FeatureField

field = FeatureField(
    name="age",
    dtype="float64",     # "int64", "float64", "utf8", "bool", "datetime", "categorical"
    nullable=True,
    description="Customer age in years",
)

FeatureSchema

Defines a complete feature set with entity identification and optional timestamps.

from kailash_ml.types import FeatureSchema, FeatureField

schema = FeatureSchema(
    name="customer_features",
    features=[
        FeatureField(name="age", dtype="float64"),
        FeatureField(name="tenure_months", dtype="float64"),
        FeatureField(name="plan_type", dtype="categorical"),
    ],
    entity_id_column="customer_id",
    timestamp_column="event_time",   # optional, enables point-in-time queries
    version=1,
)

# Round-trip serialization
d = schema.to_dict()
restored = FeatureSchema.from_dict(d)

ModelSignature

Captures the input/output schema of a trained model.

from kailash_ml.types import ModelSignature, FeatureSchema, FeatureField

signature = ModelSignature(
    input_schema=FeatureSchema(
        name="churn_input",
        features=[FeatureField(name="age", dtype="float64")],
        entity_id_column="customer_id",
    ),
    output_columns=["prediction", "probability"],
    output_dtypes=["int64", "float64"],
    model_type="classifier",   # "classifier", "regressor", "ranker"
)

MetricSpec

Records a single evaluation metric with its value and context.

from kailash_ml.types import MetricSpec

metric = MetricSpec(
    name="f1",
    value=0.87,
    split="test",              # "train", "val", "test"
    higher_is_better=True,
)

MLToolProtocol

A typing.Protocol that defines the interface for ML tools accessible to Kaizen agents via MCP. Implementors include InferenceServer and ModelRegistry. The protocol requires three methods:

• predict(model_name, features) -- single-record prediction
• get_metrics(model_name) -- model evaluation metrics
• get_model_info(model_name) -- model metadata and stage

AgentInfusionProtocol

A typing.Protocol that defines the interface for agent-augmented engine methods. Implementors are Kaizen Delegate agents (installed via kailash-ml[agents]). Consumers are engines that accept optional agent augmentation. The protocol requires four methods:

• suggest_model(data_profile, task_type) -- model family recommendations
• suggest_features(data_profile, existing_features) -- feature engineering guidance
• interpret_results(experiment_results) -- trial result analysis
• interpret_drift(drift_report) -- drift report interpretation

---

Interop Module

The kailash_ml.interop module is the sole conversion point between polars and external frameworks. No engine file contains direct conversion logic.

from kailash_ml.interop import (
    to_sklearn_input,       # polars -> (numpy X, numpy y) for sklearn
    from_sklearn_output,    # numpy predictions -> polars DataFrame
    to_lgb_dataset,         # polars -> LightGBM Dataset
    to_hf_dataset,          # polars -> HuggingFace Dataset
    polars_to_arrow,        # polars -> PyArrow Table
    from_arrow,             # PyArrow Table -> polars DataFrame
    to_pandas,              # polars -> pandas DataFrame
    from_pandas,            # pandas DataFrame -> polars DataFrame
    polars_to_dict_records, # polars -> list of dicts
    dict_records_to_polars, # list of dicts -> polars DataFrame
)

# Example: prepare data for sklearn
X, y = to_sklearn_input(polars_df, label_col="target")

# Example: convert to pandas for libraries that require it
pandas_df = to_pandas(polars_df)
polars_df = from_pandas(pandas_df)

Optional dependencies (lightgbm, datasets, pyarrow, pandas) are imported lazily so that import kailash_ml never fails due to missing extras.

---

Agent Integration

Six Kaizen agents provide LLM-augmented ML workflows. They follow the LLM-first rule: the LLM does all reasoning, tools are dumb data endpoints.

Double Opt-In Pattern

Agent augmentation requires two explicit steps:

1. Install the agents extra: pip install kailash-ml[agents]
2. Enable in configuration: pass agent=True to the engine config

Without both steps, engines run in pure algorithmic mode.

Five Mandatory Guardrails

Every agent-augmented operation enforces:

1. Confidence scores -- every recommendation includes a self-assessed confidence value (0.0 to 1.0)
2. Cost budget -- cumulative LLM cost is capped at max_llm_cost_usd (raises GuardrailBudgetExceededError if exceeded)
3. Human approval gate -- auto_approve=False by default; the engine pauses for human confirmation before applying agent recommendations
4. Baseline comparison -- a pure algorithmic baseline runs alongside the agent, so you can compare agent-augmented results against non-augmented results
5. Audit trail -- all agent decisions are logged to _kml_agent_audit_log in the database

Agent Catalog

Agent	Purpose	Tier
DataScientistAgent	Data profiling and ML strategy guidance	P0
RetrainingDecisionAgent	Retrain, monitor, or rollback decisions	P0
ModelSelectorAgent	Model family and config recommendations	P1
ExperimentInterpreterAgent	Trial result analysis and next steps	P1
FeatureEngineerAgent	Feature design, generation, and pruning	P2
DriftAnalystAgent	Drift interpretation and action plans	P2

Example: Agent-Augmented AutoML

from kailash_ml.automl import AutoMLEngine, AutoMLConfig

config = AutoMLConfig(
    task_type="classification",
    agent=True,              # Opt-in 1: enable agent augmentation
    auto_approve=False,      # Human approval gate (default)
    max_llm_cost_usd=5.0,   # Cost budget cap
)
engine = AutoMLEngine(
    config=config,
    tenant_id="acme",
    actor_id="alice@acme",
    connection=conn_mgr,           # ConnectionManager for `_kml_automl_trials` audit
)
result = await engine.run(
    space=param_specs,
    trial_fn=user_trial_fn,
    source_tag="agent",
)

# result.best_trial is the TrialRecord with the highest metric_value
print(result.best_trial)
print(result.cumulative_cost_microdollars)

Agent Tools

The kailash_ml.agents.tools module provides dumb data endpoints that agents call. These tools fetch data, compute statistics, and return raw results -- they contain zero decision logic. Available tools include:

• profile_data, get_column_stats, sample_rows (DataScientistAgent)
• compute_feature, check_target_correlation (FeatureEngineerAgent)
• list_available_trainers, get_model_metadata (ModelSelectorAgent)
• get_trial_details, compare_trials (ExperimentInterpreterAgent)
• get_drift_history, get_feature_distribution (DriftAnalystAgent)
• get_prediction_accuracy, trigger_retraining (RetrainingDecisionAgent)

---

Reinforcement Learning Module

Requires pip install kailash-ml[rl] (Stable-Baselines3, Gymnasium, PyTorch).

from kailash_ml.rl import RLTrainer, EnvironmentRegistry, PolicyRegistry
from kailash_ml.rl.trainer import RLTrainingConfig

# Register environment
env_reg = EnvironmentRegistry()
env_reg.register("CartPole-v1")

# Configure policy
policy_reg = PolicyRegistry()
policy_config = policy_reg.get_spec("PPO")

# Train
trainer = RLTrainer(env_registry=env_reg, policy_registry=policy_reg)
config = RLTrainingConfig(
    algorithm="PPO",           # PPO, SAC, DQN, A2C, TD3, DDPG
    policy_type="MlpPolicy",
    total_timesteps=100_000,
    n_eval_episodes=10,
    eval_freq=10_000,
    seed=42,
)
result = await trainer.train("CartPole-v1", "ppo-policy", config=config)
print(f"Mean reward: {result.mean_reward:.1f}")

---

ONNX Bridge

Models registered in ModelRegistry are automatically eligible for ONNX export. The OnnxBridge handles the full pipeline: compatibility check, export, and validation.

from kailash_ml.bridge.onnx_bridge import OnnxBridge

bridge = OnnxBridge()

# Check if a model can be exported
compat = bridge.check_compatibility(trained_model, framework="sklearn")
print(f"ONNX compatible: {compat.compatible}")

# Export
if compat.compatible:
    export_result = bridge.export(
        trained_model, framework="sklearn", n_features=10, output_path="model.onnx"
    )
    print(f"Export success: {export_result.success}")

    # Validate the exported artifact against the original
    validation = bridge.validate(trained_model, export_result.onnx_path, sample_input)
    print(f"Max diff: {validation.max_diff:.6f}")

ONNX export failure is non-fatal. The model falls back to native Python inference. Check model.onnx_status to determine export status.

Framework	Approximate Success Rate	Notes
sklearn	~90%	All standard estimators
LightGBM	~95%	Full support
PyTorch	70-85%	Feedforward networks; dynamic graphs may fail

---

Dashboard

kailash-ml includes a web dashboard for viewing experiments, runs, metrics, and models.

Launch

# CLI entry point
kailash-ml-dashboard

# Or from Python
from kailash_ml.dashboard import MLDashboard

dashboard = MLDashboard(db_url="sqlite:///ml.db")
dashboard.serve(host="0.0.0.0", port=5000)

The dashboard is backed by ExperimentTracker and ModelRegistry engines, served via Starlette and uvicorn. It provides a read-only view of your ML experiments -- no state mutation from the UI.

---

Drift Monitoring

from kailash.db.connection import ConnectionManager
from kailash_ml.engines.drift_monitor import DriftMonitor, DriftSpec

conn = ConnectionManager("sqlite:///ml.db")
await conn.initialize()

# Thresholds are set at construction; tenant_id is required.
monitor = DriftMonitor(
    conn,
    tenant_id="default",
    psi_threshold=0.2,  # PSI > 0.2 = severe drift
    ks_threshold=0.05,  # KS test significance level
)

# Set a reference distribution (e.g., from your training data)
await monitor.set_reference_data(
    "churn_model_v1",
    reference_df,
    feature_columns=["age", "tenure_months", "monthly_spend"],
)

# Check drift against new production data
report = await monitor.check_drift("churn_model_v1", current_df)

print(f"Drift severity: {report.overall_severity}")  # "none", "moderate", "severe"
drifted = [f.feature_name for f in report.feature_results if f.drift_detected]
print(f"Features drifted: {drifted}")

for feature_result in report.feature_results:
    print(f"  {feature_result.feature_name}: PSI={feature_result.psi:.4f}, "
          f"KS={feature_result.ks_statistic:.4f}, drift={feature_result.drift_type}")

---

Experiment Tracking

from kailash.db.connection import ConnectionManager
from kailash_ml.engines.experiment_tracker import ExperimentTracker

# Option 1: Standalone (factory -- manages its own connection)
async with await ExperimentTracker.create("sqlite:///ml.db") as tracker:
    exp_name = await tracker.create_experiment("my-experiment")
    async with tracker.run(exp_name, run_name="baseline") as run:
        await run.log_metric("accuracy", 0.95)
        await run.log_param("n_estimators", "100")

# Option 2: Shared connection (caller manages ConnectionManager lifecycle)
conn = ConnectionManager("sqlite:///ml.db")
await conn.initialize()
tracker = ExperimentTracker(conn, artifact_root="./artifacts")

exp_name = await tracker.create_experiment("churn_classification")

async with tracker.run("churn_classification", run_name="rf_baseline") as run:
    # Log parameters
    await run.log_param("n_estimators", "100")
    await run.log_param("max_depth", "5")

    # Log metrics (supports step-based logging for training curves)
    await run.log_metric("accuracy", 0.92)
    await run.log_metric("f1", 0.87)

# Query results
runs = await tracker.list_runs("churn_classification")
best_run = runs[0]
print(f"Best accuracy: {best_run.metrics['accuracy']}")

await conn.close()

---

Hyperparameter Search

from kailash_ml.engines.hyperparameter_search import (
    HyperparameterSearch,
    SearchSpace,
    SearchConfig,
    ParamDistribution,
)

# Define search space
space = SearchSpace(params=[
    ParamDistribution("n_estimators", "int_uniform", low=50, high=500),
    ParamDistribution("max_depth", "int_uniform", low=3, high=15),
    ParamDistribution("learning_rate", "log_uniform", low=0.001, high=0.3),
    ParamDistribution("subsample", "uniform", low=0.6, high=1.0),
])

config = SearchConfig(
    strategy="bayesian",        # "grid", "random", "bayesian", "successive_halving"
    n_trials=50,
    metric_to_optimize="f1",
    direction="maximize",
)

searcher = HyperparameterSearch(pipeline)
result = await searcher.search(
    data=data,
    schema=schema,
    base_model_spec=base_model_spec,
    search_space=space,
    config=config,
    eval_spec=eval_spec,
    experiment_name="my-search",
)

print(f"Best params: {result.best_params}")
print(f"Best metrics: {result.best_metrics}")
print(f"Best trial: #{result.best_trial_number}")

---

Ensemble Methods

from kailash_ml.engines.ensemble import EnsembleEngine

ensemble = EnsembleEngine()

# Engines are polars-native: pass a DataFrame + target column, not sklearn X/y.

# Blend: weighted average of predictions from multiple models
blend_result = ensemble.blend(
    models=[model_a, model_b, model_c],
    data=train_df,
    target="churned",
    weights=[0.5, 0.3, 0.2],
)

# Stack: train a meta-learner on base model outputs
stack_result = ensemble.stack(
    models=[model_a, model_b, model_c],
    data=train_df,
    target="churned",
    meta_model_class="sklearn.linear_model.LogisticRegression",
)

# Bag: bootstrap aggregation (single base model + n_estimators)
bag_result = ensemble.bag(
    model=base_model,
    data=train_df,
    target="churned",
    n_estimators=10,
)

---

Configuration

Database Setup

kailash-ml uses ConnectionManager from the core Kailash SDK. All engines share the same connection.

from kailash.db.connection import ConnectionManager

# SQLite (default, zero config)
conn = ConnectionManager("sqlite:///ml.db")

# In-memory SQLite (for testing)
conn = ConnectionManager("sqlite:///:memory:")

# PostgreSQL
conn = ConnectionManager("postgresql://user:pass@localhost:5432/mldb")

await conn.initialize()

Engine Initialization

Most engines create their tables on construction (auto_migrate) or lazily on first use; the legacy FeatureStore engine is the one that exposes an explicit initialize().

from kailash_ml.engines.feature_store import FeatureStore  # legacy write engine
from kailash_ml.engines.model_registry import ModelRegistry, LocalFileArtifactStore
from kailash_ml.engines.experiment_tracker import ExperimentTracker
from kailash_ml import DriftMonitor

feature_store = FeatureStore(conn)
await feature_store.initialize()  # legacy engine: explicit table creation

registry = ModelRegistry(conn, LocalFileArtifactStore("./artifacts"))  # ready on construct
tracker = ExperimentTracker(conn, artifact_root="./experiment_artifacts")  # lazy init
monitor = DriftMonitor(conn, tenant_id="default")  # ready on construct; tenant required

Model Class Allowlist

For security, TrainingPipeline and EnsembleEngine restrict model class imports to these prefixes:

• sklearn.*
• lightgbm.*
• xgboost.*
• catboost.*
• torch.*
• pytorch_lightning.*
• lightning.*
• kailash_ml.*

Attempting to load a model class outside this allowlist raises a ValueError. This prevents arbitrary code execution via model class strings.

CLI Entry Points

# Launch the experiment dashboard
kailash-ml-dashboard

# Configure GPU runtime (CUDA/ROCm detection)
kailash-ml-gpu-setup

---

Troubleshooting

"Module kailash_ml has no attribute ..."

Engines are lazy-loaded. Most engines have a single canonical import path under kailash_ml.<engine_module>. Use that path; the top-level lazy re-export is kept for backwards compatibility only.

# Canonical 1.0+ read surface — the top-level default since 2.0.0
from kailash_ml import FeatureStore  # → kailash_ml.features.FeatureStore
# Equivalent explicit form:
from kailash_ml.features import FeatureStore
# Constructor: FeatureStore(dataflow: DataFlow, *, default_tenant_id=None)
# See specs/ml-feature-store.md § 1.1.

# Legacy 0.x write-capable surface — explicit import only (no top-level access)
from kailash_ml.engines.feature_store import FeatureStore
# Constructor: FeatureStore(conn: ConnectionManager, *, table_prefix=...)
# Retained for callers needing the self-contained write path
# (register_features / store / compute / get_training_set / get_features_lazy /
# list_schemas). See MIGRATION.md for the canonical migration recipe.

Migration note (2.0.0 cutover): the top-level from kailash_ml import FeatureStore now resolves to the canonical DataFlow-bridge read surface kailash_ml.features.FeatureStore (constructor FeatureStore(dataflow, *, default_tenant_id=None)). The 1.7.2 bridge release emitted a DeprecationWarning while the top level still resolved to the legacy ConnectionManager-coupled engine; that warning is removed at 2.0.0 because the resolution has flipped. The legacy engine remains importable via its explicit module path kailash_ml.engines.feature_store.FeatureStore for the write/registry/ training-set surface the canonical read store does not expose.

"kailash-ml-protocols not found"

Type contracts have been merged into kailash_ml.types. Update your imports:

# Old (deprecated)
from kailash_ml_protocols import FeatureSchema, FeatureField

# New
from kailash_ml.types import FeatureSchema, FeatureField

ONNX export fails silently

ONNX export failure is by design non-fatal. Check the export status:

version = await registry.get_model("my_model")
print(version.onnx_status)  # "success", "failed", "not_attempted"

Common causes: unsupported model type, dynamic computation graph (PyTorch), custom sklearn estimators.

Agent features not available

Agent augmentation requires both installation and configuration:

pip install kailash-ml[agents]

config = AutoMLConfig(agent=True)  # Must explicitly opt in

If kailash-kaizen is not installed, setting agent=True will raise an ImportError with a helpful message.

"Database is locked" errors

Use ConnectionManager (not bare aiosqlite.connect()). ConnectionManager configures WAL mode, busy timeouts, and connection pooling automatically.

LLM cost budget exceeded

Agent-augmented engines track cumulative LLM cost. If you hit the limit:

config = AutoMLConfig(
    agent=True,
    max_llm_cost_usd=10.0,  # Increase the budget (default varies by engine)
)

---

Contributing

kailash-ml is part of the Kailash Python SDK monorepo.

git clone https://github.com/terrene-foundation/kailash-py.git
cd kailash-py/packages/kailash-ml
uv venv && uv sync
uv run pytest tests/ -x

See the monorepo CONTRIBUTING.md for guidelines on pull requests, testing, and code style.

---

License

Apache-2.0. Copyright 2026 Terrene Foundation.