Skip to navigation

Turn On MLOps

View as Markdown

Turn on MLOps by creating one kumo.MLOps configuration and running your existing workflow inside with kumo.mlops(cfg):. When you finish this page, your graph snapshot, training table, and model are stored under an artifact root you choose, your fit is recorded in your experiment tracker, and you can load the stored model and snapshot in a new process.

This page continues from Training. The examples use the same names as that page: graph from Graphs, query from Predictive Queries, and the model plan from Training.

Install the MLOps Extras

Add the extra for your experiment tracker to the extras your data sources already need, such as s3. Also add the extra for your artifact root’s storage, described in Choose an Artifact Root: s3, gcs, databricks for a Unity Catalog volume, or snowflake for a stage:

TrackerInstall command
Self-managed MLflow serverpip install "kumo-engine[mlflow]"
Databrickspip install "kumo-engine[databricks,mlflow]"; in a Databricks notebook, use %pip install and then restart Python as described in Databricks Notebook Setup. Outside Databricks, also add the extra for the cloud the workspace runs on: s3 (AWS), adls (Azure), or gcs (GCP), for example kumo-engine[databricks,mlflow,s3]. Without it, registering or loading a Unity Catalog model fails with an error that names the missing client and the extra
SageMaker managed MLflowpip install "kumo-engine[sagemaker]", which includes the mlflow and s3 extras
Snowflake (preview)pip install "kumo-engine[snowflake]"

Then confirm that the installed Kumo Engine includes the MLOps API:

import kumo.engine as kumo
print(kumo.__version__, hasattr(kumo, "MLOps"))

The command prints the installed version and True. If it prints False, the installed version predates MLOps; install a newer Kumo Engine build.

Choose an Artifact Root

The artifact root is durable storage for everything MLOps stores. Keep it separate from local scratch, which holds working files for the current process only.

Storageartifact_root value
Local or mounted disk"/path/to/kumo" or "file:///path/to/kumo"
Amazon S3"s3://<bucket>/kumo"
Google Cloud Storage"gs://<bucket>/kumo"
Unity Catalog volume, from code running on a Databricks cluster"/Volumes/<catalog>/<schema>/<volume>/kumo"
Unity Catalog volume, from any machine(dbx, "/Volumes/<catalog>/<schema>/<volume>/kumo"), where dbx is a kumo.DatabricksConnector
Snowflake internal stage (preview)(sf, "@ML.MODELS.KUMO/kumo"), where sf is a kumo.SnowflakeConnector

A URI or local path uses your environment’s default credentials, such as an instance role. To use a connector’s own credentials instead, pass a (connector, path) pair whose connector root_dir contains the path, for example:

s3 = kumo.S3Connector(
root_dir="s3://<bucket>/kumo",
identity=kumo.S3Identity(role_arn="<role-arn>"),
)
artifact_root = (s3, "s3://<bucket>/kumo")

For connectors and identities, see Connect Your Data. For a Snowflake stage, create the stage with server-side encryption first, for example CREATE STAGE ML.MODELS.KUMO ENCRYPTION = (TYPE = 'SNOWFLAKE_SSE').

The identity that runs your code needs permission to list, read, write, and delete objects under the root. On Amazon S3, rights scoped to the root’s prefix are enough. With <prefix> the root’s path in the bucket, such as kumo, grant s3:GetObject, s3:PutObject, s3:DeleteObject, and s3:AbortMultipartUpload on <bucket>/<prefix>/*, and s3:ListBucket on the bucket with an s3:prefix condition for <prefix> and <prefix>/*.

Run Your Workflow With MLOps

The following example stores artifacts in S3 and records the fit on a self-managed MLflow server. Replace <bucket> with your bucket and <mlflow-host> with your tracking server’s address. "churn" is the MLflow experiment that receives the runs. For Databricks, SageMaker, and Snowflake, change only the tracking argument as shown in Track Runs.

import kumo.engine as kumo
from kumo.engine.integrations.mlflow import MlflowTracking
cfg = kumo.MLOps(
"s3://<bucket>/kumo",
tracking=MlflowTracking("churn", target="http://<mlflow-host>:5000"),
)
with kumo.mlops(cfg):
snapshot = graph.snapshot()
training_table = query.generate_training_table(snapshot)
trainer = kumo.Trainer(plan)
fit_result = trainer.fit(training_table, snapshot=snapshot)

The first argument of kumo.MLOps is the artifact root. Create the snapshot and the training table inside the block as well, so that they are stored and the model’s lineage can point to them.

Creating kumo.MLOps checks the artifact root, the tracking server, and the experiment before any work starts, so a wrong address, a missing permission, or an unreachable server fails immediately. The tracking argument is optional; without it, Kumo stores artifacts but records no runs.

Check the Result

Print the identifiers of what the block stored and recorded, and keep them with your job’s records:

print("Model id:", fit_result.id)
print("Snapshot id:", fit_result.snapshot_id)
print("Training table id:", training_table.id)
print("Fit run:", fit_result.tracking.run_name, fit_result.tracking.url)

fit_result.tracking describes the run in your tracker: its platform, experiment, run_id, run_name, and url. It is None when cfg has no tracker. Open the URL to see the run; url is None on Snowflake, on SageMaker, and for a local MLflow store, where you find the run by its name instead. Track Runs explains what it contains.

The artifact root now contains one directory per stored artifact. List it with your storage tool, for example aws s3 ls s3://<bucket>/kumo/:

s3://<bucket>/kumo/
├── snapshots/<graph id>/<data version>/
├── training_tables/<training table id>/
└── models/<model id>/

Each stored artifact includes a _provenance.json file that records where it came from: the model names its training table, the table names its snapshot, and the snapshot names its source tables and their versions. Prediction runs appear under predictions/<prediction id>/ after you call trainer.predict() inside the block.

Read Artifacts Back

A scoring job in a new process can reuse what a training job stored, without rebuilding it. Create the same kumo.MLOps configuration in that process, then pass it where a connector usually goes. Replace <model id> and <snapshot id> with the identifiers you printed above:

snapshot = kumo.GraphSnapshot.load_from_connector(cfg, id="<snapshot id>")
trainer = kumo.Trainer.load_from_connector(cfg, path="models/<model id>")

A model takes a path relative to the artifact root; a snapshot takes its id. The loaded trainer is ready to score snapshot. To load a model by a registry alias, such as champion, instead of by its id, see Register and Promote Models.

To retrain on exactly the rows a training job used, load its stored training table and pass it to trainer.fit(). Recreate graph and query first, as described in Recreate the Prediction Context, and replace <training table id> with the id you printed above:

training_table = kumo.TrainingTable.load_from_connector(
cfg,
path="training_tables/<training table id>",
pquery=query,
snapshot=snapshot,
)

A stored prediction run loads the same way. <prediction id> is the id of the PredictionResult that trainer.predict() returned:

prediction = kumo.PredictionResult.load_from_connector(cfg, path="predictions/<prediction id>")

Control Where MLOps Applies

MLOps applies only to calls that run inside a kumo.mlops block or receive an mlops= argument:

  • Block. with kumo.mlops(cfg): covers the calls inside it on the current thread or asyncio task.
  • Single call. Pass mlops=cfg to graph.snapshot(), query.generate_training_table(), trainer.fit(), trainer.predict(), trainer.score_drift(), or trainer.evaluate(). A call’s mlops= replaces the enclosing block, and an inner block replaces an outer one.
  • Outside a block. A call outside any block runs without MLOps, even on a model trained inside one. For a scheduled job, wrap the whole job in one block.
  • Exports. export_config=, output_config=, and export() with an output configuration write only to the destination you give them, as they do without MLOps.

To store only some artifact kinds, pass a set to store=, for example store={"model", "prediction"}. The kinds are "snapshot", "training_table", "model", and "prediction"; the default stores all four.

Handle Failures

By default (on_error="warn"), a failed store, tracking call, or registration produces a warning, and the call returns its result as usual. For scheduled jobs, pass on_error="raise" so the job fails instead of continuing:

cfg = kumo.MLOps(
"s3://<bucket>/kumo",
tracking=MlflowTracking("churn", target="http://<mlflow-host>:5000"),
on_error="raise",
)

With "raise", the call finishes its own work first and then raises a RuntimeError whose cause is the MLOps failure. If trainer.fit() raises this way, the trained model is still available as trainer.result, the TrainingResult that fit() would have returned.

Next Steps