fxn 0.0.54__tar.gz → 0.0.55__tar.gz

{fxn-0.0.54 → fxn-0.0.55}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fxn
-Version: 0.0.54
+Version: 0.0.55
 Summary: Run prediction functions locally in Python. Register at https://fxn.ai.
 Author-email: "NatML Inc." <hi@fxn.ai>
 License:                                  Apache License

{fxn-0.0.54 → fxn-0.0.55}/fxn/beta/__init__.py

@@ -6,8 +6,9 @@
 from .metadata import (
     CoreMLInferenceMetadata, LiteRTInferenceMetadata, LlamaCppInferenceMetadata,
     OnnxInferenceMetadata, OnnxRuntimeInferenceSessionMetadata, OpenVINOInferenceMetadata,
-    QnnInferenceMetadata, QnnInferenceBackend, QnnInferenceQuantization,
-    # Deprecated
-    ONNXInferenceMetadata, ONNXRuntimeInferenceSessionMetadata
+    QnnInferenceMetadata, QnnInferenceBackend, QnnInferenceQuantization, TensorRTInferenceMetadata    
 )
-from .services import RemoteAcceleration
+from .services import RemoteAcceleration
+
+# DEPRECATED
+from .metadata import ONNXInferenceMetadata, ONNXRuntimeInferenceSessionMetadata

{fxn-0.0.54 → fxn-0.0.55}/fxn/beta/metadata.py

@@ -61,6 +61,11 @@ class CoreMLInferenceMetadata (BaseModel):
         description="Positional inputs to the model.",
         exclude=True
     )
+    output_keys: list[str] | None = Field(
+        default=None,
+        description="Model output dictionary keys. Use this if the model returns a dictionary.",
+        exclude=True
+    )
     model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
 
 class OnnxInferenceMetadata (BaseModel):
@@ -80,6 +85,11 @@ class OnnxInferenceMetadata (BaseModel):
         description="Positional inputs to the model.",
         exclude=True
     )
+    output_keys: list[str] | None = Field(
+        default=None,
+        description="Model output dictionary keys. Use this if the model returns a dictionary.",
+        exclude=True
+    )
     model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
 
 class OnnxRuntimeInferenceSessionMetadata (BaseModel):
@@ -118,6 +128,11 @@ class LiteRTInferenceMetadata (BaseModel):
         description="Positional inputs to the model.",
         exclude=True
     )
+    output_keys: list[str] | None = Field(
+        default=None,
+        description="Model output dictionary keys. Use this if the model returns a dictionary.",
+        exclude=True
+    )
     model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
 
 class OpenVINOInferenceMetadata (BaseModel):
@@ -137,6 +152,11 @@ class OpenVINOInferenceMetadata (BaseModel):
         description="Positional inputs to the model.",
         exclude=True
     )
+    output_keys: list[str] | None = Field(
+        default=None,
+        description="Model output dictionary keys. Use this if the model returns a dictionary.",
+        exclude=True
+    )
     model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
 
 QnnInferenceBackend = Literal["cpu", "gpu", "htp"]
@@ -161,6 +181,11 @@ class QnnInferenceMetadata (BaseModel):
         description="Positional inputs to the model.",
         exclude=True
     )
+    output_keys: list[str] | None = Field(
+        default=None,
+        description="Model output dictionary keys. Use this if the model returns a dictionary.",
+        exclude=True
+    )
     backend: QnnInferenceBackend = Field(
         default="cpu",
         description="QNN backend to execute the model.",
@@ -173,6 +198,51 @@ class QnnInferenceMetadata (BaseModel):
     )
     model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
 
+CudaArchitecture = Literal[
+    "sm_80", "sm_86", "sm_87",  # Ampere (A100)
+    "sm_89",                    # Ada Lovelace (L40)
+    "sm_90",                    # Hopper (H100)
+    "sm_100",                   # Blackwell (B200)
+]
+
+TensorRTPrecision = Literal["fp32", "fp16", "int8", "int4"]
+
+class TensorRTInferenceMetadata (BaseModel):
+    """
+    Metadata required to lower a PyTorch model for inference on Nvidia GPUs with TensorRT.
+
+    Members:
+        model (torch.nn.Module): PyTorch module to apply metadata to.
+        model_args (tuple[Tensor,...]): Positional inputs to the model.
+        cuda_arch (CudaArchitecture): Target CUDA architecture for the TensorRT engine. Defaults to `sm_80` (Ampere).
+        precision (TensorRTPrecision): TensorRT engine inference precision. Defaults to `fp16`.
+    """
+    kind: Literal["meta.inference.tensorrt"] = "meta.inference.tensorrt"
+    model: Annotated[object, BeforeValidator(_validate_torch_module)] = Field(
+        description="PyTorch module to apply metadata to.",
+        exclude=True
+    )
+    model_args: Annotated[list[object], BeforeValidator(_validate_torch_tensor_args)] = Field(
+        description="Positional inputs to the model.",
+        exclude=True
+    )
+    output_keys: list[str] | None = Field(
+        default=None,
+        description="Model output dictionary keys. Use this if the model returns a dictionary.",
+        exclude=True
+    )
+    cuda_arch: CudaArchitecture = Field(
+        default="sm_80",
+        description="Target CUDA architecture for the TensorRT engine.",
+        exclude=True
+    )
+    precision: TensorRTPrecision = Field(
+        default="fp16",
+        description="TensorRT engine inference precision.",
+        exclude=True
+    )
+    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
+
 class LlamaCppInferenceMetadata (BaseModel):
     """
     Metadata required to lower a Llama.cpp model for LLM inference.

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/__init__.py

@@ -9,7 +9,7 @@ from ..logging import TracebackMarkupConsole
 from ..version import __version__
 
 from .auth import app as auth_app
-from .compile import compile_predictor
+from .compile import compile_predictor, triage_predictor
 from .misc import cli_options
 from .predictions import create_prediction
 from .predictors import archive_predictor, delete_predictor, retrieve_predictor
@@ -86,6 +86,14 @@ app.add_typer(
     rich_help_panel="Beta"
 )
 
+# Insiders
+app.command(
+    name="triage",
+    help="Triage a compile error.",
+    rich_help_panel="Insiders",
+    hidden=True
+)(triage_predictor)
+
 # Run
 if __name__ == "__main__":
     app()

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/auth.py

@@ -12,7 +12,7 @@ from ..function import Function
 app = Typer(no_args_is_help=True)
 
 @app.command(name="login", help="Login to Function.")
-def login (
+def login(
     access_key: str=Argument(..., help="Function access key.", envvar="FXN_ACCESS_KEY")
 ):
     fxn = Function(access_key=access_key)
@@ -22,18 +22,18 @@ def login (
     print_json(data=user)
 
 @app.command(name="status", help="Get current authentication status.")
-def auth_status ():
+def auth_status():
     fxn = Function(get_access_key())
     user = fxn.users.retrieve()
     user = user.model_dump() if user else None
     print_json(data=user)
 
 @app.command(name="logout", help="Logout from Function.")
-def logout ():
+def logout():
     _set_access_key(None)
     print("Successfully logged out of Function")
 
-def get_access_key () -> str:
+def get_access_key() -> str:
     """
     Get the CLI access key.
 
@@ -46,7 +46,7 @@ def get_access_key () -> str:
     with open(credentials_path) as f:
         return f.read()
 
-def _set_access_key (key: str):
+def _set_access_key(key: str):
     """
     Set the CLI access key.
 

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/compile.py

@@ -9,9 +9,11 @@ from inspect import getmembers, getmodulename, isfunction
 from pathlib import Path
 from pydantic import BaseModel
 from rich import print as print_rich
+from rich.panel import Panel
 import sys
 from typer import Argument, Option
 from typing import Callable, Literal
+from typing_extensions import Annotated
 from urllib.parse import urlparse, urlunparse
 
 from ..client import FunctionAPIError
@@ -21,16 +23,39 @@ from ..sandbox import EntrypointCommand
 from ..logging import CustomProgress, CustomProgressTask
 from .auth import get_access_key
 
-class CompileError (Exception):
-    pass
-
-def compile_predictor (
+def compile_predictor(
     path: str=Argument(..., help="Predictor path."),
     overwrite: bool=Option(False, "--overwrite", help="Whether to delete any existing predictor with the same tag before compiling."),
 ):
     run_async(_compile_predictor_async(path, overwrite=overwrite))
 
-async def _compile_predictor_async (
+def triage_predictor(
+    reference_code: Annotated[str, Argument(help="Predictor compilation reference code.")]
+):
+    fxn = Function(get_access_key())
+    error = fxn.client.request(
+        method="GET",
+        path=f"/predictors/triage?referenceCode={reference_code}",
+        response_type=_TriagedCompileError
+    )
+    user_panel = Panel(
+        error.user,
+        title="User Error",
+        title_align="left",
+        highlight=True,
+        border_style="bright_red"
+    )
+    internal_panel = Panel(
+        error.internal,
+        title="Internal Error",
+        title_align="left",
+        highlight=True,
+        border_style="gold1"
+    )
+    print_rich(user_panel)
+    print_rich(internal_panel)
+
+async def _compile_predictor_async(
     path: str,
     *,
     overwrite: bool
@@ -78,11 +103,11 @@ async def _compile_predictor_async (
                         task_queue.push_log(event)
                     elif isinstance(event, _ErrorEvent):
                         task_queue.push_error(event)
-                        raise CompileError(event.data.error)
+                        raise _CompileError(event.data.error)
     predictor_url = _compute_predictor_url(fxn.client.api_url, spec.tag)
     print_rich(f"\n[bold spring_green3]🎉 Predictor is now being compiled.[/bold spring_green3] Check it out at [link={predictor_url}]{predictor_url}[/link]")
 
-def _load_predictor_func (path: str) -> Callable[...,object]:
+def _load_predictor_func(path: str) -> Callable[...,object]:
     if "" not in sys.path:
         sys.path.insert(0, "")
     path: Path = Path(path).resolve()
@@ -97,7 +122,7 @@ def _load_predictor_func (path: str) -> Callable[...,object]:
     main_func = next(func for _, func in getmembers(module, isfunction) if hasattr(func, "__predictor_spec"))
     return main_func
 
-def _compute_predictor_url (api_url: str, tag: str) -> str:
+def _compute_predictor_url(api_url: str, tag: str) -> str:
     parsed_url = urlparse(api_url)
     hostname_parts = parsed_url.hostname.split(".")
     if hostname_parts[0] == "api":
@@ -107,32 +132,39 @@ def _compute_predictor_url (api_url: str, tag: str) -> str:
     predictor_url = urlunparse(parsed_url._replace(netloc=netloc, path=f"{tag}"))
     return predictor_url
 
-class _Predictor (BaseModel):
+class _Predictor(BaseModel):
     tag: str
 
-class _LogData (BaseModel):
+class _LogData(BaseModel):
     message: str
     level: int = 0
     status: Literal["success", "error"] = "success"
     update: bool = False
 
-class _LogEvent (BaseModel):
+class _LogEvent(BaseModel):
     event: Literal["log"]
     data: _LogData
 
-class _ErrorData (BaseModel):
+class _ErrorData(BaseModel):
     error: str
 
-class _ErrorEvent (BaseModel):
+class _ErrorEvent(BaseModel):
     event: Literal["error"]
     data: _ErrorData
 
+class _CompileError(Exception):
+    pass
+
+class _TriagedCompileError(BaseModel):
+    user: str
+    internal: str
+
 class ProgressLogQueue:
 
-    def __init__ (self):
+    def __init__(self):
         self.queue: list[tuple[int, CustomProgressTask]] = []
 
-    def push_log (self, event: _LogEvent):
+    def push_log(self, event: _LogEvent):
         # Check for update
         if event.data.update and self.queue:
             current_level, current_task = self.queue[-1]
@@ -149,15 +181,15 @@ class ProgressLogQueue:
         task.__enter__()
         self.queue.append((event.data.level, task))
 
-    def push_error (self, error: _ErrorEvent):
+    def push_error(self, error: _ErrorEvent):
         while self.queue:
             _, current_task = self.queue.pop()
             current_task.__exit__(RuntimeError, None, None)
 
-    def __enter__ (self):
+    def __enter__(self):
         return self
     
-    def __exit__ (self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type, exc_value, traceback):
         while self.queue:
             _, current_task = self.queue.pop()
             current_task.__exit__(None, None, None)

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/misc.py

@@ -9,22 +9,22 @@ from webbrowser import open as open_browser
 
 from ..version import __version__
 
-def _explore (value: bool):
+def _explore(value: bool):
     if value:
         open_browser("https://fxn.ai/explore")
         raise Exit()
 
-def _learn (value: bool):
+def _learn(value: bool):
     if value:
         open_browser("https://docs.fxn.ai")
         raise Exit()
 
-def _version (value: bool):
+def _version(value: bool):
     if value:
         print(__version__)
         raise Exit()
 
-def cli_options (
+def cli_options(
     explore: bool = Option(None, "--explore", callback=_explore, help="Explore predictors on Function."),
     learn: bool = Option(None, "--learn", callback=_learn, help="Learn about Function."),
     version: bool = Option(None, "--version", callback=_version, help="Get the Function CLI version.")

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/predictions.py

@@ -24,7 +24,7 @@ def create_prediction (
 ):
     run_async(_predict_async(tag, quiet=quiet, context=context))
 
-async def _predict_async (tag: str, quiet: bool, context: Context):
+async def _predict_async(tag: str, quiet: bool, context: Context):
     # Preload
     with CustomProgress(transient=True, disable=quiet):
         fxn = Function(get_access_key())
@@ -42,16 +42,7 @@ async def _predict_async (tag: str, quiet: bool, context: Context):
             prediction = fxn.predictions.create(tag, inputs=inputs)
     _log_prediction(prediction)
 
-def _parse_value (value: str):
-    """
-    Parse a value from a CLI argument.
-
-    Parameters:
-        value (str): CLI input argument.
-
-    Returns:
-        bool | int | float | str | Path: Parsed value.
-    """
+def _parse_value (value: str) -> float | int | bool | str | Image.Image | BytesIO:
     # Boolean
     if value == "true":
         return True
@@ -81,14 +72,14 @@ def _parse_value (value: str):
     # String
     return value
 
-def _log_prediction (prediction: Prediction):
+def _log_prediction(prediction: Prediction):
     images = [value for value in prediction.results or [] if isinstance(value, Image.Image)]
     prediction.results = [_serialize_value(value) for value in prediction.results] if prediction.results is not None else None
     print_json(data=prediction.model_dump())
     for image in images:
         image.show()
 
-def _serialize_value (value):
+def _serialize_value(value) -> str:
     if isinstance(value, ndarray):
         return array_repr(value)
     if isinstance(value, Image.Image):

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/predictors.py

@@ -10,7 +10,7 @@ from ..function import Function
 from ..logging import CustomProgress, CustomProgressTask
 from .auth import get_access_key
 
-def retrieve_predictor (
+def retrieve_predictor(
     tag: str=Argument(..., help="Predictor tag.")
 ):
     with CustomProgress(transient=True):
@@ -20,7 +20,7 @@ def retrieve_predictor (
             predictor = predictor.model_dump() if predictor else None
             print_json(data=predictor)
 
-def archive_predictor (
+def archive_predictor(
     tag: str=Argument(..., help="Predictor tag.")
 ):
     with CustomProgress():
@@ -34,7 +34,7 @@ def archive_predictor (
                 path=f"/predictors/{tag}/archive"
             )
 
-def delete_predictor (
+def delete_predictor(
     tag: str=Argument(..., help="Predictor tag.")
 ):
     with CustomProgress():

{fxn-0.0.54 → fxn-0.0.55}/fxn/cli/sources.py

@@ -14,7 +14,7 @@ from ..function import Function
 from ..logging import CustomProgress, CustomProgressTask
 from .auth import get_access_key
 
-def retrieve_source (
+def retrieve_source(
     prediction: Annotated[str, Option(help="Prediction identifier. If specified, this MUST be from a prediction returned by the Function API.")] = None,
     predictor: Annotated[str, Option(help="Predictor tag. If specified, a prediction will be made with this predictor before retrieving the source.")] = None,
     output: Annotated[Path, Option(help="Path to output source file.")] = Path("predictor.cpp")
@@ -37,7 +37,7 @@ def retrieve_source (
             source.code = str(output.resolve())
             print_json(data=source.model_dump(mode="json", by_alias=True))
 
-class _PredictionSource (BaseModel):
+class _PredictionSource(BaseModel):
     tag: str
     target: str
     code: str

{fxn-0.0.54 → fxn-0.0.55}/fxn/compile.py

@@ -14,7 +14,7 @@ from typing import Any, Callable, Literal, ParamSpec, TypeVar, cast
 from .beta import (
     CoreMLInferenceMetadata, LiteRTInferenceMetadata, LlamaCppInferenceMetadata,
     OnnxInferenceMetadata, OnnxRuntimeInferenceSessionMetadata, OpenVINOInferenceMetadata,
-    QnnInferenceMetadata
+    QnnInferenceMetadata, TensorRTInferenceMetadata
 )
 from .sandbox import Sandbox
 from .types import PredictorAccess
@@ -36,7 +36,8 @@ CompileMetadata = (
     OnnxInferenceMetadata               |
     OnnxRuntimeInferenceSessionMetadata |
     OpenVINOInferenceMetadata           |
-    QnnInferenceMetadata
+    QnnInferenceMetadata                |
+    TensorRTInferenceMetadata
 )
 
 P = ParamSpec("P")

{fxn-0.0.54 → fxn-0.0.55}/fxn/services/prediction.py

@@ -117,6 +117,21 @@ class PredictionService:
                 with prediction:
                     yield self.__to_prediction(tag, prediction)
 
+    def delete (self, tag: str) -> bool:
+        """
+        Delete a predictor that is loaded in memory.
+
+        Parameters:
+            tag (str): Predictor tag.
+
+        Returns:
+            bool: Whether the predictor was successfully deleted from memory.
+        """
+        if tag not in self.__cache:
+            return False
+        with self.__cache.pop(tag):
+            return True
+
     def __create_raw_prediction (
         self,
         tag: str,
@@ -245,7 +260,7 @@ class PredictionService:
                     tmp_file.write(chunk)
                     completed += len(chunk)
                     task.update(total=size, completed=completed)
-            Path(tmp_file.name).replace(path)
+        Path(tmp_file.name).replace(path)
         return path
 
     def __get_resource_path (self, resource: PredictionResource) -> Path:

{fxn-0.0.54 → fxn-0.0.55}/fxn/version.py

@@ -3,4 +3,4 @@
 #   Copyright © 2025 NatML Inc. All Rights Reserved.
 #
 
-__version__ = "0.0.54"
+__version__ = "0.0.55"

{fxn-0.0.54 → fxn-0.0.55}/fxn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fxn
-Version: 0.0.54
+Version: 0.0.55
 Summary: Run prediction functions locally in Python. Register at https://fxn.ai.
 Author-email: "NatML Inc." <hi@fxn.ai>
 License:                                  Apache License