LCF-106: Integrate AmSC D3 Triton Models

alcf_ai/src/alcf_ai/cli.py

@@ -10,6 +10,7 @@
 
 from .auth import cli as auth_cli
 from .client import InferenceClient
+from .d3_triton import cli as d3_triton_cli
 from .sam3 import cli as sam3_cli
 
 logger = logging.getLogger(__name__)
@@ -24,6 +25,7 @@ class CliState(TypedDict, total=False):
 _cli_state: CliState = {}
 
 cli.add_typer(auth_cli, name="auth", help="Login and get access tokens")
+cli.add_typer(d3_triton_cli, name="d3-triton", help="Use the Triton HEP inference service")
 cli.add_typer(sam3_cli, name="sam3", help="Use the SAM3 image segmentation service")
 
 

alcf_ai/src/alcf_ai/client.py

@@ -6,7 +6,7 @@
 from pydantic import BaseModel
 
 from .auth import get_inference_authorizer
-from .resources import ClientResource, ClusterResource, Sam3Resource
+from .resources import ClientResource, ClusterResource, D3TritonResource, Sam3Resource
 from .transfer import TransferResult, https_put_to_collection, run_globus_transfer
 
 DEFAULT_BASE_URL = os.environ.get(
@@ -59,6 +59,12 @@ def sam3(self) -> "Sam3Resource":
             "sam3", Sam3Resource("sophia/sam3service", self)
         )
 
+    @property
+    def d3_triton(self) -> "D3TritonResource":
+        return self._resources.setdefault(  # type: ignore[return-value]
+            "d3_triton", D3TritonResource("sophia/triton/amsc-d3", self)
+        )
+
     def list_endpoints(self) -> dict[str, Any]:
         resp = self.get("list-endpoints")
         resp.raise_for_status()

alcf_ai/src/alcf_ai/d3_triton.py

@@ -0,0 +1,62 @@
+import logging
+from pathlib import Path
+
+import typer
+
+from .auth import STAGING_COLLECTION_ROOT
+
+logger = logging.getLogger(__name__)
+
+cli = typer.Typer(no_args_is_help=True)
+
+
+@cli.command()
+def submit(
+    model_name: str = typer.Argument(..., help="Triton model name (e.g. DoubleMetricLearning)"),
+    input_path: Path = typer.Argument(..., help="Path to the input .npz file"),
+    from_collection_id: str | None = typer.Option(
+        None, help="Globus collection ID to stage in from (default: HTTPS upload)"
+    ),
+    to_collection_id: str | None = typer.Option(
+        None, help="Globus collection ID to stage out results to"
+    ),
+) -> None:
+    """
+    Stage in an .npz input file, run Triton HEP inference, and stage out the result.
+    """
+    from .cli import _cli_state
+
+    client = _cli_state["client"]
+
+    input_path = input_path.expanduser().resolve()
+
+    logger.info(f"Staging in {input_path}")
+    stagein = client.stage_in(
+        input_path, Path(input_path.name), from_collection_id=from_collection_id
+    )
+    logger.info(f"Stage in complete: {stagein}")
+
+    remote_input = STAGING_COLLECTION_ROOT + str(stagein.destination_path)
+    remote_output = remote_input.rsplit(".", 1)[0] + ".output.npz"
+
+    logger.info(f"Submitting inference request for model {model_name!r}...")
+    resp = client.d3_triton.submit(
+        model_name=model_name,
+        input_path=remote_input,
+        output_path=remote_output,
+    )
+
+    logger.info(f"Polling on inference task {resp.task_id!r}...")
+    result = client.d3_triton.poll_task_result(resp.task_id)
+    logger.info(f"Inference completed: {result}")
+
+    if to_collection_id:
+        output_filename = Path(result["output_path"]).name
+        local_output = input_path.with_suffix(".output.npz")
+        logger.info(f"Staging out result file: {output_filename}")
+        stageout = client.stage_out(
+            to_collection_id,
+            Path(output_filename),
+            local_output,
+        )
+        logger.info(f"Stage out complete: {stageout}")

alcf_ai/src/alcf_ai/resources/__init__.py

@@ -1,5 +1,6 @@
 from .cluster import ClusterResource
+from .d3_triton import D3TritonResource
 from .resource import ClientResource
 from .sam3 import Sam3Resource
 
-__all__ = ["ClusterResource", "Sam3Resource", "ClientResource"]
+__all__ = ["ClusterResource", "D3TritonResource", "Sam3Resource", "ClientResource"]

alcf_ai/src/alcf_ai/resources/d3_triton.py

@@ -0,0 +1,84 @@
+import logging
+import time
+from typing import Any
+
+from pydantic import BaseModel
+
+from .resource import ClientResource
+
+logger = logging.getLogger(__name__)
+
+
+class D3TritonRequest(BaseModel):
+    model_name: str
+    input_path: str
+    output_path: str
+    outputs: list[str] | None = None
+
+
+class SubmitTaskResponse(BaseModel):
+    task_id: str
+
+
+class D3TritonResource(ClientResource):
+    class TaskPending(Exception): ...
+
+    def submit(
+        self,
+        model_name: str,
+        input_path: str,
+        output_path: str,
+        outputs: list[str] | None = None,
+    ) -> SubmitTaskResponse:
+        """
+        Submit a Triton HEP inference request. The input_path and output_path
+        are filesystem paths on the compute node (typically under the staging
+        collection root).
+        """
+        payload = D3TritonRequest(
+            model_name=model_name,
+            input_path=input_path,
+            output_path=output_path,
+            outputs=outputs,
+        )
+        resp = self._client.post(
+            f"{self.name}/process", json=payload.model_dump(mode="json")
+        )
+        resp.raise_for_status()
+        return SubmitTaskResponse.model_validate(resp.json())
+
+    def get_task_result(self, task_id: str) -> dict[str, Any]:
+        """
+        Get the result of a submitted inference task. Raises
+        D3TritonResource.TaskPending if the inference has not yet finished.
+
+        Returns the dict produced by the Globus Compute function:
+        ``{"model_name": ..., "output_path": ...}``
+        """
+        resp = self._client.get(f"{self.name}/tasks/{task_id}")
+
+        if resp.status_code == 202 and b"pending" in resp.content:
+            raise D3TritonResource.TaskPending
+        elif resp.status_code >= 400:
+            resp.raise_for_status()
+
+        result: dict[str, Any] = resp.json().get("result")
+        if result is not None:
+            return result
+
+        raise RuntimeError(f"Unexpected D3 Triton inference response: {resp}")
+
+    def poll_task_result(
+        self, task_id: str, timeout: int = 300
+    ) -> dict[str, Any]:
+        """
+        Poll on the inference task for up to ``timeout`` seconds.
+        """
+        start = time.monotonic()
+        logger.info(f"Polling on inference {task_id=}")
+        while time.monotonic() - start < timeout:
+            try:
+                return self.get_task_result(task_id)
+            except D3TritonResource.TaskPending:
+                time.sleep(1)
+        raise TimeoutError(f"{task_id=} not finished in {timeout=}")

resource_server_async/schemas/__init__.py

@@ -1,9 +1,11 @@
 from .clusters import CheckMaintenanceResult, ClusterStatus
+from .d3_triton import D3TritonRequest
 from .data_transfer import GlobusStagingAreaPrepared
 from .endpoints import ClusterSummary, ListEndpointsResponse
 from .sam3 import Sam3Request
 
 __all__ = [
+    "D3TritonRequest",
     "Sam3Request",
     "ListEndpointsResponse",
     "GlobusStagingAreaPrepared",

resource_server_async/schemas/d3_triton.py

@@ -0,0 +1,8 @@
+from ninja import Schema
+
+
+class D3TritonRequest(Schema):
+    model_name: str
+    input_path: str
+    output_path: str
+    outputs: list[str] | None = None

resource_server_async/views/__init__.py

@@ -3,6 +3,7 @@
 from .anthropic import router as anthropic_router
 from .batch import router as batch_router
 from .core import router as core_router
+from .d3_triton import router as d3_triton_router
 from .data import router as data_router
 from .openai import router as openai_router
 from .sam3 import router as sam3_router
@@ -12,6 +13,7 @@
 router.add_router("/", anthropic_router, tags=["anthropic"])
 router.add_router("/", batch_router, tags=["batch"])
 router.add_router("/", core_router, tags=["core"])
+router.add_router("/", d3_triton_router, tags=["d3-triton"])
 router.add_router("/data", data_router, tags=["data"])
 router.add_router("/", openai_router, tags=["openai"])
 router.add_router("/", sam3_router, tags=["sam3"])

resource_server_async/views/d3_triton.py

@@ -0,0 +1,55 @@
+import logging
+
+from ninja import Router
+
+from ..clusters import BaseCluster
+from ..endpoints import BaseEndpoint, GlobusComputeEndpoint
+from ..schemas.auth import AuthedRequest
+from ..schemas.d3_triton import D3TritonRequest
+from ..schemas.endpoints import (
+    SubmitTaskAsyncResponse,
+    SubmitTaskResult,
+)
+
+router = Router()
+log = logging.getLogger(__name__)
+
+
+@router.post("/sophia/triton/amsc-d3/process", response=SubmitTaskAsyncResponse)
+async def d3_triton_infer(
+    request: AuthedRequest, payload: D3TritonRequest
+) -> SubmitTaskAsyncResponse:
+    """
+    Submit a Triton HEP inference request to Globus Compute endpoint.
+    """
+    cluster = await BaseCluster.load_adapter("sophia")
+    (await cluster.check_maintenance()).raise_if_down()
+
+    endpoint = await BaseEndpoint.load_adapter(
+        cluster.cluster_name, "triton", "amsc-d3"
+    )
+    assert isinstance(endpoint, GlobusComputeEndpoint)
+    log.info(f"endpoint_slug: {endpoint.endpoint_slug} - user: {request.auth.username}")
+
+    endpoint.check_permission(request.auth)
+
+    data = payload.model_dump()
+    task_response = await endpoint.submit_task_async(data)
+    return task_response
+
+
+@router.get("/sophia/triton/amsc-d3/tasks/{task_id}", response=SubmitTaskResult)
+async def d3_triton_get_task_result(
+    request: AuthedRequest, task_id: str
+) -> SubmitTaskResult:
+    cluster = await BaseCluster.load_adapter("sophia")
+    (await cluster.check_maintenance()).raise_if_down()
+
+    endpoint = await BaseEndpoint.load_adapter(
+        cluster.cluster_name, "triton", "amsc-d3"
+    )
+    assert isinstance(endpoint, GlobusComputeEndpoint)
+    log.info(f"endpoint_slug: {endpoint.endpoint_slug} - user: {request.auth.username}")
+
+    endpoint.check_permission(request.auth)
+    return await endpoint.get_task_result(task_id)