7.1 Register MLflow Models and Package Feature Retrieval Specs

Register MLflow Models and Package Feature Retrieval Specs

Quick Answer: An MLflow model is a folder, not a lone pickle. Flavors such as sklearn, pytorch, and pyfunc write MLmodel, conda.yaml (or pip requirements), the serialized weights, plus optional signature and input example. Log with mlflow.<flavor>.log_model (or autolog) inside a job, then register with az ml model create using --type mlflow_model and a runs:/ or azureml://jobs/ path. Put feature_retrieval_spec.yaml at the artifact root so online endpoints look up the same features the trainer used. Do not register a raw pickle as custom_model if you need no-code deploy, Swagger, or the Responsible AI dashboard.

Exam AI-300 Domain 2 asks you to register an MLflow model and package a feature retrieval specification with the model artifact. Chapters 5 and 6 already logged metrics, trained command jobs, and compared runs. This section is how a trained artifact becomes a versioned workspace (or registry) model asset that endpoints, Responsible AI, and promotion can consume.

What an MLflow model actually is

MLflow treats a model as a directory contract, not a single file. Typical contents of that folder:

• MLmodel — YAML that names the flavor (sklearn, pytorch, tensorflow, xgboost, pyfunc, and others), the Python function flavor for generic loading, and the signature if you logged one.
• Weights — model.pkl, a PyTorch data/ tree, an ONNX file, or a custom artifact the flavor understands.
• Environment — conda.yaml and often requirements.txt so deploy can rebuild the runtime without a hand-authored environment asset.
• Optional input example — a sample row or tensor that Azure Machine Learning studio uses to generate Swagger and to drive the Test pane on a no-code online endpoint.

Azure Machine Learning is compatible with MLflow 2.16 and earlier. MLflow 2.17 and later change artifact repositories and the LoggedModels API in ways the azureml-mlflow plugin does not currently support. Pin mlflow 2.16.2 or earlier (and matching framework pins such as xgboost 2.1.1 or earlier when you use that flavor) in training environments.

Logging a model is not the same as mlflow.log_artifact("model.pkl"). Artifact logging stores a blob. Model logging stores a loadable, deployable package. Microsoft lists the operational payoffs of that package:

• mlflow.<flavor>.load_model and mlflow.pyfunc.load_model restore an object with predict.
• Pipeline steps can take type: mlflow_model inputs and outputs.
• You can deploy without a scoring script or a hand-authored environment.
• Managed online endpoints get Swagger, so studio Test works.
• The Responsible AI dashboard in the next section requires a registered MLflow sklearn-flavor model.

mlflow.autolog() captures models for many frameworks when you call fit. Switch to mlflow.<flavor>.log_model(..., signature=infer_signature(X, y), input_example=X.sample(n=1)) when autolog infers a bad signature (common with tensors), when you need a custom conda environment, or when you wrap several artifacts in mlflow.pyfunc.log_model. save_model writes the same folder to local disk without an active run; log_model writes it into the run. For custom logic, inherit mlflow.pyfunc.PythonModel, implement predict, and log the wrapper. Serializable sklearn-API objects can still use mlflow.sklearn.log_model even when they were not trained with scikit-learn.

Trap: uploading model.pkl with type: custom_model. That is a legal v2 asset, but you lose no-code deploy, automatic Swagger, Responsible AI, and flavor-aware loading. Models registered with SDK/CLI v1 also land as custom. This skill-measured bullet wants mlflow_model.

Register with SDK/CLI v2

The workspace model registry versions assets by name. Registering the same name again creates the next integer version. Three types matter on this exam:

Supported path syntax:

Register from a job so the model keeps lineage to the trainer:

az ml model create --name claims-fraud --version 3 --path runs:/<run-id>/model --type mlflow_model

The equivalent job URI is azureml://jobs/<job-id>/outputs/artifacts/paths/model/ — artifacts is the reserved name of the default MLflow artifact location. SDK v2 uses Model(path="runs:/<run-id>/model/", name="claims-fraud", type=AssetTypes.MLFLOW_MODEL) then ml_client.models.create_or_update. You can also call mlflow.register_model(f"runs:/{run_id}/{artifact_path}", model_name) or mlflow.register_model(f"file://{abs_path}", name) from a local MLflow folder produced by save_model.

MLflow registration has sharp edges in Azure Machine Learning: you can register only into the same workspace that tracked the run; organizational registries are not supported through the MLflow client; cross-workspace copies are not supported through MLflow. Studio Models → Register still works for local files, job output, or datastore and lets you pick MLflow, Triton, or unspecified.

Jobs can take models as inputs (download or ro_mount for MLflow) and write models as outputs (upload or rw_mount) with type: mlflow_model. That is how a pipeline train step emits an MLflow folder that a later register step — or a one-line az ml model create from the job URI — consumes.

Feature retrieval specification

A feature retrieval specification is a portable YAML list of features that already exist in managed feature store feature sets. Features can come from multiple feature sets and even from multiple feature stores. Training uses the spec to build a point-in-time joined training table. Inference uses the same spec to look up values from the online store, so the endpoint does not invent a different feature list than the trainer.

Generate the spec with the feature store SDK (generate_feature_retrieval_spec after resolve_feature_uri). The file name cannot change: it must be feature_retrieval_spec.yaml, and it must sit in the root of the model artifact:

<model folder>/
  MLmodel
  conda.yaml
  model.pkl
  feature_retrieval_spec.yaml

The exam pairs this file with registration for two reasons:

1. Lineage. Studio shows model ↔ feature-set lineage only when that YAML exists in the registered artifact. Without it, the model detail page and the feature-set detail page do not link.
2. Online scoring. The scoring script's init() loads the spec from AZUREML_MODEL_DIR with FeatureStoreClient.resolve_feature_retrieval_spec, then init_online_lookup / get_online_features. Online lookup requires azureml-featurestore 1.2.1 or later and a managed identity that can read the feature store.

If you used the built-in feature retrieval pipeline component, the spec is already copied under the training-data folder root. Training code should copy it into the model output:

shutil.copy(os.path.join(args.training_data, "feature_retrieval_spec.yaml"), args.model_output)

You can also pass the spec as a job input and copy it in the script. For batch inference, the same built-in component accepts the registered model as input_model and reads the packaged spec from the artifact — you do not have to pass a separate spec folder.

The spec is optional if you call get_offline_features() with an explicit feature list during experimentation. It is not optional once you want studio lineage and consistent online lookup for a registered production model.

Exam scenario

A fraud team trains a gradient-boosted classifier on point-in-time features from transactions:3 and accounts:2. The command job logs with mlflow.xgboost.log_model, passes signature=infer_signature(...) and a one-row input_example, and copies feature_retrieval_spec.yaml next to MLmodel. After the job succeeds, MLOps runs az ml model create --name fraud-gbm --path azureml://jobs/<job>/outputs/artifacts/paths/model --type mlflow_model. The managed online endpoint's init() resolves that spec and looks up features from the online store using the request's account_id. Studio's model page shows lineage to those feature sets, and no-code Test works because of the signature.

Common trap

Registering model.pkl as custom_model because "the file is the model," then wondering why no-code deploy, Swagger, Responsible AI, and feature-set lineage are missing. Sibling traps: renaming feature_retrieval_spec.yaml; burying the spec in a subfolder; logging only mlflow.log_artifact so there is no MLmodel; using MLflow 2.17+ with azureml-mlflow; expecting mlflow.register_model to publish into an organizational registry; hand-coding a different feature list in the scoring script than the spec the trainer packaged.