MLflow Utils

MLflow integration utilities for the DaviesTechLabs AI/ML platform.

Installation

pip install -r requirements.txt

Or from Gitea:

pip install git+https://git.daviestechlabs.io/daviestechlabs/mlflow.git

Modules

Module	Description
client.py	MLflow client configuration and helpers
tracker.py	General MLflowTracker for experiments
inference_tracker.py	Async inference metrics for NATS handlers
model_registry.py	Model Registry with KServe metadata
kfp_components.py	Kubeflow Pipeline MLflow components
experiment_comparison.py	Compare experiments and runs
cli.py	Command-line interface

Quick Start

from mlflow_utils import get_mlflow_client, MLflowTracker

# Simple tracking
with MLflowTracker(experiment_name="my-experiment") as tracker:
    tracker.log_params({"learning_rate": 0.001})
    tracker.log_metrics({"accuracy": 0.95})

Inference Tracking

For NATS handlers (chat-handler, voice-assistant):

from mlflow_utils import InferenceMetricsTracker
from mlflow_utils.inference_tracker import InferenceMetrics

tracker = InferenceMetricsTracker(
    experiment_name="voice-assistant-prod",
    batch_size=100,  # Batch metrics before logging
)

# During request handling
metrics = InferenceMetrics(
    request_id="uuid",
    total_latency=1.5,
    llm_latency=0.8,
    input_tokens=150,
    output_tokens=200,
)
await tracker.log_inference(metrics)

Model Registry

Register models with KServe deployment metadata:

from mlflow_utils.model_registry import register_model_for_kserve

register_model_for_kserve(
    model_name="my-qlora-adapter",
    model_uri="runs:/abc123/model",
    kserve_runtime="kserve-vllm",
    gpu_type="amd-strixhalo",
)

Kubeflow Components

Use in KFP pipelines:

from mlflow_utils.kfp_components import (
    log_experiment_component,
    register_model_component,
)

CLI

# List experiments
python -m mlflow_utils.cli list-experiments

# Compare runs
python -m mlflow_utils.cli compare-runs --experiment "qlora-training"

# Export metrics
python -m mlflow_utils.cli export --run-id abc123 --output metrics.json

Configuration

Environment Variable	Default	Description
MLFLOW_TRACKING_URI	http://mlflow.mlflow.svc.cluster.local:80	MLflow server
MLFLOW_EXPERIMENT_NAME	default	Default experiment
MLFLOW_ENABLE_ASYNC	true	Async logging for handlers

Module Structure

mlflow_utils/
├── __init__.py              # Public API
├── client.py                # Connection management
├── tracker.py               # General experiment tracker
├── inference_tracker.py     # Async inference metrics
├── model_registry.py        # Model registration + KServe
├── kfp_components.py        # Kubeflow components
├── experiment_comparison.py # Run comparison tools
└── cli.py                   # Command-line interface

Related

• handler-base - Uses inference tracker
• kubeflow - KFP components
• argo - Training workflows
• homelab-design - Architecture docs