diff --git a/alcf_ai/src/alcf_ai/cli.py b/alcf_ai/src/alcf_ai/cli.py index 785b6317..059af6cf 100644 --- a/alcf_ai/src/alcf_ai/cli.py +++ b/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") diff --git a/alcf_ai/src/alcf_ai/client.py b/alcf_ai/src/alcf_ai/client.py index bd85d484..c0080e9a 100644 --- a/alcf_ai/src/alcf_ai/client.py +++ b/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() diff --git a/alcf_ai/src/alcf_ai/d3_triton.py b/alcf_ai/src/alcf_ai/d3_triton.py new file mode 100644 index 00000000..eebb72e6 --- /dev/null +++ b/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}") diff --git a/alcf_ai/src/alcf_ai/resources/__init__.py b/alcf_ai/src/alcf_ai/resources/__init__.py index 2e097596..fe9d9059 100644 --- a/alcf_ai/src/alcf_ai/resources/__init__.py +++ b/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"] diff --git a/alcf_ai/src/alcf_ai/resources/d3_triton.py b/alcf_ai/src/alcf_ai/resources/d3_triton.py new file mode 100644 index 00000000..43de5aaf --- /dev/null +++ b/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=}") diff --git a/resource_server_async/schemas/__init__.py b/resource_server_async/schemas/__init__.py index c0d4f0f8..064dbf15 100644 --- a/resource_server_async/schemas/__init__.py +++ b/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", diff --git a/resource_server_async/schemas/d3_triton.py b/resource_server_async/schemas/d3_triton.py new file mode 100644 index 00000000..052ec6d5 --- /dev/null +++ b/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 diff --git a/resource_server_async/views/__init__.py b/resource_server_async/views/__init__.py index 6e49c015..b9b9c6c8 100644 --- a/resource_server_async/views/__init__.py +++ b/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"]) diff --git a/resource_server_async/views/d3_triton.py b/resource_server_async/views/d3_triton.py new file mode 100644 index 00000000..033ca313 --- /dev/null +++ b/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)