leapp 0.5.1__py3-none-any.whl

leapp/__init__.py

@@ -0,0 +1,49 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+LEAPP - Lightweight Export Annotations for Policy Pipelines
+
+A Python package for tracing and exporting computational graphs from PyTorch code.
+LEAPP is specifically designed for robotics and autonomous agent applications, allowing 
+you to trace and export complex policy pipelines with interconnected components to 
+various formats including PyTorch JIT, ONNX, and generate visualization and YAML specifications.
+"""
+
+from .export_manager import ExportManager
+from .inference_manager import InferenceManager
+from .leapp import annotate, start, stop, compile_graph
+from .utils.enums import InputKindEnum, OutputKindEnum
+from .utils.tensor_description import TensorSemantics
+
+__version__ = "0.5.1"
+__config_version__ = "1.1"
+__author__ = "Frank Lai"
+__email__ = "frlai@nvidia.com"
+
+__all__ = [
+    "ExportManager",
+    "InferenceManager",
+    "InputKindEnum",
+    "OutputKindEnum",
+    "annotate",
+    "start",
+    "stop",
+    "compile_graph",
+    "__version__",
+    "TensorSemantics",
+]

leapp/backends/export_backend.py

@@ -0,0 +1,538 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import abc
+import hashlib
+import os
+import shutil
+from typing import Callable, Tuple, Any
+
+import numpy as np
+import torch
+import onnx
+import onnxruntime as ort
+
+from leapp.utils.logging import _get_logger
+from leapp.backends.module_builder import ModuleBuilder
+
+
+class SimplifiedONNXProgram:
+    """Wrapper for ONNX models.
+
+    This class mimics the behavior of ONNXProgram generated by torch.onnx.export
+    with onnx-dynamo, allowing the model to be called directly for inference.
+
+    Keeps source files on disk and copies on save. Cleans up temp dir on delete.
+    """
+
+    def __init__(self, onnx_model_path, temp_dir=None):
+        """Initialize the ONNX program.
+
+        Args:
+            onnx_model_path: Path to an ONNX file.
+            temp_dir: If provided, this temp directory will be cleaned up when the object is deleted.
+        """
+        model_path = str(onnx_model_path)
+        self._source_dir = os.path.dirname(os.path.abspath(model_path))
+        self._source_filename = os.path.basename(model_path)
+        self._temp_dir = temp_dir
+
+        # Session created lazily on first __call__
+        self._session = None
+        self._input_names = None
+        self._output_metas = None
+        self._active_provider = None
+
+    @staticmethod
+    def _torch_dtype_to_numpy_dtype(dtype):
+        mapping = {
+            torch.bool: np.bool_,
+            torch.uint8: np.uint8,
+            torch.int8: np.int8,
+            torch.int16: np.int16,
+            torch.int32: np.int32,
+            torch.int64: np.int64,
+            torch.float16: np.float16,
+            torch.float32: np.float32,
+            torch.float64: np.float64,
+        }
+        return mapping.get(dtype)
+
+    @staticmethod
+    def _onnx_type_to_torch_dtype(type_str):
+        mapping = {
+            "tensor(bool)": torch.bool,
+            "tensor(uint8)": torch.uint8,
+            "tensor(int8)": torch.int8,
+            "tensor(int16)": torch.int16,
+            "tensor(int32)": torch.int32,
+            "tensor(int64)": torch.int64,
+            "tensor(float16)": torch.float16,
+            "tensor(float)": torch.float32,
+            "tensor(double)": torch.float64,
+        }
+        return mapping.get(type_str)
+
+    @staticmethod
+    def _output_shape_is_static(shape):
+        if shape is None:
+            return False
+        return all(isinstance(dim, int) and dim >= 0 for dim in shape)
+
+    def _can_use_cuda_iobinding(self, args):
+        if self._active_provider != 'CUDAExecutionProvider':
+            return False
+
+        if not args:
+            return False
+
+        first_tensor = args[0]
+        if not isinstance(first_tensor, torch.Tensor) or not first_tensor.is_cuda:
+            return False
+
+        for tensor in args:
+            if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda:
+                return False
+            if tensor.device != first_tensor.device:
+                return False
+
+        if self._output_metas is None:
+            return False
+
+        for output_meta in self._output_metas:
+            if not self._output_shape_is_static(output_meta.shape):
+                return False
+            if self._onnx_type_to_torch_dtype(output_meta.type) is None:
+                return False
+
+        return True
+
+    def _run_with_standard_inference(self, args, output_device):
+        input_dict = {}
+        for name, tensor in zip(self._input_names, args):
+            if isinstance(tensor, torch.Tensor):
+                input_dict[name] = tensor.detach().cpu().numpy()
+            else:
+                input_dict[name] = tensor
+
+        outputs = self._session.run(None, input_dict)
+        return tuple(torch.from_numpy(out).to(output_device) for out in outputs)
+
+    def _run_with_cuda_iobinding(self, args):
+        binding = self._session.io_binding()
+
+        prepared_inputs = []
+        for name, tensor in zip(self._input_names, args):
+            prepared = tensor.detach().contiguous()
+            prepared_inputs.append(prepared)
+            device_id = prepared.device.index or 0
+            np_dtype = self._torch_dtype_to_numpy_dtype(prepared.dtype)
+            if np_dtype is None:
+                raise TypeError(
+                    f"Unsupported CUDA input dtype for ONNX Runtime I/O binding: {prepared.dtype}"
+                )
+            binding.bind_input(
+                name=name,
+                device_type='cuda',
+                device_id=device_id,
+                element_type=np_dtype,
+                shape=tuple(prepared.shape),
+                buffer_ptr=prepared.data_ptr(),
+            )
+
+        output_device = prepared_inputs[0].device
+        outputs = []
+        for output_meta in self._output_metas:
+            torch_dtype = self._onnx_type_to_torch_dtype(output_meta.type)
+            if torch_dtype is None:
+                raise TypeError(
+                    f"Unsupported ONNX output dtype for CUDA I/O binding: {output_meta.type}"
+                )
+
+            output_tensor = torch.empty(
+                tuple(output_meta.shape),
+                dtype=torch_dtype,
+                device=output_device,
+            ).contiguous()
+            binding.bind_output(
+                name=output_meta.name,
+                device_type='cuda',
+                device_id=output_device.index or 0,
+                element_type=self._torch_dtype_to_numpy_dtype(torch_dtype),
+                shape=tuple(output_tensor.shape),
+                buffer_ptr=output_tensor.data_ptr(),
+            )
+            outputs.append(output_tensor)
+
+        self._session.run_with_iobinding(binding)
+        return tuple(outputs)
+
+    def __del__(self):
+        """Clean up temp directory if we own it."""
+        if getattr(self, '_temp_dir', None) is not None:
+            try:
+                if os.path.exists(self._temp_dir):
+                    shutil.rmtree(self._temp_dir)
+            except Exception:
+                pass  # Silently ignore cleanup errors
+
+    def _get_providers(self):
+        """Get execution providers, always preferring CUDA when available."""
+        available_providers = ort.get_available_providers()
+
+        if 'CUDAExecutionProvider' in available_providers:
+            return ['CUDAExecutionProvider', 'CPUExecutionProvider']
+
+        _get_logger().warning(
+            "CUDA execution provider not available. Falling back to CPU."
+        )
+        return ['CPUExecutionProvider']
+
+    def _get_source_size(self):
+        """Get total size of all files in the source directory."""
+        total_size = 0
+        files_info = []
+        for f in os.listdir(self._source_dir):
+            file_path = os.path.join(self._source_dir, f)
+            if os.path.isfile(file_path):
+                size = os.path.getsize(file_path)
+                total_size += size
+                files_info.append(f"{f} ({size / (1024**2):.2f} MB)")
+
+        return total_size
+
+    def save(self, destination, include_initializers=True, keep_initializers_as_inputs=False):
+        """Save the ONNX model to disk.
+
+        For large models (>2GB), re-saves with external data format.
+        For small models, copies files directly.
+        """
+        # 2GB threshold for protobuf limit
+        SIZE_THRESHOLD = 2 * 1024 * 1024 * 1024
+
+        dest_dir = os.path.dirname(os.path.abspath(destination))
+        dest_filename = os.path.basename(destination)
+        src_path = os.path.join(self._source_dir, self._source_filename)
+
+        total_size = self._get_source_size()
+
+        if total_size > SIZE_THRESHOLD:
+            # Large model - load and re-save with external data
+            _get_logger().info(
+                f"Large model (~{total_size / (1024**3):.2f} GB) - saving with external data format"
+            )
+            # Ensure destination directory exists
+            os.makedirs(dest_dir, exist_ok=True)
+
+            model_proto = onnx.load(src_path, load_external_data=True)
+            data_filename = dest_filename + ".data"
+            data_path = os.path.join(dest_dir, data_filename)
+
+            # Delete existing data file to avoid appending (ONNX appends instead of overwriting)
+            if os.path.exists(data_path):
+                os.remove(data_path)
+
+            _get_logger().info(
+                f"Saving model to {destination} with external data in {data_filename}")
+            onnx.save(
+                model_proto,
+                destination,
+                save_as_external_data=True,
+                all_tensors_to_one_file=True,
+                location=data_filename,
+                size_threshold=1024,
+                convert_attribute=True,  # Also convert Constant node attributes
+            )
+
+            # Verify external data file was created
+            expected_data_path = os.path.join(dest_dir, data_filename)
+            if os.path.exists(expected_data_path):
+                _get_logger().info(
+                    f"External data file created: {expected_data_path}")
+            else:
+                _get_logger().warning(
+                    f"External data file NOT found: {expected_data_path}")
+        else:
+            # Small model - just copy files
+            shutil.copy2(src_path, destination)
+
+            # Copy ALL other files from source directory (external data files)
+            for f in os.listdir(self._source_dir):
+                if f != self._source_filename:
+                    src_file = os.path.join(self._source_dir, f)
+                    if os.path.isfile(src_file):
+                        dst_file = os.path.join(dest_dir, f)
+                        shutil.copy2(src_file, dst_file)
+
+    def _ensure_session(self):
+        """Create the inference session if not already created."""
+        if self._session is None:
+            model_path = os.path.join(self._source_dir, self._source_filename)
+            providers = self._get_providers()
+            sess_options = ort.SessionOptions()
+            # ORT_ENABLE_ALL can silently corrupt results for certain graph
+            # patterns (e.g. Gemm chains produced by FX make_fx decomposition).
+            # ORT_ENABLE_BASIC is safe and still applies constant folding.
+            sess_options.graph_optimization_level = (
+                ort.GraphOptimizationLevel.ORT_ENABLE_BASIC)
+            self._session = ort.InferenceSession(
+                model_path, sess_options, providers=providers)
+            self._input_names = [
+                inp.name for inp in self._session.get_inputs()]
+            self._output_metas = list(self._session.get_outputs())
+            self._active_provider = self._session.get_providers(
+            )[0] if self._session.get_providers() else 'CPUExecutionProvider'
+
+    def __call__(self, *args):
+        """Run inference on the ONNX model.
+
+        Args:
+            *args: Input tensors in order, matching the model's input signature.
+
+        Returns:
+            List of output tensors as torch.Tensor objects on the same device as inputs.
+        """
+        self._ensure_session()
+
+        if len(args) != len(self._input_names):
+            raise ValueError(
+                f"Expected {len(self._input_names)} inputs, got {len(args)}. "
+                f"Input names: {self._input_names}"
+            )
+
+        # Determine output device from input tensors
+        output_device = 'cpu'
+        for tensor in args:
+            if isinstance(tensor, torch.Tensor) and tensor.is_cuda:
+                output_device = tensor.device
+                break
+
+        if self._can_use_cuda_iobinding(args):
+            return self._run_with_cuda_iobinding(args)
+
+        return self._run_with_standard_inference(args, output_device)
+
+
+def prepare_tensors_for_export(tensors):
+    """
+    Prepare tensors for export by cloning them to escape inference mode.
+
+    Tensors created under torch.inference_mode() cannot participate in autograd,
+    which causes torch.export.export() (used by dynamo) to fail with:
+    "RuntimeError: Inference tensors cannot be saved for backward."
+
+    Cloning creates new tensors that are not marked as inference tensors.
+
+    Args:
+        tensors: A sequence of tensors (or other values) to prepare.
+
+    Returns:
+        A tuple of prepared tensors (cloned if they were torch.Tensor).
+    """
+    prepared = []
+    for t in tensors:
+        if isinstance(t, torch.Tensor):
+            if hasattr(t, 'original_clone'):
+                prepared.append(t.original_clone())
+            else:
+                prepared.append(t.clone())
+        else:
+            prepared.append(t)
+    return tuple(prepared)
+
+
+class ExportBackend(abc.ABC):
+    def __init__(self, node_context, backend_params=None):
+        self.node_context = node_context
+        if backend_params is None:
+            self.backend_params = {}
+        else:
+            self.backend_params = backend_params
+
+        self.module_builder = ModuleBuilder(node_context)
+
+        self.compiled_model = None
+        self.compiled_module = None
+        self.runtime_device = None
+    
+    def override_module_builder(self, module_builder: Callable):
+        self.module_builder = module_builder
+
+    def _verify_model_location_and_get_hash(self, model_path):
+        if not os.path.exists(model_path):
+            _get_logger().error(f"Model file not found at {model_path}")
+            return None, None
+        with open(model_path, 'rb') as f:
+            file_data = f.read()
+            md5sum = hashlib.md5(file_data).hexdigest()
+            sha256sum = hashlib.sha256(file_data).hexdigest()
+        return md5sum, sha256sum
+
+    def _copy_model_to_path(self, model_path, save_path):
+        if not os.path.exists(save_path):
+            os.makedirs(save_path)
+        if not os.path.exists(model_path):
+            return None
+
+        # Check if save_path is the same as the directory containing model_path
+        model_dir = os.path.dirname(os.path.abspath(model_path))
+        save_dir = os.path.abspath(save_path)
+        if model_dir == save_dir:
+            # No need to copy if already in the same directory
+            return model_path
+
+        # Get the filename from model_path
+        filename = os.path.basename(model_path)
+        # Create the full destination path
+        dest_path = os.path.join(save_path, filename)
+        # Copy the file
+        shutil.copy2(model_path, dest_path)
+
+        return dest_path
+
+    @abc.abstractmethod
+    def get_backend_model_type(self):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_backend_metadata(self) -> dict:
+        """Optional backend-specific metadata to include in node YAML parameters."""
+        return {}
+
+    @abc.abstractmethod
+    def compile(self, m: torch.nn.Module = None) -> Any:
+        '''
+        Compiles the model.
+
+        This function should return the compiled model. The resulting compiled model 
+        can be used in recombination with other models.
+        '''
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def save(self, save_path: str) -> Tuple[str, str, str]:
+        '''
+        Save the compiled model to the given path
+
+        This function should apply all necessary optimizations
+        '''
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def load(self, model_path: str, sha256sum: str):
+        raise NotImplementedError
+    
+    def _select_runtime_device(self):
+        return 'cuda' if torch.cuda.is_available() else 'cpu'
+
+    def _load_onnx(self, model_path: str, sha256sum: str):
+        _, actual_sha256sum = self._verify_model_location_and_get_hash(
+            model_path)
+        if actual_sha256sum != sha256sum:
+            raise ValueError(
+                f"SHA256 checksum mismatch for {model_path}: "
+                f"expected {sha256sum}, got {actual_sha256sum}"
+            )
+        model = SimplifiedONNXProgram(model_path)
+        self.compiled_model = model
+        has_ort_cuda = 'CUDAExecutionProvider' in ort.get_available_providers()
+        self.runtime_device = 'cuda' if (
+            torch.cuda.is_available() and has_ort_cuda
+        ) else 'cpu'
+        self.compiled_module = None # ONNX models cannot be represented as a module for reexport.
+    
+    def _load_torchscript(self, model_path: str, sha256sum: str):
+        _, actual_sha256sum = self._verify_model_location_and_get_hash(
+            model_path)
+        if actual_sha256sum != sha256sum:
+            raise ValueError(
+                f"SHA256 checksum mismatch for {model_path}: "
+                f"expected {sha256sum}, got {actual_sha256sum}"
+            )
+        device = self._select_runtime_device()
+        model = torch.jit.load(model_path, map_location=device)
+        model = model.to(device)
+        self.compiled_model = model.eval()
+        self.runtime_device = device
+        self.compiled_module = model
+
+
+class NoneExportBackend(ExportBackend):
+    """Null export backend for nodes without an active export backend.
+
+    This backend represents "no export" behavior in normal usage. It can also
+    optionally act as a thin wrapper for existing model artifacts when
+    ``backend_params['model_path']`` is provided, including checksum validation
+    and backend-specific loading.
+    """
+
+    def get_backend_metadata(self):
+        return {}
+
+    def compile(self, m: torch.nn.Module = None):
+        if "model_path" not in self.backend_params or self.backend_params['model_path'] is None:
+            _get_logger().warning(
+                f"No model path provided for {self.node_context.name}")
+            _get_logger().warning("if this is intentional, please provide a path to the correct model "
+                                  "in the generated yaml file. Otherwise, please manually fill in the backend parameters.")
+        else:
+            _, sha256sum = self._verify_model_location_and_get_hash(self.backend_params['model_path'])
+            self.load(self.backend_params['model_path'], sha256sum, self.get_backend_model_type())
+
+
+    def save(self, save_path: str) -> Tuple[str, str, str]:
+        if "model_path" not in self.backend_params or self.backend_params['model_path'] is None:
+            return None, None, None
+        md5sum, sha256sum = self._verify_model_location_and_get_hash(
+            self.backend_params['model_path'])
+        model_path = self.backend_params['model_path']
+
+        if "copy_original_model" in self.backend_params and self.backend_params['copy_original_model'] is True:
+            model_path = self._copy_model_to_path(model_path, save_path)
+
+        return model_path, md5sum, sha256sum
+
+    def load(self, model_path: str, sha256sum: str, model_type=None):
+        if model_type is None: 
+            return
+        if model_type == "onnx":
+            self._load_onnx(model_path, sha256sum)
+        elif model_type == "jit":
+            self._load_torchscript(model_path, sha256sum)
+        else:
+            raise ValueError(f"Unsupported model type: {model_type}")
+
+    def get_backend_model_type(self):
+        if "model_path" not in self.backend_params:
+            return None
+
+        path = self.backend_params['model_path']
+        suffix = path.split('.')[-1]
+        if suffix == 'pt':
+            return "jit"
+        elif suffix == 'pt2':
+            return "torchscript2"
+        elif suffix == 'onnx':
+            return "onnx"
+        elif suffix == 'cpp' or suffix == "cc":
+            return "cpp"
+        elif suffix == 'py':
+            return "py"
+        elif suffix == 'engine' or suffix == 'plan':
+            return 'trt'
+        else:
+            raise ValueError(
+                f"Unsupported model file suffix: {suffix}")