Metadata-Version: 2.4
Name: dbkmlflow
Version: 0.3.0
Summary: A high-level client for managing the ML lifecycle with MLflow and Databricks Unity Catalog.
Author-email: Adolfo Morales <adolfo.morales.gonzalez@gmail.com>
License: MIT License
        
        Copyright (c) 2026 Adolf770
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://github.com/Adolf770/dbkmlflow
Project-URL: Bug Tracker, https://github.com/Adolf770/dbkmlflow/issues
Keywords: mlflow,databricks,unity catalog,mlops
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: mlflow<4,>=3.1
Requires-Dist: pandas
Provides-Extra: dev
Requires-Dist: pytest>=7.4; extra == "dev"
Requires-Dist: pytest-cov>=4.1; extra == "dev"
Requires-Dist: scikit-learn>=1.3; extra == "dev"
Requires-Dist: ruff>=0.6; extra == "dev"
Requires-Dist: mypy>=1.11; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Dynamic: license-file

# dbkmlflow

A high-level client for managing the ML lifecycle with MLflow, usable with a
local or self-hosted MLflow backend, or with Databricks Unity Catalog.

## Features

- One facade (`MLflowManager`) for experiments, runs, model registration,
  aliasing and lifecycle operations (promote, compare, rollback).
- Works against plain open-source MLflow (e.g. `sqlite:///`) or against
  Databricks Unity Catalog — the same code, pointed at a different backend.
- Every public method is keyword-only (except an obvious leading positional
  argument such as the model or the model name), so call sites stay
  self-describing and immune to argument-order mistakes.
- Typed (`py.typed` ships in the wheel) and covered by an honest test suite
  running against a real MLflow backend, not a mocked client.

## Installation

```bash
pip install dbkmlflow
```

## Configuration

### Local or self-hosted MLflow

No Databricks account is required. Point the manager at any MLflow-supported
backend, for example a local SQLite file:

```python
from dbkmlflow import MLflowManager

manager = MLflowManager(tracking_uri="sqlite:///mlflow.db")
```

With only `tracking_uri` given, the model registry defaults to that same
backend — no separate `registry_uri` is needed for local use. The one
exception is Databricks: `tracking_uri="databricks"` resolves to the
`databricks-uc` registry, not to `databricks`, so MLflow's own mapping is
left in charge there rather than being second-guessed.

> **Model names must be three-level everywhere.** `catalog.schema.model` is
> enforced on every backend, including local SQLite — see
> [Model naming](#model-naming-is-enforced-everywhere). Every example below
> uses `main.default.<name>`; substitute your own catalog and schema.

### Databricks Unity Catalog

Both environment variables are required. Without `MLFLOW_REGISTRY_URI=databricks-uc`,
models are logged to Databricks but never actually registered in Unity
Catalog — they stay as plain run artifacts:

```bash
export DATABRICKS_HOST="https://<your-databricks-instance>"
export DATABRICKS_TOKEN="<your-personal-access-token>"
export MLFLOW_TRACKING_URI="databricks"
export MLFLOW_REGISTRY_URI="databricks-uc"
```

With these set, `MLflowManager()` (no arguments) picks them up automatically.
Unity Catalog also requires registered model names to use the three-level
`catalog.schema.model` form, and requires every registered version to carry a
signature — `dbkmlflow` enforces both up front, before the call ever reaches
Databricks, so failures are immediate and explain themselves instead of
surfacing as a generic backend rejection.

### Tag governance (optional)

No tags are required by default — a public library must not reject a first
call over tag names it invented. To enforce a required set for your team:

```python
manager = MLflowManager(required_tags={"owner", "project"})
```

or `export DBKMLFLOW_REQUIRED_TAGS="owner,project"`. `start_run(..., tags=...)`
and `set_tags(...)` then raise `ConfigurationError` if any required tag is
missing, before anything is written.

### Default alias (optional)

**There is no default alias.** A model registered without an explicit
`alias=` gets a version and no alias at all — the library does not invent a
naming convention on your behalf. (Version 0.1.0 silently applied
`developer`; that was removed.)

To have one applied automatically:

```python
manager = MLflowManager(default_alias="challenger")
```

or `export DBKMLFLOW_DEFAULT_ALIAS="challenger"`. When set, `log_model(...)`,
`log_pyfunc_model(...)` and `track(...)` apply it whenever `alias=` is
omitted. Precedence is `alias=` argument → `default_alias` → no alias.

One asymmetry worth knowing: **`rollback()` assumes `alias="champion"`** — it
is the only method with a built-in alias default, because rolling back is a
production operation. If your production alias is named something else, pass
it explicitly or the rollback will fail looking for an alias that does not
exist.

## Thread safety

**Not thread-safe.** MLflow's active run is process-global state, and the
flavor `log_model` functions (`mlflow.sklearn.log_model`, etc.) take no run
id — they always act on whatever run is currently active. `MLflowManager`
scopes MLflow's tracking/registry URIs around each call so that multiple
*managers* in one process don't clobber each other's configuration, but it
cannot make concurrent *runs* on separate threads safe: use one manager, and
one run at a time, per process.

---

## Usage examples: Local MLflow (SQLite)

### Quick start — train, register, load and predict

```python
import pandas as pd
from sklearn.ensemble import RandomForestRegressor

from dbkmlflow import MLflowManager

X = pd.DataFrame({"f1": [1, 2, 3], "f2": [4, 5, 6]})
y = pd.Series([1.0, 2.0, 3.0])
model = RandomForestRegressor(n_estimators=10).fit(X, y)

manager = MLflowManager(tracking_uri="sqlite:///mlflow.db")

manager.start_run(experiment_name="demo")
manager.log_params({"n_estimators": 10})
manager.log_metrics({"r2": 0.95})
version = manager.log_model(
    model,
    artifact_path="model",
    flavor="sklearn",
    registered_model_name="main.default.demo_model",
    alias="champion",
    input_example=X.head(1),
)
manager.end_run()
print("registered version:", version.version)

loaded_model, metadata = manager.load_model("main.default.demo_model", alias="champion")
print("predictions:", loaded_model.predict(X.head(1)))
print("loaded version:", metadata["model_info"]["version"])
```

`dbkmlflow` automatically infers the correct MLflow flavor (like `sklearn`, `xgboost`, or `lightgbm`) by inspecting your model's module, so you don't even need to pass `flavor="sklearn"` anymore. If you prefer, you can still pass standard tags like `"RandomForest"` and it will be mapped properly.

### Using the singleton helper

For a shared manager rather than a local variable, use the singleton helper
instead of constructing `MLflowManager` directly:

```python
from dbkmlflow import get_mlflow_manager, reset_manager

manager = get_mlflow_manager(tracking_uri="sqlite:///mlflow.db")
# ... later, in the same process, from anywhere:
manager = get_mlflow_manager()  # returns the same instance
reset_manager()  # ends any active run and clears it, so the next call
                  # can reconfigure (e.g. in tests, between environments)
```

### Using as a context manager

```python
from dbkmlflow import MLflowManager

manager = MLflowManager(tracking_uri="sqlite:///mlflow.db")
with manager.run(experiment_name="ctx_demo"):
    manager.log_params({"learning_rate": 0.01})
    manager.log_metrics({"accuracy": 0.92})
    version = manager.log_model(
        model,
        artifact_path="model",
        registered_model_name="main.default.ctx_model",
        alias="champion",
        input_example=X.head(1),
    )
# end_run() is called automatically when exiting the context
```

### Decorator & Autologging

`@track` wraps a training function: it starts the run, executes, and ends the
run — including when the function raises.

```python
manager = MLflowManager(tracking_uri="sqlite:///mlflow.db")

@manager.track(experiment_name="demo")
def train_model(X, y):
    manager.log_metrics({"r2": 0.95})
    return RandomForestRegressor().fit(X, y)

model = train_model(X, y)
```

To register the trained model in the same call, add `registered_model_name`
**and** an `input_example` (or `signature`) — registering without one is
rejected, so the decorator needs it too:

```python
@manager.track(
    experiment_name="demo",
    registered_model_name="main.default.tracked_model",
    alias="challenger",
    input_example=X.head(1),          # required whenever registering
)
def train_model(X, y):
    return RandomForestRegressor().fit(X, y), {"r2": 0.95}
```

Return either the model alone, or a `(model, metrics)` tuple — the metrics
dict is logged automatically before the model is registered.

To enable framework autologging, call `manager.autolog()` before training;
it forwards to `mlflow.autolog()` inside the manager's scope.

---

## Usage examples: Databricks Unity Catalog

### Setup

```bash
# Set these environment variables before running your script:
export DATABRICKS_HOST="https://adb-1234567890.12.azuredatabricks.net"
export DATABRICKS_TOKEN="dapi1234567890abcdef"
export MLFLOW_TRACKING_URI="databricks"
export MLFLOW_REGISTRY_URI="databricks-uc"
```

### Model naming is enforced everywhere

Registered model names **must** use the three-level `catalog.schema.model`
form, on every backend — Unity Catalog, a self-hosted server, or local SQLite.
This is deliberate: it is a naming standard, not a Unity Catalog technicality.
If a local run accepted `my_model` while production required
`main.default.my_model`, local runs would stop predicting what production
accepts, and the failure would surface at deployment instead of at development
time.

```python
manager.log_model(model, artifact_path="model",
                  registered_model_name="my_model")          # ConfigurationError
manager.log_model(model, artifact_path="model",
                  registered_model_name="main.default.my_model")   # OK
```

Projects that have not adopted the standard can relax it per manager:

```python
MLflowManager(uc_validation="auto")   # three-level only on Unity Catalog
MLflowManager(uc_validation="never")  # no structural check at all
```

### Train and register a model

```python
import pandas as pd
from sklearn.ensemble import RandomForestRegressor

from dbkmlflow import MLflowManager

X = pd.DataFrame({"f1": [1, 2, 3], "f2": [4, 5, 6]})
y = pd.Series([1.0, 2.0, 3.0])
model = RandomForestRegressor(n_estimators=10).fit(X, y)

# Environment variables are picked up automatically — no arguments needed
manager = MLflowManager()

manager.start_run(experiment_name="/Users/user@company.com/demo")
manager.log_params({"n_estimators": 10})
manager.log_metrics({"r2": 0.95})
version = manager.log_model(
    model,
    artifact_path="model",
    flavor="sklearn",
    registered_model_name="my_catalog.my_schema.demo_model",  # Three-level name required
    alias="champion",
    input_example=X.head(1),  # Required for Unity Catalog
)
manager.end_run()
print("Registered version:", version.version)
```

### Load a model from Unity Catalog

```python
from dbkmlflow import MLflowManager

manager = MLflowManager()

loaded_model, metadata = manager.load_model(
    "my_catalog.my_schema.demo_model",
    alias="champion",
)
print("predictions:", loaded_model.predict(X.head(1)))

# Or load by specific version number:
loaded_model, metadata = manager.load_model(
    "my_catalog.my_schema.demo_model",
    version=1,
)
```

### Explicit URIs (without environment variables)

```python
from dbkmlflow import MLflowManager

manager = MLflowManager(
    tracking_uri="databricks",
    registry_uri="databricks-uc",
)
```

---

## Model lifecycle: aliases, promotion, comparison, rollback

```python
# Point an alias at a specific version, or move one between aliases.
manager.set_alias("main.default.demo_model", alias="challenger", version=2)
manager.promote("main.default.demo_model", from_alias="challenger", to_alias="champion")
manager.delete_alias("main.default.demo_model", alias="challenger")

# Compare the challenger against the current champion and promote if better.
# This wraps compare_models and promote in a single call.
promoted = manager.compare_and_promote(
    "main.default.demo_model",
    metric="r2",
    higher_is_better=True,
    champion_alias="champion",
    challenger_alias="challenger"
)

# Roll @champion back to the newest earlier version not marked defective.
manager.rollback("main.default.demo_model", reason="accuracy drift in production")

# Params, metrics and tags for a specific version or alias.
metadata = manager.get_model_metadata("main.default.demo_model", alias="champion")
```

`compare_models` and `compare_and_promote` never silently promote on a backend
failure: a genuinely missing champion alias is required before a challenger is
promoted by default.

## Downloading and saving artifacts

```python
# Download a registered version's artifacts to a local directory.
local_path = manager.download_model_artifacts(
    "main.default.demo_model", alias="champion", dst_path="./downloaded_model"
)

# Save a model to a local directory without touching the registry at all.
# This static method can be called directly on the class.
MLflowManager.save_local(
    model, path="./local_model", metrics={"r2": 0.95}
)
# Re-running against the same path fails by default (MLflow requires an
# empty destination); pass overwrite=True to replace it.
MLflowManager.save_local(
    model, path="./local_model", overwrite=True
)
```

> **Note:** For safety, `save_model_locally` refuses to overwrite dangerous
> paths (`.`, home directory, filesystem root) even when `overwrite=True`.

## Custom (pyfunc) models

```python
version = manager.log_pyfunc_model(
    my_python_model,
    artifact_path="model",
    config={"threshold": 0.5},
    registered_model_name="main.default.custom_model",
    alias="champion",
    input_example=X.head(1),
)
```

## Complete workflow example (local)

Here's a full training pipeline example using a local SQLite backend:

```python
import pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score, mean_squared_error

from dbkmlflow import MLflowManager

# 1. Prepare data
X = pd.DataFrame({
    "feature_1": range(100),
    "feature_2": range(100, 200),
    "feature_3": [x * 0.5 for x in range(100)],
})
y = pd.Series([x * 2.1 + 3 for x in range(100)])
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# 2. Train model
model = RandomForestRegressor(n_estimators=50, max_depth=10)
model.fit(X_train, y_train)
predictions = model.predict(X_test)

# 3. Calculate metrics
r2 = r2_score(y_test, predictions)
mse = mean_squared_error(y_test, predictions)

# 4. Register with MLflow
manager = MLflowManager(tracking_uri="sqlite:///mlflow.db")

manager.start_run(experiment_name="training_pipeline")
manager.log_params({"n_estimators": 50, "max_depth": 10})
manager.log_metrics({"r2": r2, "mse": mse})

version = manager.log_model(
    model,
    artifact_path="model",
    flavor="sklearn",
    registered_model_name="main.default.production_model",
    alias="challenger",
    input_example=X_train.head(1),
)
manager.end_run()

# 5. Compare against champion and promote if better
result = manager.compare_models(
    "main.default.production_model", metric="r2", higher_is_better=True
)
if result.should_promote:
    manager.promote(
        "main.default.production_model", from_alias="challenger", to_alias="champion"
    )
    print(f"Promoted version {result.challenger_version} to champion!")
else:
    print(f"Challenger (r2={result.challenger_metric:.4f}) did not beat "
          f"champion (r2={result.champion_metric:.4f})")

# 6. Download champion model locally for serving
local_path = manager.download_model_artifacts(
    "main.default.production_model", alias="champion", dst_path="./serving_model"
)
print(f"Champion model downloaded to: {local_path}")
```

## Complete workflow example (Databricks Unity Catalog)

```python
import pandas as pd
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.metrics import accuracy_score, f1_score

from dbkmlflow import MLflowManager

# Environment variables must be set:
# DATABRICKS_HOST, DATABRICKS_TOKEN, MLFLOW_TRACKING_URI, MLFLOW_REGISTRY_URI

# 1. Initialize (picks up env vars automatically)
manager = MLflowManager()

# 2. Train
X_train = pd.DataFrame({"f1": range(100), "f2": range(100, 200)})
y_train = pd.Series([0] * 50 + [1] * 50)
model = GradientBoostingClassifier(n_estimators=100).fit(X_train, y_train)

# 3. Register with three-level model name
manager.start_run(experiment_name="/Users/user@company.com/classification")
manager.log_params({"n_estimators": 100, "model_type": "GradientBoosting"})
manager.log_metrics({
    "accuracy": accuracy_score(y_train, model.predict(X_train)),
    "f1": f1_score(y_train, model.predict(X_train)),
})

version = manager.log_model(
    model,
    artifact_path="model",
    flavor="sklearn",
    registered_model_name="prod_catalog.ml_schema.classifier_v1",
    alias="challenger",
    input_example=X_train.head(1),
)
manager.end_run()

# 4. Promote if better than champion
result = manager.compare_models(
    "prod_catalog.ml_schema.classifier_v1",
    metric="f1",
    higher_is_better=True,
)
if result.should_promote:
    manager.promote(
        "prod_catalog.ml_schema.classifier_v1",
        from_alias="challenger",
        to_alias="champion",
    )

# 5. Load champion for inference
loaded_model, metadata = manager.load_model(
    "prod_catalog.ml_schema.classifier_v1",
    alias="champion",
)
predictions = loaded_model.predict(X_train.head(5))
print(f"Version: {metadata['model_info']['version']}, Predictions: {predictions}")
```

## Error handling

Every method translates MLflow's own `MlflowException` into one of
`dbkmlflow`'s own exception types before it reaches the caller — no raw
`MlflowException` crosses the public boundary, and the original is always
preserved as `__cause__`:

- `ConfigurationError` — invalid arguments, a rejected model name, or a
  missing required tag.
- `RunError` — an operation needed an active run and there wasn't one.
- `ModelNotFoundError` — the model, version or alias genuinely does not
  exist.
- `RegistryError` — the registry rejected the call for any other reason
  (permissions, network, backend failure). Kept distinct from
  `ModelNotFoundError` so a permissions error is never reported as a missing
  model.
- `ArtifactError` — logging, downloading or saving artifacts failed.

All of them subclass `DbkmlflowError`.

```python
from dbkmlflow import MLflowManager, ModelNotFoundError, ConfigurationError

manager = MLflowManager(tracking_uri="sqlite:///mlflow.db")

try:
    model, meta = manager.load_model("main.default.nonexistent_model", alias="champion")
except ModelNotFoundError as e:
    print(f"Model not found: {e}")
except ConfigurationError as e:
    print(f"Configuration error: {e}")
```

## Environment variables reference

| Variable | Description | Example |
|---|---|---|
| `MLFLOW_TRACKING_URI` | MLflow tracking server URI | `sqlite:///mlflow.db` or `databricks` |
| `MLFLOW_REGISTRY_URI` | MLflow model registry URI | `databricks-uc` |
| `DATABRICKS_HOST` | Databricks workspace URL | `https://adb-123.12.azuredatabricks.net` |
| `DATABRICKS_TOKEN` | Databricks personal access token | `dapi1234567890abcdef` |
| `DBKMLFLOW_REQUIRED_TAGS` | Comma-separated required tags | `owner,project,team` |
| `DBKMLFLOW_DEFAULT_ALIAS` | Alias applied when `alias=` is omitted (no default) | `challenger` |

## Development

```bash
pip install -e ".[dev]"
ruff check src tests
mypy src
pytest --cov=dbkmlflow --cov-report=json --cov-report=term-missing
python scripts/check_coverage.py  # enforces an 85% floor per file, not just overall
```

See `CHANGELOG.md` for what changed in each release.
