feat: add KServe and Triton inference adapters (#2908)

* feat: add KServe and Triton inference adapters

* fix: allow remote backend without core transformers

* test: support Python 3.10 TOML parsing

---------

Co-authored-by: Maziyar Panahi <5762953+maziyarpanahi@users.noreply.github.com>

Author: maziyarpanahi

CHANGELOG.md

@@ -30,6 +30,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   optional Outlook MSG input.
 - Added committed Android OpenMedKit release-AAR and offline cold-start budgets,
   with blocking Gradle/CI gates and measured values in the Android job summary.
+- Added a Triton ONNX model-repository generator and configuration-selected
+  KServe V2 HTTP/gRPC inference backend with local tokenization and decoding,
+  mocked local/remote span-parity coverage, and no bundled serving runtime.
 - Added opt-in, no-PHI OpenTelemetry spans and aggregate histograms for all ten
   core privacy-pipeline stages, with lazy optional imports, no exporter by
   default, shared `Timer` measurements, synthetic leakage regression tests, and

NOTICE

@@ -22,6 +22,13 @@ License: Apache-2.0
 Usage: Planned structured-data profiling and quasi-identifier discovery for
 tabular anonymization workflows.
 
+Name: KServe V2 inference protocol
+Source: https://github.com/kserve/open-inference-protocol
+License: Apache-2.0
+Usage: OpenMed ships a minimal, locally namespaced protobuf subset for the
+standard ModelInfer request and response wire fields. It provides a client-only
+Triton/KServe integration and does not bundle an inference server.
+
 Name: medspaCy ConText
 Source: https://github.com/medspacy/medspacy
 License: Apache-2.0

docs/brand/system/publication.yml

@@ -106,6 +106,7 @@ classification:
     - serving/tracing.md
     - serving/resilience.md
     - serving/backpressure.md
+    - serving/kserve-triton.md
     - deploy/hardened-image.md
     - offline-install.md
     - serving/async-jobs.md

docs/configuration.md

@@ -51,6 +51,12 @@ export HF_TOKEN=hf_xxx
 export OPENMED_TORCH_DEVICE=cuda:1
 ```
 
+An explicitly selected `backend="remote"` adds KServe V2 endpoint, model,
+protocol, timeout, TLS, and local-tokenizer settings. It is never selected
+automatically and is incompatible with `local_only` or `OPENMED_OFFLINE`.
+See [KServe and Triton model repositories](serving/kserve-triton.md) for the
+full configuration and deployment boundary.
+
 For Indic personal-name pseudonymization, enable
 `transliteration_aware_name_matching` and reuse the same setting when reopening
 a file-backed surrogate vault. The collision threshold and optional local

docs/model-loader.md

@@ -98,6 +98,15 @@ print(encoding["tokens"][:10])
 
 `load_model` lets you access the underlying HF `AutoModel` + `AutoTokenizer` for workflows that outgrow pipelines.
 
+## User-operated KServe and Triton backends
+
+Set `OpenMedConfig.backend="remote"` to preserve the normal
+`ModelLoader.create_pipeline` call while running the ONNX graph through a
+user-operated KServe V2 or Triton endpoint. OpenMed still tokenizes and decodes
+locally, and the remote backend is never part of automatic backend selection.
+See [KServe and Triton model repositories](serving/kserve-triton.md) for the
+repository generator, complete config, and privacy boundary.
+
 ## Releasing cached models
 
 Long-running services can release cached model references when VRAM or RAM should be reclaimed:

docs/serving/grpc.md

@@ -23,9 +23,11 @@ make grpc-proto
 
 ## Contract
 
-The source protobuf contract is committed at
-`openmed/service/proto/openmed.proto`. Generated Python modules are committed
-under `openmed/service/proto/generated/`.
+The OpenMed service contract is committed at
+`openmed/service/proto/openmed.proto`. The client-only KServe V2 wire subset
+used by the [remote inference backend](kserve-triton.md) is committed at
+`openmed/service/proto/kserve_v2.proto`. Generated Python modules for both are
+committed under `openmed/service/proto/generated/`.
 
 The `OpenMedService` exposes:
 

docs/serving/kserve-triton.md

@@ -0,0 +1,187 @@
+# KServe and Triton model repositories
+
+OpenMed can package an exported ONNX token classifier as a Triton model
+repository and use a user-operated Triton or KServe endpoint as its inference
+backend. Tokenization, source offsets, and entity decoding remain in the
+OpenMed process. The server receives numeric tensors and returns `logits`.
+
+OpenMed does not install or start Triton or KServe. Operate the serving system
+separately and apply the authentication, network isolation, and transport
+controls required for your deployment.
+
+## Install client and export dependencies
+
+Repository generation needs ONNX support. Remote HTTP and gRPC inference use
+client libraries in the `triton` extra:
+
+```bash
+pip install 'openmed[onnx,triton]'
+```
+
+The extra contains no KServe or Triton server distribution. HTTP follows the
+KServe V2 JSON protocol. gRPC uses a small, wire-compatible subset of the
+KServe V2 protobuf contract and the standard
+`inference.GRPCInferenceService/ModelInfer` method.
+
+## 1. Export and build the repository
+
+Start with an OpenMed ONNX token-classification artifact produced by
+`openmed.onnx.convert`. For `max_batch_size > 0`, every graph input and output
+must have a dynamic leading batch dimension. The graph must expose
+`input_ids`, `attention_mask`, and a `logits` output.
+
+```python
+from pathlib import Path
+
+from openmed.onnx import convert
+from openmed.service.backends import write_triton_model_repository
+
+artifact_dir = Path("build/openmed-pii-onnx")
+convert(
+    "OpenMed/OpenMed-PII-SuperClinical-Small-44M-v1",
+    artifact_dir,
+)
+
+repository = write_triton_model_repository(
+    artifact_dir / "model.onnx",
+    "build/model-repository",
+    model_name="openmed_pii",
+    version=1,
+    max_batch_size=8,
+)
+print(repository.model_dir)
+```
+
+The adapter derives tensor names, data types, and non-batch dimensions from
+the ONNX graph and renders a deterministic `config.pbtxt` using Triton's
+`onnxruntime` backend:
+
+```text
+build/model-repository/
+└── openmed_pii/
+    ├── config.pbtxt
+    └── 1/
+        └── model.onnx
+```
+
+For a graph with external tensor data, the adapter uses Triton's documented
+multi-file ONNX form and copies every referenced sidecar:
+
+```text
+build/model-repository/openmed_pii/1/
+└── model.onnx/
+    ├── model.onnx
+    └── weights.bin
+```
+
+Absolute, parent-relative, and symlink-escaping external-data paths are
+rejected. Existing numeric model versions are never overwritten, and a new
+version must have the same derived schema as the existing config.
+
+Recheck a copied repository before mounting it:
+
+```python
+from openmed.service.backends import validate_triton_model_repository
+
+validate_triton_model_repository(
+    "build/model-repository",
+    model_name="openmed_pii",
+    version=1,
+)
+```
+
+## 2. Serve it in your infrastructure
+
+Mount `build/model-repository` into your own Triton deployment and point
+`--model-repository` at that mount. For KServe, configure its Triton runtime or
+storage initializer to expose the same repository. Verify the V2 model-ready
+endpoint before routing OpenMed traffic:
+
+```bash
+curl --fail https://triton.example/v2/models/openmed_pii/ready
+```
+
+Server installation, GPU tuning, autoscaling, credentials, and lifecycle
+management remain outside OpenMed.
+
+## 3. Select remote inference through configuration
+
+Keep the original artifact or model repository available to the OpenMed
+client. Its fast tokenizer and `id2label` metadata are used locally; only the
+ONNX graph executes on the remote server.
+
+```python
+from openmed.core import ModelLoader, OpenMedConfig
+
+config = OpenMedConfig(
+    backend="remote",
+    remote_inference_endpoint="https://triton.example",
+    remote_inference_protocol="http",
+    remote_inference_model_name="openmed_pii",
+    remote_inference_model_version="1",
+    remote_inference_tokenizer="build/openmed-pii-onnx",
+    remote_inference_timeout_seconds=30,
+    remote_inference_verify_tls=True,
+)
+
+loader = ModelLoader(config)
+pipeline = loader.create_pipeline(
+    "OpenMed/OpenMed-PII-SuperClinical-Small-44M-v1",
+    task="token-classification",
+    aggregation_strategy="simple",
+)
+entities = pipeline("Synthetic patient Alice Nguyen")
+```
+
+No pipeline call-site change is required. `OpenMedConfig.backend` selects the
+remote implementation through the normal `ModelLoader.create_pipeline` path.
+The pipeline accepts one string or a batch and returns the same entity
+dictionary shape and source spans as local ONNX inference.
+
+The settings can also be stored in the normal flat TOML configuration:
+
+```toml
+backend = "remote"
+remote_inference_endpoint = "https://kserve.example"
+remote_inference_protocol = "http"
+remote_inference_model_name = "openmed_pii"
+remote_inference_model_version = "1"
+remote_inference_tokenizer = "/models/openmed-pii-onnx"
+remote_inference_timeout_seconds = 30.0
+remote_inference_verify_tls = true
+```
+
+For Triton gRPC, use a gRPC target and protocol:
+
+```python
+config = OpenMedConfig(
+    backend="remote",
+    remote_inference_endpoint="grpcs://triton.example:8001",
+    remote_inference_protocol="grpc",
+    remote_inference_model_name="openmed_pii",
+    remote_inference_tokenizer="/models/openmed-pii-onnx",
+)
+```
+
+`grpc://` and bare `host:port` targets are plaintext. `grpcs://` uses TLS and
+does not permit certificate verification to be disabled. HTTP deployments
+should likewise keep `remote_inference_verify_tls=true`.
+
+## Privacy and offline boundary
+
+The adapter does not log source text, request tensors, response bodies, or raw
+server errors. It sends token IDs, attention masks, and tokenizer-required
+token-type IDs rather than the original string. These tensors can still encode
+sensitive content, so treat the endpoint as inside the trusted PHI boundary,
+use transport encryption, and enforce authentication and network policy in
+your serving infrastructure.
+
+Remote inference is explicit only: it is never auto-selected. It is also
+rejected when `local_only=True` or `OPENMED_OFFLINE=1`, because either setting
+promises a no-egress runtime.
+
+Protocol references:
+
+- [Triton model repository layout](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/model_repository.html)
+- [Triton model configuration](https://docs.nvidia.com/deeplearning/triton-inference-server/user-guide/docs/user_guide/model_configuration.html)
+- [KServe V2 inference protocol](https://kserve.github.io/website/docs/concepts/architecture/data-plane/v2-protocol)

mkdocs.yml

@@ -33,6 +33,7 @@ nav:
       - REST Tracing: serving/tracing.md
       - Serving Resilience: serving/resilience.md
       - Serving Backpressure: serving/backpressure.md
+      - KServe & Triton Model Repositories: serving/kserve-triton.md
       - Hardened Service Image: deploy/hardened-image.md
       - Offline & Air-Gapped Installation: offline-install.md
       - Async REST Jobs & Webhooks: serving/async-jobs.md

openmed/core/backends.py

@@ -230,12 +230,46 @@ def create_pipeline(
         return OnnxTokenClassificationPipeline(model)
 
 
+class RemoteInferenceBackend:
+    """Backend using a user-operated KServe V2 or Triton endpoint."""
+
+    def __init__(self, config: Any = None) -> None:
+        self._config = config
+
+    def is_available(self) -> bool:
+        from openmed.service.backends.remote_inference import (
+            remote_inference_dependencies_available,
+        )
+
+        return remote_inference_dependencies_available(self._config)
+
+    def create_pipeline(
+        self,
+        model_name: str,
+        task: str = "token-classification",
+        aggregation_strategy: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Callable:
+        from openmed.service.backends.remote_inference import (
+            create_remote_inference_pipeline,
+        )
+
+        return create_remote_inference_pipeline(
+            model_name,
+            config=self._config,
+            task=task,
+            aggregation_strategy=aggregation_strategy,
+            **kwargs,
+        )
+
+
 # -- Backend registry and auto-detection ------------------------------------
 
 _BACKENDS: Dict[str, type] = {
     "hf": HuggingFaceBackend,
     "mlx": MLXBackend,
     "onnx": OnnxBackend,
+    "remote": RemoteInferenceBackend,
 }
 
 
@@ -246,7 +280,8 @@ def get_backend(
     """Return the requested backend, or auto-detect the best available one.
 
     Args:
-        name: ``"hf"``, ``"mlx"``, ``"onnx"``, or ``None`` for auto-detect.
+        name: ``"hf"``, ``"mlx"``, ``"onnx"``, ``"remote"``, or ``None``
+            for auto-detect.
         config: OpenMedConfig to pass to the backend.
 
     Auto-detection order:
@@ -271,6 +306,11 @@ def get_backend(
                     "with: pip install 'openmed[onnx-runtime]'. The low_resource "
                     "profile does not fall back to a Torch backend."
                 )
+            if name == "remote":
+                raise RuntimeError(
+                    "Backend 'remote' is not available. Install its client-only "
+                    "dependencies with: pip install 'openmed[triton]'."
+                )
             raise RuntimeError(
                 f"Backend {name!r} is not available. Install its dependencies first."
             )

openmed/core/capabilities.py

@@ -2,7 +2,8 @@
 
 OpenMed ships many optional integrations behind ``pip install openmed[...]``
 extras (``mlx``, ``coreml``, ``onnx``, ``gliner``, ``spacy``, ``presidio``,
-``hf``, ``multimodal``, ``service``, ``mcp``, ...). Historically each seam
+``hf``, ``multimodal``, ``service``, ``triton``, ``mcp``, ...). Historically
+each seam
 guarded its optional imports differently: some raised a bare :class:`ImportError`
 with an ad-hoc message, others raised one of several ``MissingDependencyError``
 variants, and there was no single, importless way to probe which backends are
@@ -197,6 +198,13 @@ def primary_package(self) -> str:
         description="FastAPI REST service surface",
         install="fastapi",
     ),
+    "triton": BackendSpec(
+        name="triton",
+        extra="triton",
+        modules=("grpc", "httpx", "numpy", "transformers", "google"),
+        description="KServe V2 HTTP/gRPC remote inference clients",
+        install="grpcio",
+    ),
     "mcp": BackendSpec(
         name="mcp",
         extra="mcp",