torch¶

async def load(self) -> None:
    """Load TorchScript model and prepare for inference (async).

    Performs:
        - Load model from disk (in thread pool)
        - Setup device
        - Move model to device
        - Set evaluation mode
        - Apply precision
        - Warmup inference (in thread pool)

    Raises:
        FileNotFoundError: If model file does not exist.
        RuntimeError: If device is not available.
    """
    logger.info(f"Loading model from {self.model_path}")

    # Load model in thread pool
    loop = asyncio.get_event_loop()
    self.model = await loop.run_in_executor(None, torch.jit.load, str(self.model_path))  # type: ignore

    # Setup device (reuse mixin)
    self._torch_device = self._setup_torch_device()

    # Configure model (reuse mixin)
    self.model.to(self._torch_device)
    self.model.eval()
    self.model = self._apply_precision_to_model(self.model)

    logger.info(f"Model loaded on {self._torch_device}")

    # Warmup
    await self._warmup()

async def infer(self, input: torch.Tensor) -> t.Any:
    """Run inference on a single input (async).

    Automatically handles:
        - Moving input to correct device
        - Converting to correct precision
        - Adding batch dimension (if configured)
        - Removing batch dimension (if added)

    Args:
        input: Input tensor. Can be 3D (C, H, W) if auto_add_batch_dim=True,
            or 4D (1, C, H, W) otherwise.

    Returns:
        Model output.  Type depends on model architecture (tensor or tuple).

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        async with runtime:
            # 3D input (auto_add_batch_dim=True)
            input = torch.randn(3, 224, 224)
            output = await runtime.infer(input)

            # 4D input
            input = torch.randn(1, 3, 224, 224)
            output = await runtime.infer(input)
        ```
    """
    if self.model is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Prepare input (reuse mixin)
    input = self._prepare_input(input, self._torch_device)
    input, added_batch = self._add_batch_dim_if_needed(input, self.auto_add_batch_dim)

    # Inference in thread pool
    loop = asyncio.get_event_loop()
    output = await loop.run_in_executor(None, self._infer_sync, input)

    # Post-process (reuse mixin)
    return self._remove_batch_dim_if_added(output, added_batch)

async def infer_batch(self, inputs: list[torch.Tensor]) -> list[t.Any]:
    """Run inference on a batch of inputs (async).

    Concatenates inputs into a single batch tensor for efficient
    processing, then splits the output back into individual results.

    Args:
        inputs: List of input tensors. Each should have shape (1, C, H, W).

    Returns:
        List of outputs, one per input. Each maintains batch dimension.

    Raises:
        RuntimeError: If model is not loaded.

    Example:
        ```python
        async with runtime:
            batch = [
                torch.randn(1, 3, 224, 224),
                torch.randn(1, 3, 224, 224),
                torch.randn(1, 3, 224, 224),
            ]

            # Efficient batch processing
            outputs = await runtime.infer_batch(batch)

            # outputs[0], outputs[1], outputs[2]
        ```
    """
    if self.model is None:
        raise RuntimeError("Model not loaded. Call load() first.")

    # Concatenate
    batch = torch.cat(inputs, dim=0).to(self._torch_device)

    if self.precision == Precision.FP16:
        batch = batch.half()

    # Inference in thread pool
    loop = asyncio.get_event_loop()
    batch_output = await loop.run_in_executor(None, self._infer_sync, batch)

    # Split (reuse mixin)
    return self._split_batch_output(batch_output, len(inputs))

async def unload(self) -> None:
    """Unload model and free resources (async).

    Performs:
        - Release model from memory
        - Clear CUDA cache (if using CUDA)

    Safe to call multiple times.
    """
    logger.info("Unloading model")
    self.model = None

    if self.device.type.value == "cuda":
        torch.cuda.empty_cache()

    logger.info("Model unloaded")