mlflow-triton 0.4.0__tar.gz

mlflow_triton-0.4.0/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

mlflow_triton-0.4.0/PKG-INFO

@@ -0,0 +1,459 @@
+Metadata-Version: 2.4
+Name: mlflow-triton
+Version: 0.4.0
+Summary: MLflow logging wrappers and portable Triton model repositories
+License-Expression: Apache-2.0
+Project-URL: Homepage, https://github.com/ai-lab-tech/mlflow-plugin
+Project-URL: Documentation, https://github.com/ai-lab-tech/mlflow-plugin/blob/main/packages/mlflow-triton/README.md
+Project-URL: Repository, https://github.com/ai-lab-tech/mlflow-plugin
+Project-URL: Issues, https://github.com/ai-lab-tech/mlflow-plugin/issues
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlflow<4,>=3.14
+Requires-Dist: protobuf>=3.20.3
+Dynamic: license-file
+
+# 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.
+
+```bash
+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`:
+
+```python
+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:
+
+```text
+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):
+
+```python
+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
+
+| Function | Model argument | Required `input_example` |
+| --- | --- | --- |
+| `sklearn.log_model` | `sk_model`: already-fitted estimator or complete preprocessing `Pipeline`. | A non-empty DataFrame or NumPy array accepted by the saved prediction interface. Preserve column names, order and dtypes. |
+| `pytorch.log_model` | `pytorch_model`: instantiated `torch.nn.Module` with trained weights loaded; call `.eval()` first. A `state_dict` or filename alone is insufficient. | A non-empty NumPy input, or named tensor input supported by MLflow's PyTorch prediction interface. Match the model's input dtype and dimensions. |
+| `transformers.log_model` | `transformers_model`: supported Transformers pipeline, or component dictionary such as `{"model": model, "tokenizer": tokenizer}`. Include the task's required tokenizer/processor assets. | A non-empty task input supported by MLflow pyfunc, such as a list of text strings or a DataFrame with the task's input fields. |
+
+`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`.
+
+| Argument | Type / default | Meaning |
+| --- | --- | --- |
+| `name` | `str`, `"model"` | MLflow logged-model name. Also supplies the Triton model name unless `triton_model_name` is set. |
+| `triton_model_name` | `str` or `None`, `None` | Triton directory and `config.pbtxt` model name. Defaults to `name`; must start with an ASCII letter or underscore, followed by ASCII letters, digits, `_` or `-`. |
+| `device` | `"cpu"` or `"gpu"`, `"cpu"` | Target Triton instance kind. sklearn accepts only `"cpu"`; PyTorch and Transformers accept both. Default validation uses the supplied model's current device. This does not move the model, choose a GPU index or start Triton. |
+| `validate_saved_model` | `bool`, `False` | By default, validate predictions and tensor mappings using the supplied model without reloading weights. Set `True` to instead reload the saved artifact on CPU and validate it, requiring memory for another copy. |
+| `signature` | MLflow `ModelSignature` or `None`, `None` | Explicit input **and output** schemas, or inference from the model and example. Explicit schemas must be compatible with the generated tensor mapping. `False` is rejected. |
+| `registered_model_name` | `str` or `None`, `None` | Create a version under this MLflow Registry name. Omit to log without registering. |
+| `await_registration_for` | seconds, `300` | MLflow wait timeout for registration, when registration is requested. |
+| `run_id` | `str` or `None`, `None` | Optional MLflow run association, forwarded to `Model.log`. Without it, MLflow applies its active-run/default behavior. |
+| `metadata` | `dict` or `None`, `None` | Model metadata, passed to native packaging and MLflow logging. |
+| `tags` | `dict` or `None`, `None` | Tags for the MLflow logged model. |
+| `params` | `dict` or `None`, `None` | Parameters recorded with the MLflow logged model. These do not configure inference; use Transformers `model_config` for generation settings. |
+
+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.
+
+| Argument | Applies to | Default / purpose |
+| --- | --- | --- |
+| `code_paths` | All three | Native default `None`. List of local inference-code files/directories to package, e.g. custom architecture modules. Include their importable package structure; do not pass training entry points. |
+| `pip_requirements` | All three | Default `None`: use native framework defaults without model-loading dependency inference. Supply a requirements list or requirements-file path for custom dependencies and a reproducible serving environment. With `validate_saved_model=True`, native dependency inference is also enabled when no explicit environment is supplied. |
+| `extra_pip_requirements` | All three | Native default `None`. Add requirements to native defaults/inference. |
+| `conda_env` | All three | Native default `None`. Explicit environment dictionary or environment-file path. Follow native MLflow restrictions on combining it with pip requirement arguments. |
+| `pyfunc_predict_fn` | sklearn | `"predict"`. Select the estimator method used by the MLflow prediction interface, e.g. `"predict_proba"` when supported by the estimator and MLflow. This also determines the output used to infer the signature. |
+| `serialization_format` | sklearn | Native MLflow default. Choose a format supported by the installed sklearn flavor. |
+| `serialization_format` | PyTorch | Plugin default `"pickle"`. An explicit alternative must be supported by the installed native flavor and pass saved-model validation. The plugin does not automatically export or guarantee graph compatibility. |
+| `pickle_module` | PyTorch | Native default `None`. Optional serialization module for native PyTorch packaging. |
+| `task` | Transformers | Native default `None`. Task identifier such as `"text-classification"` or `"text-generation"`. Supply it when native packaging cannot infer the task, especially with custom models/component dictionaries. |
+| `processor` | Transformers | Native default `None`. Task-specific processor where required; model/tokenizer components otherwise belong in the pipeline or component dictionary. |
+| `model_config` | Transformers | Native default `None`. Saved inference configuration, e.g. `{"max_new_tokens": 32, "do_sample": False}`. Options must be supported by the task and MLflow prediction interface. |
+| `save_pretrained` | Transformers | Native default `True`. Must not be `False`: Triton needs the model and tokenizer assets saved locally. |
+
+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.
+
+```python
+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:
+
+```python
+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/`.
+
+```python
+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.
+
+```python
+triton.log_model(
+    triton_model_path,
+    name="triton-model",
+    registered_model_name=None,
+    await_registration_for=300,
+    run_id=None,
+    metadata=None,
+)
+```
+
+| Argument | Required / default | Meaning |
+| --- | --- | --- |
+| `triton_model_path` | Required | Path to one model-name directory containing at least one positive numbered version directory. Symbolic links are rejected. Pass the model directory, not the parent repository containing multiple models. |
+| `name` | `"triton-model"` | MLflow logged-model name. The existing Triton directory name and configuration are preserved. |
+| `registered_model_name` | `None` | Optional MLflow Registry name. |
+| `await_registration_for` | `300` | Registration wait timeout in seconds. |
+| `run_id` | `None` | Optional MLflow run association. |
+| `metadata` | `None` | Optional model metadata dictionary. |
+
+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](https://github.com/ai-lab-tech/mlflow-plugin/blob/main/examples/README.md) and the
+[repository overview](https://github.com/ai-lab-tech/mlflow-plugin/blob/main/README.md). 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](https://github.com/ai-lab-tech/mlflow-plugin/blob/main/packages/mlflow-triton-control/README.md). 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](https://github.com/ai-lab-tech/mlflow-plugin#send-an-example-request)
+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](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/protocol/extension_binary_data.html)
+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.
+
+| Argument | Required | Meaning |
+| --- | --- | --- |
+| `architecture_init_kwargs` | No | JSON object documenting constructor arguments, supplied by the caller. Defaults to unknown; never guessed or passed to the constructor. |
+
+```python
+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.