MLflow Kubernetes Operator Python Client

A Python client library for managing MLflow Kubernetes Operator custom resources in Kubernetes.

Installation

pip install mlflow-k8s-operator

Quick Start

Basic Usage

from mlflow_k8s_operator import MLflowServerClient, MLflowServer, TrackingConfig, BackendConfig, ArtifactStoreConfig

# Initialize the client
client = MLflowServerClient()

# Create an MLflowServer
mlflow_server = MLflowServer(
    name="my-mlflow",
    namespace="mlflow",
    version="2.11.0",
    tracking=TrackingConfig(replicas=2),
    backend=BackendConfig(
        type="postgresql",
        postgresql={
            "host": "postgres.mlflow.svc.cluster.local",
            "database": "mlflow",
            "credentials_secret": "mlflow-db-credentials"
        }
    ),
    artifact_store=ArtifactStoreConfig(
        type="s3",
        s3={
            "bucket": "my-mlflow-artifacts",
            "region": "us-east-1",
            "credentials_secret": "aws-credentials"
        }
    )
)

# Create the resource
client.create(mlflow_server)

# Get the resource
resource = client.get("my-mlflow", "mlflow")
print(resource)

# Check if ready
if client.is_ready("my-mlflow", "mlflow"):
    print("MLflow server is ready!")

# Wait until ready
client.wait_until_ready("my-mlflow", "mlflow", timeout=300)

# List all servers
servers = client.list(namespace="mlflow")
print(servers)

# Delete the resource
client.delete("my-mlflow", "mlflow")

Create from YAML

# Create from YAML string
yaml_content = """
apiVersion: mlops.NotHarshhaa.io/v1alpha1
kind: MLflowServer
metadata:
  name: my-mlflow
  namespace: mlflow
spec:
  version: "2.11.0"
  tracking:
    replicas: 2
  backend:
    type: postgresql
    postgresql:
      host: postgres.mlflow.svc.cluster.local
      database: mlflow
      credentialsSecret: mlflow-db-credentials
  artifactStore:
    type: s3
    s3:
      bucket: my-mlflow-artifacts
      region: us-east-1
      credentialsSecret: aws-credentials
"""

client.create_from_yaml(yaml_content, namespace="mlflow")

# Or from YAML file
client.create_from_yaml_file("mlflow-server.yaml", namespace="mlflow")

Advanced Configuration

from mlflow_k8s_operator import (
    MLflowServer,
    TrackingConfig,
    BackendConfig,
    ArtifactStoreConfig,
    IngressConfig,
    AutoscalingConfig,
    SchedulingConfig,
    PodDisruptionBudgetConfig,
    ProbesConfig
)

mlflow_server = MLflowServer(
    name="advanced-mlflow",
    namespace="mlflow",
    version="2.11.0",
    tracking=TrackingConfig(
        replicas=2,
        additional_args=["--serve-artifacts"],
        probes=ProbesConfig(
            liveness_probe={
                "httpGet": {"path": "/health", "port": 5000},
                "initialDelaySeconds": 30,
                "periodSeconds": 10
            }
        ),
        pod_labels={"environment": "production"}
    ),
    backend=BackendConfig(
        type="postgresql",
        postgresql={
            "host": "postgres.mlflow.svc.cluster.local",
            "database": "mlflow",
            "credentials_secret": "mlflow-db-credentials"
        }
    ),
    artifact_store=ArtifactStoreConfig(
        type="s3",
        s3={
            "bucket": "my-mlflow-artifacts",
            "region": "us-east-1",
            "credentials_secret": "aws-credentials"
        }
    ),
    ingress=IngressConfig(
        enabled=True,
        host="mlflow.example.com",
        tls={"enabled": True, "issuer": "letsencrypt-prod"}
    ),
    autoscaling=AutoscalingConfig(
        enabled=True,
        min_replicas=2,
        max_replicas=10,
        target_cpu_utilization_percentage=70
    ),
    scheduling=SchedulingConfig(
        node_selector={"nodepool": "ml-workload"},
        priority_class_name="high-priority"
    ),
    pod_disruption_budget=PodDisruptionBudgetConfig(
        enabled=True,
        min_available=1
    )
)

client.create(mlflow_server)

Using Custom kubeconfig

# Load from specific kubeconfig file
client = MLflowServerClient(
    kubeconfig="/path/to/kubeconfig",
    context="my-context"
)

Development

Install in development mode

cd python-client
pip install -e ".[dev]"

Run tests

pytest

Format code

black .
isort .

Type checking

mypy mlflow_k8s_operator

License

Apache 2.0 - See LICENSE for details.

Contributing

Contributions are welcome! Please read CONTRIBUTING.md before submitting.