mlflow-triton-control 0.2.0__tar.gz

mlflow_triton_control-0.2.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_control-0.2.0/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.4
+Name: mlflow-triton-control
+Version: 0.2.0
+Summary: MLflow deployment target for Triton Control
+Classifier: License :: OSI Approved :: Apache Software License
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlflow<4,>=3
+Requires-Dist: httpx<1,>=0.27
+Dynamic: license-file
+
+# mlflow-triton-control
+
+MLflow deployment target for models that already exist in a Triton S3 model
+repository. Install this Python package in the environment running the MLflow
+CLI or Python API. Triton Control creates the Kubernetes deployment and resolves
+the selected user-owned S3 profile.
+
+```bash
+python -m build
+python -m pip install dist/mlflow_triton_control-*.whl
+mlflow deployments help -t triton-control
+```
+
+## How MLflow finds the plugin
+
+`mlflow deployments create` is a command provided by MLflow. This package
+registers the deployment target through its Python package metadata:
+
+```toml
+[project.entry-points."mlflow.deployments"]
+triton-control = "mlflow_triton_control.deployment_client"
+```
+
+For `--target triton-control://triton-control:8000`, MLflow looks up the
+`triton-control` entry point among installed packages, loads this module, and
+calls `TritonControlDeploymentClient.create_deployment()`. That class inherits
+from MLflow's `BaseDeploymentClient` and sends the deployment request to the
+Triton Control API. MLflow reads entry-point metadata; it does not search
+through every package's source files. The package must be installed in the
+same Python environment as the `mlflow` CLI.
+
+## 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_control 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.
+
+The current deployment client still accepts an already-published `s3://` Triton
+repository URI. Deploying directly from `models:/...` and copying that version
+to the selected S3 profile are separate planned changes. The current Iris
+workflow therefore continues to use the S3 URI.
+
+## Argo Workflows
+
+Submit the workflow through Triton Control's authenticated Argo proxy. Mark the
+deployment template with these Workflow annotations:
+
+```yaml
+metadata:
+  name: train-iris
+  annotations:
+    triton-control.ai/mlflow-deploy-template: deploy
+    triton-control.ai/mlflow-deployment-name: iris-classifier
+    triton-control.ai/mlflow-s3-profile-name: workflow-training
+```
+
+The backend resolves the profile name and its linked Argo artifact repository
+for the submitting user. It sets `spec.artifactRepositoryRef` and injects
+`TRITON_CONTROL_TOKEN`, `TRITON_CONTROL_S3_PROFILE_ID`, and
+`TRITON_CONTROL_S3_BUCKET` into the named `script` or `container` template.
+Use the injected ID with `-C s3_profile_id="$TRITON_CONTROL_S3_PROFILE_ID"`.
+The token is valid for 60 minutes and limited
+to the specified deployment and S3 profile. The Workflow receives a five-minute
+completion TTL unless it already defines one, and the Secret receives a
+Workflow owner reference for garbage collection.
+
+The opted-in Workflow must have a fixed `metadata.name` and run in the Triton
+Control namespace. Direct submissions to Argo bypass this token injection.

mlflow_triton_control-0.2.0/README.md

@@ -0,0 +1,86 @@
+# mlflow-triton-control
+
+MLflow deployment target for models that already exist in a Triton S3 model
+repository. Install this Python package in the environment running the MLflow
+CLI or Python API. Triton Control creates the Kubernetes deployment and resolves
+the selected user-owned S3 profile.
+
+```bash
+python -m build
+python -m pip install dist/mlflow_triton_control-*.whl
+mlflow deployments help -t triton-control
+```
+
+## How MLflow finds the plugin
+
+`mlflow deployments create` is a command provided by MLflow. This package
+registers the deployment target through its Python package metadata:
+
+```toml
+[project.entry-points."mlflow.deployments"]
+triton-control = "mlflow_triton_control.deployment_client"
+```
+
+For `--target triton-control://triton-control:8000`, MLflow looks up the
+`triton-control` entry point among installed packages, loads this module, and
+calls `TritonControlDeploymentClient.create_deployment()`. That class inherits
+from MLflow's `BaseDeploymentClient` and sends the deployment request to the
+Triton Control API. MLflow reads entry-point metadata; it does not search
+through every package's source files. The package must be installed in the
+same Python environment as the `mlflow` CLI.
+
+## 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_control 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.
+
+The current deployment client still accepts an already-published `s3://` Triton
+repository URI. Deploying directly from `models:/...` and copying that version
+to the selected S3 profile are separate planned changes. The current Iris
+workflow therefore continues to use the S3 URI.
+
+## Argo Workflows
+
+Submit the workflow through Triton Control's authenticated Argo proxy. Mark the
+deployment template with these Workflow annotations:
+
+```yaml
+metadata:
+  name: train-iris
+  annotations:
+    triton-control.ai/mlflow-deploy-template: deploy
+    triton-control.ai/mlflow-deployment-name: iris-classifier
+    triton-control.ai/mlflow-s3-profile-name: workflow-training
+```
+
+The backend resolves the profile name and its linked Argo artifact repository
+for the submitting user. It sets `spec.artifactRepositoryRef` and injects
+`TRITON_CONTROL_TOKEN`, `TRITON_CONTROL_S3_PROFILE_ID`, and
+`TRITON_CONTROL_S3_BUCKET` into the named `script` or `container` template.
+Use the injected ID with `-C s3_profile_id="$TRITON_CONTROL_S3_PROFILE_ID"`.
+The token is valid for 60 minutes and limited
+to the specified deployment and S3 profile. The Workflow receives a five-minute
+completion TTL unless it already defines one, and the Secret receives a
+Workflow owner reference for garbage collection.
+
+The opted-in Workflow must have a fixed `metadata.name` and run in the Triton
+Control namespace. Direct submissions to Argo bypass this token injection.

mlflow_triton_control-0.2.0/mlflow_triton_control/__init__.py

@@ -0,0 +1,5 @@
+"""MLflow deployment target for Triton Control."""
+
+from .deployment_client import TritonControlDeploymentClient
+
+__all__ = ["TritonControlDeploymentClient"]

mlflow_triton_control-0.2.0/mlflow_triton_control/deployment_client.py

@@ -0,0 +1,126 @@
+"""MLflow deployment client backed by Triton Control's deployment API."""
+
+from __future__ import annotations
+
+import os
+import re
+import time
+from urllib.parse import quote, urlsplit
+
+import httpx
+from mlflow.deployments import BaseDeploymentClient
+from mlflow.exceptions import MlflowException
+
+_NAME = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\Z")
+_MODEL = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]*\Z")
+
+
+def target_help() -> str:
+    return (
+        "Target: triton-control://<triton-control-host>[:port]. "
+        "Install this package in the MLflow CLI environment. "
+        "Use -m s3://<profile-bucket>/<repository-prefix>/<model-name> and "
+        "-C s3_profile_id=<owned-profile-id> -C image=<triton-image>. "
+        "The model must already exist in Triton repository format. "
+        "Set TRITON_CONTROL_TOKEN to a user-bound API token. "
+        "MLflow models:/ and runs:/ artifacts, update, predict and local serving are unsupported."
+    )
+
+
+def run_local(target, name, model_uri, flavor=None, config=None):
+    raise MlflowException("Local Triton serving is not supported by triton-control")
+
+
+class TritonControlDeploymentClient(BaseDeploymentClient):
+    """Create and manage user-owned Triton Control deployments."""
+
+    def __init__(self, target_uri):
+        super().__init__(target_uri)
+        parsed = urlsplit(target_uri)
+        if parsed.scheme != "triton-control" or not parsed.netloc or parsed.path not in {"", "/"}:
+            raise MlflowException("Target URI must be triton-control://<host>[:port]")
+        self._base_url = f"http://{parsed.netloc}"
+
+    def _request(self, method, path, **kwargs):
+        token = os.getenv("TRITON_CONTROL_TOKEN", "").strip()
+        if not token:
+            raise MlflowException("TRITON_CONTROL_TOKEN is required")
+        try:
+            with httpx.Client(base_url=self._base_url, timeout=10.0, follow_redirects=False) as client:
+                response = client.request(method, path, headers={"Authorization": f"Bearer {token}"}, **kwargs)
+            response.raise_for_status()
+            return response.json()
+        except httpx.HTTPStatusError as exc:
+            try:
+                detail = exc.response.json().get("detail", "Request failed")
+            except (ValueError, AttributeError):
+                detail = "Request failed"
+            raise MlflowException(f"Triton Control returned HTTP {exc.response.status_code}: {detail}") from exc
+        except httpx.RequestError as exc:
+            raise MlflowException(f"Triton Control is unavailable: {type(exc).__name__}") from exc
+
+    def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=None):
+        if endpoint is not None:
+            raise MlflowException("endpoint is not supported")
+        if flavor not in {None, "triton"}:
+            raise MlflowException("Only the triton flavor is supported")
+        if not _NAME.fullmatch(name):
+            raise MlflowException("name must be a lowercase Kubernetes deployment name")
+        config = config or {}
+        try:
+            profile_id = int(config["s3_profile_id"])
+            image = str(config["image"]).strip()
+        except (KeyError, TypeError, ValueError) as exc:
+            raise MlflowException("s3_profile_id and image are required") from exc
+        if profile_id <= 0 or not image:
+            raise MlflowException("s3_profile_id must be positive and image must not be empty")
+        try:
+            timeout = float(config.get("wait_timeout_seconds", 300))
+        except (TypeError, ValueError) as exc:
+            raise MlflowException("wait_timeout_seconds must be a positive number") from exc
+        if timeout <= 0:
+            raise MlflowException("wait_timeout_seconds must be a positive number")
+        uri = urlsplit(model_uri)
+        parts = uri.path.strip("/").split("/")
+        if (uri.scheme != "s3" or not uri.netloc or uri.query or uri.fragment
+                or len(parts) < 2 or any(part in {"", ".", ".."} for part in parts)
+                or not _MODEL.fullmatch(parts[-1])):
+            raise MlflowException("model_uri must be s3://<bucket>/<repository-prefix>/<model-name>")
+        prefix, model_name = "/".join(parts[:-1]), parts[-1]
+        payload = {
+            "deployment_name": name,
+            "s3_profile_id": profile_id,
+            "model_uri": model_uri,
+            "repository_prefix": prefix,
+            "model_name": model_name,
+            "image": image,
+            "model_control_mode": "explicit",
+        }
+        result = self._request("POST", "/api/deployments", json=payload)
+        deadline = time.monotonic() + timeout
+        while True:
+            current = self.get_deployment(name)
+            if current["status"] == "ready":
+                return current
+            if time.monotonic() >= deadline:
+                raise MlflowException(
+                    f"Deployment {name} (instance {result.get('instance_id')}) did not become ready: "
+                    f"{current.get('status_detail', current['status'])}"
+                )
+            time.sleep(min(3.0, max(0.0, deadline - time.monotonic())))
+
+    def get_deployment(self, name):
+        deployment = self._request("GET", f"/api/deployments/mlflow/{quote(name, safe='')}")
+        return {**deployment, "flavor": "triton"}
+
+    def list_deployments(self):
+        return self._request("GET", "/api/deployments/mlflow")
+
+    def delete_deployment(self, name):
+        return self._request("DELETE", f"/api/deployments/mlflow/{quote(name, safe='')}")
+
+    def update_deployment(self, name, model_uri=None, flavor=None, config=None, endpoint=None):
+        raise MlflowException("Updating deployments is not supported")
+
+    def predict(self, deployment_name=None, inputs=None, endpoint=None):
+        raise MlflowException("Prediction is not supported")

mlflow_triton_control-0.2.0/mlflow_triton_control/triton.py

@@ -0,0 +1,92 @@
+"""MLflow flavor for a complete Triton model repository directory."""
+
+from __future__ import annotations
+
+import shutil
+import sys
+from pathlib import Path
+
+from mlflow.artifacts import download_artifacts
+from mlflow.exceptions import MlflowException
+from mlflow.models import Model
+
+FLAVOR_NAME = "triton"
+_MODEL_DATA = "model"
+
+
+def _validate_repository(model_path: Path) -> None:
+    if not model_path.is_dir():
+        raise MlflowException(f"Triton model directory does not exist: {model_path}")
+    if model_path.is_symlink() or any(path.is_symlink() for path in model_path.rglob("*")):
+        raise MlflowException("Triton model directory must not contain symbolic links")
+    versions = [
+        path for path in model_path.iterdir()
+        if path.is_dir() and path.name.isascii() and path.name.isdecimal()
+        and int(path.name) > 0
+    ]
+    if not versions:
+        raise MlflowException(
+            "Triton model directory must contain a numbered version"
+        )
+
+
+def save_model(triton_model_path, path, mlflow_model=None):
+    """Save one Triton model directory as an MLflow Model with the triton flavor.
+
+    ``triton_model_path`` points at the model-name directory containing
+    ``config.pbtxt`` (if needed) and numbered version directories.
+    """
+    source = Path(triton_model_path).expanduser()
+    _validate_repository(source)
+    source = source.resolve()
+    destination = Path(path).expanduser().resolve()
+    if destination.exists():
+        raise MlflowException(f"MLflow model path already exists: {destination}")
+    if destination == source or destination.is_relative_to(source):
+        raise MlflowException("MLflow model path must be outside the Triton model directory")
+
+    destination.mkdir(parents=True)
+    shutil.copytree(source, destination / _MODEL_DATA)
+    model = mlflow_model or Model()
+    model.add_flavor(FLAVOR_NAME, data=_MODEL_DATA, model_name=source.name)
+    model.save(str(destination / "MLmodel"))
+
+
+def log_model(
+    triton_model_path,
+    name="triton-model",
+    registered_model_name=None,
+    await_registration_for=300,
+    run_id=None,
+    metadata=None,
+):
+    """Log a Triton repository and optionally register a model version.
+
+    Returns MLflow's ``ModelInfo``. A registry version is created only when
+    ``registered_model_name`` is given.
+    """
+    return Model.log(
+        artifact_path=None,
+        flavor=sys.modules[__name__],
+        name=name,
+        registered_model_name=registered_model_name,
+        await_registration_for=await_registration_for,
+        run_id=run_id,
+        metadata=metadata,
+        flavor_name=FLAVOR_NAME,
+        triton_model_path=triton_model_path,
+    )
+
+
+def load_model(model_uri):
+    """Download an MLflow Triton model and return its repository directory."""
+    model_root = Path(download_artifacts(artifact_uri=model_uri)).resolve()
+    model = Model.load(str(model_root / "MLmodel"))
+    flavor = model.flavors.get(FLAVOR_NAME)
+    if not flavor:
+        raise MlflowException(f"Model has no {FLAVOR_NAME} flavor: {model_uri}")
+    data_path = (model_root / flavor.get("data", "")).resolve()
+    if not data_path.is_relative_to(model_root) or data_path == model_root:
+        raise MlflowException("Triton flavor contains an invalid data path")
+    _validate_repository(data_path)
+    return data_path

mlflow_triton_control-0.2.0/mlflow_triton_control.egg-info/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.4
+Name: mlflow-triton-control
+Version: 0.2.0
+Summary: MLflow deployment target for Triton Control
+Classifier: License :: OSI Approved :: Apache Software License
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlflow<4,>=3
+Requires-Dist: httpx<1,>=0.27
+Dynamic: license-file
+
+# mlflow-triton-control
+
+MLflow deployment target for models that already exist in a Triton S3 model
+repository. Install this Python package in the environment running the MLflow
+CLI or Python API. Triton Control creates the Kubernetes deployment and resolves
+the selected user-owned S3 profile.
+
+```bash
+python -m build
+python -m pip install dist/mlflow_triton_control-*.whl
+mlflow deployments help -t triton-control
+```
+
+## How MLflow finds the plugin
+
+`mlflow deployments create` is a command provided by MLflow. This package
+registers the deployment target through its Python package metadata:
+
+```toml
+[project.entry-points."mlflow.deployments"]
+triton-control = "mlflow_triton_control.deployment_client"
+```
+
+For `--target triton-control://triton-control:8000`, MLflow looks up the
+`triton-control` entry point among installed packages, loads this module, and
+calls `TritonControlDeploymentClient.create_deployment()`. That class inherits
+from MLflow's `BaseDeploymentClient` and sends the deployment request to the
+Triton Control API. MLflow reads entry-point metadata; it does not search
+through every package's source files. The package must be installed in the
+same Python environment as the `mlflow` CLI.
+
+## 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_control 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.
+
+The current deployment client still accepts an already-published `s3://` Triton
+repository URI. Deploying directly from `models:/...` and copying that version
+to the selected S3 profile are separate planned changes. The current Iris
+workflow therefore continues to use the S3 URI.
+
+## Argo Workflows
+
+Submit the workflow through Triton Control's authenticated Argo proxy. Mark the
+deployment template with these Workflow annotations:
+
+```yaml
+metadata:
+  name: train-iris
+  annotations:
+    triton-control.ai/mlflow-deploy-template: deploy
+    triton-control.ai/mlflow-deployment-name: iris-classifier
+    triton-control.ai/mlflow-s3-profile-name: workflow-training
+```
+
+The backend resolves the profile name and its linked Argo artifact repository
+for the submitting user. It sets `spec.artifactRepositoryRef` and injects
+`TRITON_CONTROL_TOKEN`, `TRITON_CONTROL_S3_PROFILE_ID`, and
+`TRITON_CONTROL_S3_BUCKET` into the named `script` or `container` template.
+Use the injected ID with `-C s3_profile_id="$TRITON_CONTROL_S3_PROFILE_ID"`.
+The token is valid for 60 minutes and limited
+to the specified deployment and S3 profile. The Workflow receives a five-minute
+completion TTL unless it already defines one, and the Secret receives a
+Workflow owner reference for garbage collection.
+
+The opted-in Workflow must have a fixed `metadata.name` and run in the Triton
+Control namespace. Direct submissions to Argo bypass this token injection.

mlflow_triton_control-0.2.0/mlflow_triton_control.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+LICENSE
+README.md
+pyproject.toml
+mlflow_triton_control/__init__.py
+mlflow_triton_control/deployment_client.py
+mlflow_triton_control/triton.py
+mlflow_triton_control.egg-info/PKG-INFO
+mlflow_triton_control.egg-info/SOURCES.txt
+mlflow_triton_control.egg-info/dependency_links.txt
+mlflow_triton_control.egg-info/entry_points.txt
+mlflow_triton_control.egg-info/requires.txt
+mlflow_triton_control.egg-info/top_level.txt
+tests/test_deployment_client.py
+tests/test_triton_flavor.py

mlflow_triton_control-0.2.0/mlflow_triton_control.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

mlflow_triton_control-0.2.0/mlflow_triton_control.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[mlflow.deployments]
+triton-control = mlflow_triton_control.deployment_client

mlflow_triton_control-0.2.0/mlflow_triton_control.egg-info/requires.txt

@@ -0,0 +1,2 @@
+mlflow<4,>=3
+httpx<1,>=0.27

mlflow_triton_control-0.2.0/mlflow_triton_control.egg-info/top_level.txt

@@ -0,0 +1 @@
+mlflow_triton_control

mlflow_triton_control-0.2.0/pyproject.toml

@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mlflow-triton-control"
+version = "0.2.0"
+description = "MLflow deployment target for Triton Control"
+readme = "README.md"
+requires-python = ">=3.10"
+license-files = ["LICENSE"]
+classifiers = ["License :: OSI Approved :: Apache Software License"]
+dependencies = ["mlflow>=3,<4", "httpx>=0.27,<1"]
+
+[project.entry-points."mlflow.deployments"]
+triton-control = "mlflow_triton_control.deployment_client"
+
+[tool.setuptools.packages.find]
+include = ["mlflow_triton_control*"]

mlflow_triton_control-0.2.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mlflow_triton_control-0.2.0/tests/test_deployment_client.py

@@ -0,0 +1,71 @@
+"""Contract checks for the package without requiring an MLflow server."""
+
+import os
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+mlflow = types.ModuleType("mlflow")
+deployments = types.ModuleType("mlflow.deployments")
+exceptions = types.ModuleType("mlflow.exceptions")
+
+
+class BaseDeploymentClient:
+    def __init__(self, target_uri):
+        self.target_uri = target_uri
+
+
+class MlflowException(Exception):
+    pass
+
+
+deployments.BaseDeploymentClient = BaseDeploymentClient
+exceptions.MlflowException = MlflowException
+with patch.dict(sys.modules, {"mlflow": mlflow, "mlflow.deployments": deployments, "mlflow.exceptions": exceptions}):
+    from mlflow_triton_control.deployment_client import TritonControlDeploymentClient
+
+
+class DeploymentClientTests(unittest.TestCase):
+    def setUp(self) -> None:
+        self.client = TritonControlDeploymentClient("triton-control://control.test")
+
+    def test_create_derives_repository_from_model_uri_and_waits_for_model(self) -> None:
+        with (
+            patch.object(self.client, "_request") as request,
+            patch.dict(os.environ, {"TRITON_CONTROL_TOKEN": "token"}),
+        ):
+            request.side_effect = [
+                {"instance_id": 42},
+                {"name": "iris", "status": "starting"},
+                {"name": "iris", "status": "ready"},
+            ]
+            with patch("time.sleep"):
+                result = self.client.create_deployment(
+                    "iris", "s3://models/dev/iris_classifier",
+                    config={"s3_profile_id": 7, "image": "triton:latest"},
+                )
+        self.assertEqual(result["status"], "ready")
+        self.assertEqual(result["flavor"], "triton")
+        payload = request.call_args_list[0].kwargs["json"]
+        self.assertEqual(payload["repository_prefix"], "dev")
+        self.assertEqual(payload["model_name"], "iris_classifier")
+
+    def test_mlflow_registry_uri_is_rejected(self) -> None:
+        with self.assertRaises(MlflowException):
+            self.client.create_deployment(
+                "iris", "models:/iris/1", config={"s3_profile_id": 7, "image": "triton:latest"}
+            )
+
+    def test_invalid_timeout_is_rejected_before_create(self) -> None:
+        self.client._request = Mock()
+        with self.assertRaisesRegex(MlflowException, "wait_timeout_seconds must be a positive number"):
+            self.client.create_deployment(
+                "iris",
+                "s3://models/repository/iris",
+                config={"s3_profile_id": 7, "image": "triton:latest", "wait_timeout_seconds": "later"},
+            )
+        self.client._request.assert_not_called()

mlflow_triton_control-0.2.0/tests/test_triton_flavor.py

@@ -0,0 +1,111 @@
+"""Checks for packaging and retrieving a complete Triton repository."""
+
+import importlib
+import json
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+
+class MlflowException(Exception):
+    pass
+
+
+class Model:
+    log = Mock()
+
+    def __init__(self):
+        self.flavors = {}
+
+    def add_flavor(self, name, **params):
+        self.flavors[name] = params
+        return self
+
+    def save(self, path):
+        Path(path).write_text(json.dumps({"flavors": self.flavors}))
+
+    @classmethod
+    def load(cls, path):
+        result = cls()
+        result.flavors = json.loads(Path(path).read_text())["flavors"]
+        return result
+
+
+mlflow = types.ModuleType("mlflow")
+artifacts = types.ModuleType("mlflow.artifacts")
+exceptions = types.ModuleType("mlflow.exceptions")
+models = types.ModuleType("mlflow.models")
+deployments = types.ModuleType("mlflow.deployments")
+artifacts.download_artifacts = Mock()
+exceptions.MlflowException = MlflowException
+models.Model = Model
+deployments.BaseDeploymentClient = type("BaseDeploymentClient", (), {})
+with patch.dict(sys.modules, {
+    "mlflow": mlflow, "mlflow.artifacts": artifacts,
+    "mlflow.exceptions": exceptions, "mlflow.models": models,
+    "mlflow.deployments": deployments,
+}):
+    triton = importlib.import_module("mlflow_triton_control.triton")
+
+
+class TritonFlavorTests(unittest.TestCase):
+    def setUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp.cleanup)
+        self.root = Path(self.temp.name)
+        self.source = self.root / "iris_classifier"
+        (self.source / "1").mkdir(parents=True)
+        (self.source / "1" / "model.onnx").write_bytes(b"onnx model")
+        (self.source / "config.pbtxt").write_text('name: "iris_classifier"')
+
+    def test_save_and_load_preserve_repository_and_flavor(self):
+        destination = self.root / "saved"
+        triton.save_model(str(self.source) + "/", destination)
+
+        self.assertEqual((destination / "model/1/model.onnx").read_bytes(), b"onnx model")
+        self.assertEqual((destination / "model/config.pbtxt").read_text(), 'name: "iris_classifier"')
+        saved = Model.load(destination / "MLmodel")
+        self.assertEqual(saved.flavors["triton"], {"data": "model", "model_name": "iris_classifier"})
+
+        with patch.object(triton, "download_artifacts", return_value=str(destination)):
+            self.assertEqual(triton.load_model("models:/iris-serving/1"), destination / "model")
+
+    def test_log_model_forwards_registry_name_to_mlflow(self):
+        info = object()
+        with (
+            patch.dict(sys.modules, {"mlflow_triton_control.triton": triton}),
+            patch.object(Model, "log", return_value=info) as log,
+        ):
+            result = triton.log_model(self.source, registered_model_name="iris-serving")
+        self.assertIs(result, info)
+        self.assertEqual(log.call_args.kwargs["registered_model_name"], "iris-serving")
+        self.assertEqual(log.call_args.kwargs["flavor_name"], "triton")
+        self.assertIs(log.call_args.kwargs["flavor"], triton)
+        self.assertEqual(log.call_args.kwargs["triton_model_path"], self.source)
+
+    def test_rejects_missing_version_and_symbolic_links(self):
+        (self.source / "1" / "model.onnx").unlink()
+        (self.source / "1").rmdir()
+        with self.assertRaisesRegex(MlflowException, "numbered version"):
+            triton.save_model(self.source, self.root / "saved")
+        (self.source / "1").mkdir()
+        (self.source / "1" / "model.onnx").symlink_to(self.source / "config.pbtxt")
+        with self.assertRaisesRegex(MlflowException, "symbolic links"):
+            triton.save_model(self.source, self.root / "saved")
+
+    def test_rejects_non_triton_mlflow_model(self):
+        destination = self.root / "saved"
+        destination.mkdir()
+        Model().save(destination / "MLmodel")
+        with patch.object(triton, "download_artifacts", return_value=str(destination)):
+            with self.assertRaisesRegex(MlflowException, "no triton flavor"):
+                triton.load_model("models:/sklearn/1")
+
+
+if __name__ == "__main__":
+    unittest.main()