mlflow-triton

For demos and proofs of concept — not production-ready serving.

Before production use, an ML Engineer must review, validate and optimize the serving implementation and, where needed, convert the model to a suitable production format/backend. The generated Triton Python backend uses MLflow's prediction interface as a starting point. Successful example tests do not establish production readiness, throughput, latency or resource efficiency.

Core MLflow logging wrappers for sklearn, PyTorch and Hugging Face. Generate a portable Triton repository for Data Scientist demos and PoCs, without requiring Triton Control, deployment credentials or an LLM. Production optimization remains a separate ML Engineering task.

pip install mlflow-triton

Requires Python 3.10+ and MLflow >=3.14,<4. Install the framework dependencies appropriate to your saved model separately. This package registers no deployment entry point and makes no Triton Control API calls. config.pbtxt is generated in Protobuf text format, using Protobuf string escaping.

Log native models with a required example request

Use the plugin's framework module instead of mlflow.<framework>.log_model:

from mlflow_triton import sklearn, pytorch, transformers

# Already-fitted Pipeline, including its original preprocessing.
info = sklearn.log_model(
    fitted_pipeline,
    name="classifier",
    input_example=example_dataframe,
    registered_model_name="classifier",
)

# Instantiated custom architecture with its saved trained weights loaded.
# Include inference-only class/helper modules, not a training entry point.
torch_info = pytorch.log_model(
    loaded_model.eval(),
    name="network",
    input_example=example_numpy_array,
    code_paths=["inference_architecture"],
    pip_requirements=["torch==<your-version>", "numpy==<your-version>"],
)

# A supported Transformers pipeline containing the model and tokenizer/processor.
hf_info = transformers.log_model(
    loaded_pipeline,
    name="text_model",
    input_example=["Example input"],
    model_config={"max_new_tokens": 32, "do_sample": False, "return_full_text": True},
)

print(info.model_uri)

input_example is a mandatory keyword argument even when signature= is supplied. It is the model's prediction input (DataFrame, NumPy array, named tensor dictionary, or supported Transformers task input), not a Triton HTTP request envelope. MLflow stores the example with the model: use appropriate non-sensitive examples. No training is run. By default, logging executes inference using the supplied model on its current device, without reloading its weights or moving it to CPU. Allow memory for example inference and native serialization. Set validate_saved_model=True to instead reload the saved artifact on CPU for validation; this requires memory for an additional model copy. Framework packages such as torch/transformers remain optional plugin dependencies; install the versions required by your model.

The wrappers preserve an explicit signature. Otherwise they infer one with MLflow from the example and model prediction. They reject missing/incompatible schemas before logging or registration. Native flavor options such as code_paths, requirements, sklearn's pyfunc_predict_fn, and Transformers' task, processor and model_config are forwarded to native packaging. name, registered_model_name, await_registration_for, run_id, metadata, tags and logging params are supported. Logging params are tracking metadata, not generation parameters: use model_config for Transformers settings. This is an explicit wrapper API; it does not monkey-patch MLflow or autologging.

The logged artifact contains the native MLflow flavors and an additional triton flavor pointing to:

model_repository/<name>/
  config.pbtxt
  requirements.txt
  request.json
  validation.json
  1/
    model.py
    artifact/          # Complete local MLflow package, including code and assets

Native mlflow.pyfunc.load_model(info.model_uri) and native framework loaders remain usable. The embedded package is a separate copy so the Triton directory can be deployed independently; this increases artifact storage size. Dependencies are recorded, not installed into the running Triton server automatically. Build an environment from requirements.txt with a Python version matching Triton's Python backend stub; custom code dependencies must be included or installed.

The generated configuration uses backend: "python" and max_batch_size: 0. Tensor shapes include their full batch dimension. Column schemas map to separate rank-one tensors; tensor schemas retain their declared shapes. The mlflow_contract parameter in config.pbtxt records mappings between original field names and Triton-safe tensor names. Only dynamic axes already declared by the signature are retained; a single example does not prove that other sequence/image dimensions are variable. Supply an explicit signature when those dimensions need to be dynamic.

device="cpu" is the default. device="gpu" emits KIND_GPU for PyTorch or Transformers; model.py uses the assigned Triton instance device when loading. Default validation uses the supplied model's current device, independently of the Triton target. Optional saved-artifact validation loads on CPU even for a GPU target. Sharded/device-mapped models or tasks unsupported by MLflow pyfunc need a separate adapter; the wrapper does not claim universal model support.

PyTorch uses Python/pickle packaging by default, not an implicit LibTorch/PT2 export. Supply required architecture modules with code_paths or dependencies. Pre/postprocessing outside the supplied model/Pipeline is not captured. Transformers packages preserve their saved tokenizer, chat template and model configuration, and expose the MLflow task interface (which can differ from raw forward() or pipeline outputs). Serving is unary prediction, not streaming or a dedicated high-throughput LLM engine. No GGUF/Q4_K_M compatibility is implied.

Supported signatures contain NumPy-compatible numeric/string tensors or required scalar DataFrame columns. Optional columns, nested/object schemas, datetime columns, and scalar tensors require an explicit serving adapter and fail clearly rather than being flattened or cast speculatively.

validation.json identifies in-memory or saved-artifact prediction validation, records whether the saved artifact was checked, and distinguishes local mapping validation from actual Triton validation. Logging does not launch Triton, deploy a service, or establish GPU correctness. The direct triton.log_model() API below remains available. Deployment is a separate concern.

Log API reference

The three framework wrappers have these public entry points (common keyword options are abbreviated below; the Python signatures expose them explicitly):

sklearn.log_model(sk_model, *, input_example, **kwargs)
pytorch.log_model(pytorch_model, *, input_example, **kwargs)
transformers.log_model(transformers_model, *, input_example, **kwargs)

Pass the model positionally or by its exact parameter name. input_example and all other options must be keyword arguments. The tables below describe the options handled by this plugin and commonly used native packaging options.

Required arguments

input_example=None and empty examples are rejected. The example remains mandatory with an explicit signature. It is saved in the artifact and used for prediction and tensor-mapping validation; it is not training data or a Triton HTTP request. PyTorch models with any submodule still in training mode are rejected.

Common optional arguments

These arguments are accepted by all three framework wrappers through **kwargs.

The MLflow Registry version and Triton filesystem version are independent: logging always generates <triton_model_name>/1/, even when the Registry creates version 2 or higher. There is no triton_version argument.

Native packaging arguments

Remaining keyword arguments are forwarded to mlflow.<framework>.save_model, not to that framework's log_model. Availability and defaults therefore depend on the installed MLflow version. This wrapper does not accept every argument exposed by native log_model; for example, use name, not artifact_path. Unknown or incompatible options fail during native packaging.

Default logging skips MLflow's automatic dependency inference because it can reload the saved model in a subprocess. Framework defaults may omit custom dependencies; supply them with pip_requirements or extra_pip_requirements.

Custom architecture code and third-party dependencies must be available when the saved artifact is reloaded. code_paths bundles local code; requirement arguments record installable dependencies. Neither reconstructs a missing architecture. External preprocessing must already be part of the supplied prediction interface.

Examples with explicit options

The objects below are already loaded/fitted; these calls do not train them. Configure the MLflow tracking URI/experiment using MLflow's normal API before logging, if required for your environment.

from mlflow_triton import sklearn, pytorch, transformers

# sklearn: named DataFrame columns and preprocessing travel with the Pipeline.
sk_info = sklearn.log_model(
    sk_model=fitted_pipeline,
    input_example=example_dataframe.iloc[:2],
    name="classifier",
    device="cpu",
    pyfunc_predict_fn="predict_proba",
    registered_model_name="demo-classifier",
)

# PyTorch: custom architecture instance, trained weights already loaded.
pt_info = pytorch.log_model(
    pytorch_model=loaded_model.eval(),
    input_example=example_numpy_array[:2],
    name="network",
    device="gpu",
    code_paths=["inference_architecture"],
    pip_requirements="requirements-serving.txt",
)

# Hugging Face: a pipeline includes the tokenizer; use saved generation options.
hf_info = transformers.log_model(
    transformers_model=loaded_generation_pipeline,
    input_example=["Explain what a model signature describes."],
    name="text-generator",
    triton_model_name="text_generator",
    device="gpu",
    task="text-generation",
    model_config={"max_new_tokens": 32, "do_sample": False},
)

print(hf_info.model_uri)

Logging has already generated and saved the Triton repository. If you need a local directory to inspect or hand off to deployment tooling, optionally retrieve it:

from mlflow_triton import triton

repository = triton.load_model(hf_info.model_uri)
print(repository / "config.pbtxt")
print(repository / "1/model.py")

Return value and validation

All three calls return MLflow ModelInfo. Use info.model_uri to load the saved model or resolve its Triton repository, and info.registered_model_version when registration was requested. triton.load_model(info.model_uri) returns a local pathlib.Path to the model-name directory, downloading artifacts as necessary. This optional call only retrieves the generated repository; it does not load model weights, regenerate the repository or load it into a running Triton server.

Before logging/registration, the wrapper executes the example prediction using MLflow's prediction adapter around the supplied model, derives/checks the signature, saves the native model and validates the input/output tensor mapping. It reuses the existing weights. Unsupported schemas, invalid devices or failed predictions raise an exception. A successful default return does not establish that the saved artifact can be reloaded in another environment.

Use validate_saved_model=True to instead validate a CPU reload of the saved artifact, including whether its custom code can be imported locally. The original model remains in memory during this check. validation.json records the selected mode and validation device; actual Triton serving and GPU execution must be tested separately. There are no deploy, validate_triton, max_batch_size or GPU-index arguments in these wrappers.

Log a Triton model in MLflow

Like NVIDIA's Triton MLflow plugin, this package provides a triton model flavor for a complete Triton model directory. The directory must contain at least one numbered version; config.pbtxt may be included at its root. The MLflow Model contains an MLmodel file with a triton flavor entry and a copy of the entire directory under model/.

from mlflow_triton import triton

model_info = triton.log_model(
    triton_model_path="/tmp/model-repository/iris_classifier",
    name="triton-model",
    registered_model_name="iris-serving",
)
print(model_info.registered_model_version)

triton.load_model("models:/iris-serving/1") downloads the registered model and returns the local Triton model directory. Omit registered_model_name to save a model without creating a Registry version. This API is in the plugin package, not in MLflow's mlflow.triton namespace.

Arguments for an existing Triton model directory

This separate API packages an existing Triton model directory. It does not create model.py, infer a signature or change its configuration.

triton.log_model(
    triton_model_path,
    name="triton-model",
    registered_model_name=None,
    await_registration_for=300,
    run_id=None,
    metadata=None,
)

Returns MLflow ModelInfo. This API has no input_example, signature, device, tags, params or arbitrary **kwargs parameters. Configure serving directly in the supplied Triton directory before logging it.

Examples and optional deployment

See shared examples and validation and the repository overview. CPU/GPU validation is performed by the example runner, not by starting a server inside log_model().

For optional Triton Control deployment, install mlflow-triton-control. A complete Triton model repository is the handoff to any chosen deployment tooling.

Packaging and concurrent loading

Packaging checks the target name and device before saving or predicting. It builds in a temporary sibling of the destination and renames the completed directory into place, cleaning up staging on failure. The native artifact and independent serving copy remain separate; intermediate full-package copies are avoided. MLflow may still copy/upload the completed artifact during logging. TMPDIR selects the outer logging workspace; choose a disk with enough room for both model copies and MLflow's own temporary storage.

Calls through the runtime's load_predictor are serialized while MLflow device and offline environment settings are temporarily applied, including restoration after failure. This lock does not protect unrelated code that modifies os.environ, direct MLflow loads, or separately imported runtime copies. Coordinate those callers or use separate processes. Prediction requests are not serialized by this lock.

MLflow mapping in Triton configuration

Generated repositories store the MLflow-specific mapping as JSON in the parameters.mlflow_contract.string_value field of config.pbtxt. It contains a mapping version, framework flavor, input/output structure (columns, single tensor or named tensors), original field names and text/binary semantics. Tensor names link mapping entries to the standard input and output configuration fields; datatypes and shapes are read exclusively from those fields.

Triton passes the configuration to model.py.initialize() as args["model_config"]. The runtime validates and reconstructs its tensor mapping from it. Device placement still comes from Triton's assigned instance kind/device ID. No separate contract.json is generated or read by the new runtime. Unsupported mapping versions, missing tensor mappings and nonzero max_batch_size fail initialization.

Existing logged artifacts retain their original model.py and contract.json. Re-log a model to produce the new layout; do not mix old generated runtime files with new configuration. Editing a mapping or tensor schema must remain compatible with the saved model's prediction interface and requires renewed validation.

Ready-to-send inference example

Each framework wrapper generates one request.json beside config.pbtxt, using values from input_example, concrete tensor shapes and all configured output names. See the central request instructions for curl commands. No per-model shell script or instruction file is generated. Logging does not send a request.

For NaN/infinity or binary byte inputs, the example instead consists of request.bin (complete HTTP body, including its JSON prefix) and request.headers (the HTTP content type and JSON-prefix byte length). This uses Triton's binary tensor extension without changing input values. Outputs are requested as binary in that case.

Example files contain the supplied input values, just like MLflow's stored input example; use non-sensitive examples. The direct triton.log_model() API for pre-existing repositories does not generate requests because it takes no input example.

Architecture information without loading weights

PyTorch and Transformers wrappers automatically include architecture.json in the MLflow artifact and its self-contained Triton 1/artifact/ copy. Read this JSON without importing custom model code or loading weights. It records the model class, module classes, parameter/buffer shapes and dtypes (never tensor values), and relative paths to packaged Python and configuration files. Transformers reports component classes such as the tokenizer; its saved config.json contains model configuration.

mlflow_triton.pytorch.log_model(
    model,
    input_example=example,
    code_paths=["custom_network.py"],
    architecture_init_kwargs={"input_size": 16, "hidden_size": 32},
)

The same optional argument is supported by mlflow_triton.transformers.log_model. Include only non-sensitive architecture settings, never credentials or training data. This is documentation, not an automatic reconstruction recipe: module structure cannot describe arbitrary forward() logic or preprocessing. Constructor settings are caller-declared and are not verified against the model. Required custom code must still be packaged with the native flavor's code options. Existing artifacts are unchanged; the report is generated on subsequent logging. Serving inputs, outputs, weights and config.pbtxt are unaffected.