From c620383a608d62ae9462b7552429daccb0fb73c5 Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Fri, 23 Oct 2026 13:49:59 -0700 Subject: [PATCH 1/2] feat: add Warp plugins for Triton ONNX and PT2 Discover installed Warp adapters for Triton ONNX Runtime and NativeRT/PT2 models. Keep ONNX CUDA replay tensors on the GPU and support native Torch registration, stream ordering, and cleanup when models unload. Add explicit SDK selection, generic CMake argument forwarding, and install rules without embedding the absolute Warp library path in dependencies. Validate SDKs before replacing an existing cached runtime build. Validation: seven build-helper tests, two native CUDA replay tests, and two Triton integration tests covering twelve session placements and load errors. Verified SDK overrides, relocated installs, and formatting. Signed-off-by: Marc Alwan --- .../warp_operator/__init__.py | 10 +- .../warp_operator/runtime/CMakeLists.txt | 201 ++++++++++-------- .../warp_operator/runtime/README.md | 58 ++++- .../runtime/core/warp_apic_runner.cc | 168 +++++++++------ .../runtime/core/warp_apic_runner.h | 4 + .../runtime/onnx/ort_wrp_runner_op.cc | 173 +++++++++------ .../runtime/torch/torch_warp_runner_op.cc | 68 +++++- packages/leapp-runtime-rust/README.md | 4 + .../leapp-runtime-rust/src/engines/triton.rs | 7 +- .../src/engines/triton/backends/README.md | 56 ++++- .../backends/torch/common/src/model_config.cc | 8 +- .../backends/torch/common/src/model_config.h | 3 +- .../triton/backends/torch/pt2/CMakeLists.txt | 2 +- .../triton/backends/torch/pt2/src/backend.cc | 35 ++- .../torch/pt2/src/native_rt_instance.h | 5 +- .../src/engines/triton/convert.rs | 56 ++++- .../tests/functional_tests/triton.rs | 2 + .../tests/functional_tests/triton/warp.rs | 110 ++++++++++ .../tests/generators/generate_warp_models.py | 123 +++++++++++ .../tests/unit_tests/triton/convert.rs | 71 +++++++ scripts/leapp-build-warp-runtime | 153 +++---------- tests/functional_tests/test_warp_tracing.py | 101 ++++++++- tests/unit_tests/test_warp_runtime.py | 138 +++--------- 23 files changed, 1074 insertions(+), 482 deletions(-) create mode 100644 packages/leapp-runtime-rust/tests/functional_tests/triton/warp.rs create mode 100644 packages/leapp-runtime-rust/tests/generators/generate_warp_models.py diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py b/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py index 201a10e5..515ff35a 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/__init__.py @@ -120,13 +120,19 @@ def _register() -> None: and hasattr(getattr(torch.ops, NAMESPACE), OP_NAME) ): _LIB.define(_SCHEMA) - torch.library.register_fake(QUALIFIED_NAME, fx_impl.warp_runner_fake, lib=_LIB) - _LIB.impl(OP_NAME, fx_impl.warp_runner_eager, "CompositeExplicitAutograd") _get_logger().debug( f"Registered custom op {QUALIFIED_NAME} with schema: {_SCHEMA}" ) + # The native adapter may register the schema before Python is imported. + if not torch._C._dispatch_has_kernel_for_dispatch_key(QUALIFIED_NAME, "Meta"): + torch.library.register_fake(QUALIFIED_NAME, fx_impl.warp_runner_fake, lib=_LIB) + if not torch._C._dispatch_has_kernel_for_dispatch_key( + QUALIFIED_NAME, "CompositeExplicitAutograd" + ): + _LIB.impl(OP_NAME, fx_impl.warp_runner_eager, "CompositeExplicitAutograd") + if wp is not None: onnx_impl.register_onnx_lowering() _register_export_hooks() diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/CMakeLists.txt b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/CMakeLists.txt index f1ec1bff..f1916761 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/CMakeLists.txt +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/CMakeLists.txt @@ -14,22 +14,44 @@ endforeach() option(LEAPP_WARP_BUILD_ONNX "Build ONNX Runtime custom op" ON) option(LEAPP_WARP_BUILD_TORCH "Build PyTorch dispatcher extension" ON) -find_package(Python3 REQUIRED COMPONENTS Interpreter) +include(GNUInstallDirs) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux" AND NOT DEFINED CMAKE_INSTALL_RPATH) + set(CMAKE_INSTALL_RPATH "$ORIGIN") +endif() + +if(NOT DEFINED CUDAToolkit_ROOT AND DEFINED ENV{CUDAToolkit_ROOT}) + set(CUDAToolkit_ROOT "$ENV{CUDAToolkit_ROOT}") +endif() +if(CUDAToolkit_ROOT AND NOT EXISTS "${CUDAToolkit_ROOT}/bin/nvcc${CMAKE_EXECUTABLE_SUFFIX}") + message(FATAL_ERROR "CUDA compiler not found under CUDAToolkit_ROOT: ${CUDAToolkit_ROOT}") +endif() find_package(CUDAToolkit REQUIRED) find_package(Threads REQUIRED) +message(STATUS "Warp CUDA toolkit: ${CUDAToolkit_LIBRARY_ROOT} (${CUDAToolkit_VERSION})") -execute_process( - COMMAND "${Python3_EXECUTABLE}" -c "import os, warp; print(os.path.dirname(warp.__file__))" - OUTPUT_VARIABLE WARP_PACKAGE_DIR - OUTPUT_STRIP_TRAILING_WHITESPACE - COMMAND_ERROR_IS_FATAL ANY -) +function(leapp_warp_python OUTPUT CODE) + find_package(Python3 REQUIRED COMPONENTS Interpreter) + execute_process( + COMMAND "${Python3_EXECUTABLE}" -c "${CODE}" + OUTPUT_VARIABLE VALUE + OUTPUT_STRIP_TRAILING_WHITESPACE + COMMAND_ERROR_IS_FATAL ANY + ) + set(${OUTPUT} "${VALUE}" PARENT_SCOPE) +endfunction() -set(WARP_NATIVE_DIR "${WARP_PACKAGE_DIR}/native") -set(WARP_BIN_DIR "${WARP_PACKAGE_DIR}/bin") +set(WARP_ROOT "$ENV{WARP_ROOT}" CACHE PATH "Warp SDK with native/ and bin/") +if(NOT WARP_ROOT) + leapp_warp_python(WARP_ROOT "import os, warp; print(os.path.dirname(warp.__file__))") +endif() +set(WARP_NATIVE_DIR "${WARP_ROOT}/native") +set(WARP_BIN_DIR "${WARP_ROOT}/bin") if(WIN32) set(WARP_SHARED_LIBRARY "${WARP_BIN_DIR}/warp.dll") - set(WARP_IMPORT_LIBRARY "${WARP_BIN_DIR}/warp.lib" CACHE FILEPATH "Warp import library") + set(WARP_IMPORT_LIBRARY "$ENV{LEAPP_WARP_IMPORT_LIBRARY}" CACHE FILEPATH "Warp import library") + if(NOT WARP_IMPORT_LIBRARY) + set(WARP_IMPORT_LIBRARY "${WARP_BIN_DIR}/warp.lib") + endif() elseif(APPLE) set(WARP_SHARED_LIBRARY "${WARP_BIN_DIR}/libwarp.dylib") else() @@ -43,17 +65,22 @@ if(NOT EXISTS "${WARP_NATIVE_DIR}/apic.h") message(FATAL_ERROR "Warp APIC header not found: ${WARP_NATIVE_DIR}/apic.h. Install a Warp build with APIC support (Warp 1.13+ expected).") endif() -add_library(warp_native SHARED IMPORTED) -set_target_properties(warp_native PROPERTIES - IMPORTED_LOCATION "${WARP_SHARED_LIBRARY}" -) -if(WIN32) - if(NOT EXISTS "${WARP_IMPORT_LIBRARY}") - message(FATAL_ERROR "Warp import library not found: ${WARP_IMPORT_LIBRARY}. Set WARP_IMPORT_LIBRARY to a valid warp.lib.") +message(STATUS "Warp native library: ${WARP_SHARED_LIBRARY}") +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + # Warp's wheel ships warp.so without a SONAME or the usual lib prefix. + # Exact-filename linking avoids recording the absolute SDK path in DT_NEEDED. + add_library(warp_native INTERFACE) + target_link_directories(warp_native INTERFACE "${WARP_BIN_DIR}") + target_link_libraries(warp_native INTERFACE "-l:warp.so") +else() + add_library(warp_native SHARED IMPORTED) + set_target_properties(warp_native PROPERTIES IMPORTED_LOCATION "${WARP_SHARED_LIBRARY}") + if(WIN32) + if(NOT EXISTS "${WARP_IMPORT_LIBRARY}") + message(FATAL_ERROR "Warp import library not found: ${WARP_IMPORT_LIBRARY}. Set WARP_IMPORT_LIBRARY to a valid warp.lib.") + endif() + set_target_properties(warp_native PROPERTIES IMPORTED_IMPLIB "${WARP_IMPORT_LIBRARY}") endif() - set_target_properties(warp_native PROPERTIES - IMPORTED_IMPLIB "${WARP_IMPORT_LIBRARY}" - ) endif() set(LEAPP_WARP_CORE_SOURCES @@ -70,73 +97,62 @@ target_include_directories(leapp_warp_runtime_core PUBLIC "${WARP_NATIVE_DIR}" ) target_link_libraries(leapp_warp_runtime_core PUBLIC warp_native CUDA::cudart Threads::Threads) +# Use CUDA vector types when Warp headers are included with the CUDA SDK. +target_compile_definitions(leapp_warp_runtime_core PRIVATE WP_ENABLE_CUDA=1) set_target_properties(leapp_warp_runtime_core PROPERTIES POSITION_INDEPENDENT_CODE ON) if(LEAPP_WARP_BUILD_ONNX) - execute_process( - COMMAND "${Python3_EXECUTABLE}" -c "import os, onnxruntime; print(os.path.dirname(onnxruntime.__file__))" - OUTPUT_VARIABLE ONNXRUNTIME_PACKAGE_DIR - OUTPUT_STRIP_TRAILING_WHITESPACE - COMMAND_ERROR_IS_FATAL ANY - ) - execute_process( - COMMAND "${Python3_EXECUTABLE}" -c "import onnxruntime; print(onnxruntime.__version__)" - OUTPUT_VARIABLE ONNXRUNTIME_VERSION - OUTPUT_STRIP_TRAILING_WHITESPACE - COMMAND_ERROR_IS_FATAL ANY - ) - - set(ORT_CAPI_DIR "${ONNXRUNTIME_PACKAGE_DIR}/capi") - if(WIN32) - set(ORT_SHARED_LIBRARY "${ORT_CAPI_DIR}/onnxruntime.dll") - elseif(APPLE) - set(ORT_SHARED_LIBRARY "${ORT_CAPI_DIR}/libonnxruntime.${ONNXRUNTIME_VERSION}.dylib") - else() - set(ORT_SHARED_LIBRARY "${ORT_CAPI_DIR}/libonnxruntime.so.${ONNXRUNTIME_VERSION}") - endif() - if(NOT EXISTS "${ORT_SHARED_LIBRARY}") - message(FATAL_ERROR "ONNX Runtime shared library not found: ${ORT_SHARED_LIBRARY}") - endif() - - set(ORT_INCLUDE_DIR "${CMAKE_BINARY_DIR}/onnxruntime_headers") - file(MAKE_DIRECTORY "${ORT_INCLUDE_DIR}") - foreach(ORT_HEADER IN ITEMS - onnxruntime_c_api.h - onnxruntime_ep_c_api.h - onnxruntime_error_code.h - ) - set(ORT_HEADER_PATH "${ORT_INCLUDE_DIR}/${ORT_HEADER}") - if(NOT EXISTS "${ORT_HEADER_PATH}") - file(DOWNLOAD - "https://raw.githubusercontent.com/microsoft/onnxruntime/v${ONNXRUNTIME_VERSION}/include/onnxruntime/core/session/${ORT_HEADER}" - "${ORT_HEADER_PATH}" - STATUS ORT_HEADER_DOWNLOAD_STATUS - TLS_VERIFY ON - ) - list(GET ORT_HEADER_DOWNLOAD_STATUS 0 ORT_HEADER_DOWNLOAD_CODE) - list(GET ORT_HEADER_DOWNLOAD_STATUS 1 ORT_HEADER_DOWNLOAD_MESSAGE) - if(NOT ORT_HEADER_DOWNLOAD_CODE EQUAL 0) - # Split out of onnxruntime_c_api.h in 1.29; older releases 404. - if(ORT_HEADER STREQUAL "onnxruntime_error_code.h") - file(REMOVE "${ORT_HEADER_PATH}") - message(STATUS - "Optional ${ORT_HEADER} is not available for " - "ONNX Runtime ${ONNXRUNTIME_VERSION}; skipping.") - else() - message(FATAL_ERROR "Failed to download ${ORT_HEADER}: ${ORT_HEADER_DOWNLOAD_MESSAGE}") + set(ORT_INCLUDE_DIR "$ENV{ORT_INCLUDE_DIR}" CACHE PATH "ONNX Runtime C API headers") + set(LEAPP_WARP_ONNXRUNTIME_VERSION "" CACHE STRING "Target ONNX Runtime header version (defaults to the Python package)") + if(NOT ORT_INCLUDE_DIR) + set(ORT_HEADER_VERSION "${LEAPP_WARP_ONNXRUNTIME_VERSION}") + if(NOT ORT_HEADER_VERSION) + leapp_warp_python(ORT_HEADER_VERSION "import onnxruntime; print(onnxruntime.__version__)") + endif() + set(ORT_INCLUDE_DIR "${CMAKE_BINARY_DIR}/onnxruntime_headers/${ORT_HEADER_VERSION}") + file(MAKE_DIRECTORY "${ORT_INCLUDE_DIR}") + foreach(ORT_HEADER IN ITEMS + onnxruntime_c_api.h + onnxruntime_ep_c_api.h + onnxruntime_error_code.h + ) + set(ORT_HEADER_PATH "${ORT_INCLUDE_DIR}/${ORT_HEADER}") + if(NOT EXISTS "${ORT_HEADER_PATH}") + file(DOWNLOAD + "https://raw.githubusercontent.com/microsoft/onnxruntime/v${ORT_HEADER_VERSION}/include/onnxruntime/core/session/${ORT_HEADER}" + "${ORT_HEADER_PATH}" + STATUS ORT_HEADER_DOWNLOAD_STATUS + TLS_VERIFY ON + ) + list(GET ORT_HEADER_DOWNLOAD_STATUS 0 ORT_HEADER_DOWNLOAD_CODE) + list(GET ORT_HEADER_DOWNLOAD_STATUS 1 ORT_HEADER_DOWNLOAD_MESSAGE) + if(NOT ORT_HEADER_DOWNLOAD_CODE EQUAL 0) + # Split out of onnxruntime_c_api.h in 1.29; older releases 404. + if(ORT_HEADER STREQUAL "onnxruntime_error_code.h") + file(REMOVE "${ORT_HEADER_PATH}") + message(STATUS + "Optional ${ORT_HEADER} is not available for " + "ONNX Runtime ${ORT_HEADER_VERSION}; skipping.") + else() + message(FATAL_ERROR "Failed to download ${ORT_HEADER}: ${ORT_HEADER_DOWNLOAD_MESSAGE}") + endif() endif() endif() + endforeach() + endif() + foreach(ORT_HEADER IN ITEMS onnxruntime_c_api.h onnxruntime_ep_c_api.h) + if(NOT EXISTS "${ORT_INCLUDE_DIR}/${ORT_HEADER}") + message(FATAL_ERROR "ONNX Runtime header not found: ${ORT_INCLUDE_DIR}/${ORT_HEADER}. Set ORT_INCLUDE_DIR to the target SDK headers.") endif() endforeach() + message(STATUS "Warp ONNX adapter headers: ${ORT_INCLUDE_DIR}") add_library(leapp_wrp_onnx_custom_op SHARED onnx/ort_wrp_runner_op.cc) target_include_directories(leapp_wrp_onnx_custom_op PRIVATE "${ORT_INCLUDE_DIR}") target_link_libraries(leapp_wrp_onnx_custom_op PRIVATE leapp_warp_runtime_core) if(NOT WIN32) - set_target_properties(leapp_wrp_onnx_custom_op PROPERTIES - BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}" - INSTALL_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}" - ) + set_property(TARGET leapp_wrp_onnx_custom_op APPEND PROPERTY + BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}") endif() endif() @@ -147,21 +163,32 @@ if(LEAPP_WARP_BUILD_TORCH) if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) set(CMAKE_CUDA_ARCHITECTURES native CACHE STRING "CUDA architectures for Torch-enabled builds") endif() - execute_process( - COMMAND "${Python3_EXECUTABLE}" -c "import torch; print(torch.utils.cmake_prefix_path)" - OUTPUT_VARIABLE TORCH_CMAKE_PREFIX_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE - COMMAND_ERROR_IS_FATAL ANY - ) - list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}") - find_package(Torch REQUIRED) + if(NOT Torch_DIR AND DEFINED ENV{Torch_DIR}) + set(Torch_DIR "$ENV{Torch_DIR}") + endif() + if(Torch_DIR AND NOT EXISTS "${Torch_DIR}/TorchConfig.cmake") + message(FATAL_ERROR "TorchConfig.cmake not found under Torch_DIR: ${Torch_DIR}") + endif() + find_package(Torch QUIET CONFIG) + if(NOT Torch_FOUND) + leapp_warp_python(TORCH_CMAKE_PREFIX_PATH "import torch; print(torch.utils.cmake_prefix_path)") + list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}") + find_package(Torch REQUIRED CONFIG) + endif() + message(STATUS "Warp Torch SDK: ${Torch_DIR}") add_library(leapp_wrp_torch_custom_op SHARED torch/torch_warp_runner_op.cc) target_link_libraries(leapp_wrp_torch_custom_op PRIVATE leapp_warp_runtime_core "${TORCH_LIBRARIES}") target_include_directories(leapp_wrp_torch_custom_op PRIVATE "${TORCH_INCLUDE_DIRS}") if(NOT WIN32) - set_target_properties(leapp_wrp_torch_custom_op PROPERTIES - BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}" - INSTALL_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}" - ) + set_property(TARGET leapp_wrp_torch_custom_op APPEND PROPERTY + BUILD_RPATH "${WARP_BIN_DIR};${CUDAToolkit_LIBRARY_DIR}") endif() endif() + +foreach(ADAPTER IN ITEMS leapp_wrp_onnx_custom_op leapp_wrp_torch_custom_op) + if(TARGET ${ADAPTER}) + install(TARGETS ${ADAPTER} + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} + RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) + endif() +endforeach() diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/README.md b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/README.md index 595307d2..384b6fa5 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/README.md +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/README.md @@ -15,8 +15,10 @@ leapp-build-warp-runtime leapp-build-warp-runtime --status ``` -The build command validates existing dependencies without installing them, -then builds both adapters under the standard user cache: +The build command runs CMake to find and validate existing SDKs, then builds both +adapters under the standard user cache. Before replacing an existing build, it +validates the SDKs in a temporary CMake directory; failed validation leaves the +cached build untouched. It never installs dependencies: - Linux: `build/libleapp_wrp_onnx_custom_op.so` and `build/libleapp_wrp_torch_custom_op.so` @@ -25,9 +27,55 @@ then builds both adapters under the standard user cache: LEAPP discovers the cached libraries automatically. Use `LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY` and -`LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY` only to override the cache paths. Build with -the Python environment that will load the libraries because the adapters link -against that environment's Warp and PyTorch installations. +`LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY` to select other build or installation paths. +By default, SDK discovery uses installed CUDA, CMake's Torch search paths, and +Python packages from the interpreter running the helper. + +Pass CMake configuration after `--` to select explicit SDKs: + +```bash +leapp-build-warp-runtime --build-dir /path/to/build -- \ + -DWARP_ROOT=/path/to/warp \ + -DTorch_DIR=/path/to/libtorch/share/cmake/Torch \ + -DCUDAToolkit_ROOT=/usr/local/cuda \ + -DORT_INCLUDE_DIR=/path/to/onnxruntime/include +``` + +`WARP_ROOT` contains `native/` and `bin/`; `ORT_INCLUDE_DIR` contains the ONNX +Runtime C API headers. Explicit SDK locations take precedence over Python +package discovery. Use Warp headers, `warp.so`, and exported graphs from the +same Warp release, and a Torch SDK matching the process that loads the adapter. + +When no ONNX header directory is supplied, CMake downloads the release selected +by `-DLEAPP_WARP_ONNXRUNTIME_VERSION=VERSION`, defaulting to Python's ONNX Runtime +version. Downloads are cached by version. The adapter does not link another +ONNX Runtime library; its host supplies the API. + +Install the adapters with standard CMake commands: + +```bash +cmake --install /path/to/build --prefix /path/to/install +``` + +Libraries install under `CMAKE_INSTALL_LIBDIR` (`GNUInstallDirs`). Warp, Torch, +and CUDA remain in their SDK locations. Build outputs use the discovered library +paths; installed Linux adapters default to `$ORIGIN`. Configure +`CMAKE_INSTALL_RPATH` or set `LD_LIBRARY_PATH` to the selected SDK library +directories when loading installed adapters. Only the LEAPP adapters are installed. + +ONNX registers CPU and CUDA kernels. CUDA replay keeps data tensors on the GPU +and uses ONNX Runtime's compute stream; only the embedded archive is copied to +host once during loading. CPU replay stages data through CUDA device 0. +Replay currently supports CUDA device 0 only and restores the caller's CUDA +context. Calls sharing a graph are ordered across streams; warm each instance +serially before concurrent use. Torch makes strided inputs contiguous before +replay. + +Torch caches runners without retaining model constants. Expired entries are +released on subsequent calls; native model owners can call +`LEAPP_WARP_ReleaseUnusedRunners()` after releasing constants to clean up +immediately (zero indicates success). Destruction waits for GPU work before +freeing graph resources. Bundle tensors must remain immutable during inference. Windows builds also require a Warp import library. The `warp-lang` wheel does not currently include `warp.lib`; set `LEAPP_WARP_IMPORT_LIBRARY` to an diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.cc b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.cc index 41035d6d..b1f50635 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.cc +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.cc @@ -8,24 +8,31 @@ #include "warp.h" -extern "C" int cudaMemcpyAsync(void* dst, const void* src, std::size_t count, int kind, void* stream); -extern "C" int cudaMemsetAsync(void* dst, int value, std::size_t count, void* stream); -extern "C" const char* cudaGetErrorString(int error); - namespace leapp::warp_runtime { namespace { -constexpr int kCudaSuccess = 0; -constexpr int kCudaMemcpyDeviceToDevice = 3; constexpr const char* kWrpFilename = "segment.wrp"; +class ScopedCudaContext { + public: + explicit ScopedCudaContext(void* context) : previous_(wp_cuda_context_get_current()) { + wp_cuda_context_set_current(context); + } + ~ScopedCudaContext() { wp_cuda_context_set_current(previous_); } + ScopedCudaContext(const ScopedCudaContext&) = delete; + ScopedCudaContext& operator=(const ScopedCudaContext&) = delete; + + private: + void* previous_; +}; + std::string WarpErrorString() { const char* error = wp_get_error_string(); return error == nullptr || error[0] == '\0' ? "unknown Warp error" : error; } -void CheckCuda(int status, const std::string& what) { - if (status != kCudaSuccess) { +void CheckCuda(cudaError_t status, const std::string& what) { + if (status != cudaSuccess) { const char* msg = cudaGetErrorString(status); throw std::runtime_error(what + ": CUDA error " + std::to_string(status) + (msg != nullptr ? std::string(" (") + msg + ")" : std::string())); @@ -41,14 +48,25 @@ void ValidateParamSize(APICGraph* graph, const std::string& name, std::size_t by } } +void* ParamData(APICGraph* graph, const std::string& name) { + void* data = wp_apic_get_param_ptr(graph, name.c_str()); + if (data == nullptr) { + throw std::runtime_error("Missing APIC param pointer: " + name); + } + return data; +} + } // namespace WarpApicRunner::WarpApicRunner(RuntimeMetadata metadata) : metadata_(std::move(metadata)) {} WarpApicRunner::~WarpApicRunner() { if (graph_ != nullptr) { + const ScopedCudaContext context(cuda_context_); + // A caller can release a runner after enqueueing its last CUDA invocation. + cudaEventSynchronize(completed_); wp_apic_destroy_graph(graph_); - graph_ = nullptr; + cudaEventDestroy(completed_); } } @@ -76,27 +94,43 @@ void WarpApicRunner::LoadOnce(const std::uint8_t* bundle_data, std::size_t bundl throw std::runtime_error("Warp APIC bundle input is empty"); } EnsureCudaInitialized(); - temp_dir_ = std::make_unique(); - ExtractWRPBToDirectory(bundle_data, bundle_size, temp_dir_->path()); - const std::filesystem::path wrp_path = temp_dir_->path() / kWrpFilename; - wp_cuda_context_set_current(cuda_context_); - graph_ = wp_apic_load_graph(cuda_context_, wrp_path.string().c_str(), APIC_DEVICE_CUDA); - if (graph_ == nullptr) { + const ScopedCudaContext context(cuda_context_); + auto directory = std::make_unique(); + ExtractWRPBToDirectory(bundle_data, bundle_size, directory->path()); + const std::filesystem::path wrp_path = directory->path() / kWrpFilename; + std::unique_ptr graph( + wp_apic_load_graph(cuda_context_, wrp_path.string().c_str(), APIC_DEVICE_CUDA), + wp_apic_destroy_graph); + if (graph == nullptr) { throw std::runtime_error("Failed to load embedded APIC graph '" + std::string(kWrpFilename) + "': " + WarpErrorString()); } for (const auto& spec : metadata_.inputs) { - ValidateParamSize(graph_, spec.param_name, spec.num_bytes); + ValidateParamSize(graph.get(), spec.param_name, spec.num_bytes); } for (const auto& spec : metadata_.outputs) { if (spec.mask) { - ValidateParamSize(graph_, spec.param_name, spec.num_bytes); + ValidateParamSize(graph.get(), spec.param_name, spec.num_bytes); } } + CheckCuda(cudaEventCreateWithFlags(&completed_, cudaEventDisableTiming), + "create Warp completion event"); + temp_dir_ = std::move(directory); + graph_ = graph.release(); }); } +void WarpApicRunner::ValidateTensor(const TensorView& tensor) const { + if (tensor.num_bytes > 0 && tensor.data == nullptr) { + throw std::runtime_error("Warp tensor has no storage"); + } + if (tensor.is_cuda && tensor.device_index != 0) { + throw std::runtime_error("Warp replay currently requires CUDA device 0"); + } +} + void WarpApicRunner::ValidateInput(const TensorView& tensor, const InputSpec& spec) const { + ValidateTensor(tensor); if (tensor.dtype != spec.dtype) { throw std::runtime_error("Warp input '" + spec.param_name + "' dtype mismatch: got " + ElementTypeName(tensor.dtype) + ", expected " + ElementTypeName(spec.dtype)); @@ -107,6 +141,7 @@ void WarpApicRunner::ValidateInput(const TensorView& tensor, const InputSpec& sp } void WarpApicRunner::ValidateOutput(const TensorView& tensor, const OutputSpec& spec) const { + ValidateTensor(tensor); if (tensor.dtype != spec.dtype) { throw std::runtime_error("Warp output dtype mismatch at logical index " + std::to_string(spec.logical_index)); } @@ -128,69 +163,64 @@ void WarpApicRunner::Run(const RuntimeInvocation& invocation) { throw std::runtime_error("Warp output count mismatch"); } - bool device_resident = invocation.cuda_stream != nullptr; - for (const auto& tensor : invocation.inputs) { - device_resident = device_resident && tensor.is_cuda; - } - for (const auto& tensor : invocation.outputs) { - device_resident = device_resident && tensor.is_cuda; - } - - wp_cuda_context_set_current(cuda_context_); + bool uses_host_storage = false; for (std::size_t i = 0; i < metadata_.inputs.size(); ++i) { - const auto& spec = metadata_.inputs[i]; - const TensorView& tensor = invocation.inputs[i]; - ValidateInput(tensor, spec); - if (device_resident) { - void* param_ptr = wp_apic_get_param_ptr(graph_, spec.param_name.c_str()); - if (param_ptr == nullptr) { - throw std::runtime_error("Failed to get APIC input param pointer '" + spec.param_name + "'"); - } - CheckCuda(cudaMemcpyAsync(param_ptr, tensor.data, spec.num_bytes, - kCudaMemcpyDeviceToDevice, invocation.cuda_stream), - "copy into APIC param '" + spec.param_name + "'"); - } else { - if (!wp_apic_set_param(graph_, spec.param_name.c_str(), tensor.data, spec.num_bytes)) { - throw std::runtime_error("Failed to set APIC param '" + spec.param_name + "': " + WarpErrorString()); + ValidateInput(invocation.inputs[i], metadata_.inputs[i]); + uses_host_storage |= !invocation.inputs[i].is_cuda; + } + for (std::size_t i = 0; i < metadata_.outputs.size(); ++i) { + ValidateOutput(invocation.outputs[i], metadata_.outputs[i]); + uses_host_storage |= !invocation.outputs[i].is_cuda; + } + + const ScopedCudaContext context(cuda_context_); + const auto stream = static_cast(invocation.cuda_stream); + try { + // Protect graph-owned buffers even when consecutive calls use new streams. + CheckCuda(cudaStreamWaitEvent(stream, completed_, 0), "wait for Warp replay"); + for (std::size_t i = 0; i < metadata_.inputs.size(); ++i) { + const auto& spec = metadata_.inputs[i]; + const TensorView& tensor = invocation.inputs[i]; + if (spec.num_bytes != 0) { + CheckCuda(cudaMemcpyAsync(ParamData(graph_, spec.param_name), tensor.data, spec.num_bytes, + tensor.is_cuda ? cudaMemcpyDeviceToDevice : cudaMemcpyHostToDevice, + stream), + "copy into APIC param '" + spec.param_name + "'"); } } - } - if (!wp_apic_launch(graph_, device_resident ? invocation.cuda_stream : nullptr)) { - throw std::runtime_error("Failed to launch APIC graph: " + WarpErrorString()); - } - if (!device_resident) { - wp_cuda_context_synchronize(cuda_context_); - } + if (!wp_apic_launch(graph_, stream)) { + throw std::runtime_error("Failed to launch APIC graph: " + WarpErrorString()); + } - for (std::size_t i = 0; i < metadata_.outputs.size(); ++i) { - const auto& spec = metadata_.outputs[i]; - const TensorView& tensor = invocation.outputs[i]; - ValidateOutput(tensor, spec); - if (!spec.mask || spec.num_bytes == 0) { - if (tensor.data != nullptr && tensor.num_bytes > 0) { - if (device_resident) { - CheckCuda(cudaMemsetAsync(tensor.data, 0, tensor.num_bytes, invocation.cuda_stream), + for (std::size_t i = 0; i < metadata_.outputs.size(); ++i) { + const auto& spec = metadata_.outputs[i]; + const TensorView& tensor = invocation.outputs[i]; + if (spec.num_bytes == 0) { + continue; + } + if (!spec.mask) { + if (tensor.is_cuda) { + CheckCuda(cudaMemsetAsync(tensor.data, 0, spec.num_bytes, stream), "fill unused Warp output"); } else { - std::memset(tensor.data, 0, tensor.num_bytes); + std::memset(tensor.data, 0, spec.num_bytes); } + } else { + CheckCuda(cudaMemcpyAsync(tensor.data, ParamData(graph_, spec.param_name), spec.num_bytes, + tensor.is_cuda ? cudaMemcpyDeviceToDevice : cudaMemcpyDeviceToHost, + stream), + "copy from APIC param '" + spec.param_name + "'"); } - continue; } - if (device_resident) { - void* param_ptr = wp_apic_get_param_ptr(graph_, spec.param_name.c_str()); - if (param_ptr == nullptr) { - throw std::runtime_error("Failed to get APIC output param pointer '" + spec.param_name + "'"); - } - CheckCuda(cudaMemcpyAsync(tensor.data, param_ptr, spec.num_bytes, - kCudaMemcpyDeviceToDevice, invocation.cuda_stream), - "copy from APIC param '" + spec.param_name + "'"); - } else { - if (!wp_apic_get_param(graph_, spec.param_name.c_str(), tensor.data, spec.num_bytes)) { - throw std::runtime_error("Failed to get APIC param '" + spec.param_name + "': " + WarpErrorString()); - } + CheckCuda(cudaEventRecord(completed_, stream), "record Warp completion"); + if (uses_host_storage) { + CheckCuda(cudaStreamSynchronize(stream), "finish Warp host copies"); } + } catch (...) { + // Finish partial submissions before callers release their tensor storage. + cudaStreamSynchronize(stream); + throw; } } diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.h b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.h index 634f6612..1bf6affa 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.h +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/core/warp_apic_runner.h @@ -10,6 +10,8 @@ #include #include +#include + #include "apic.h" namespace leapp::warp_runtime { @@ -29,6 +31,7 @@ class WarpApicRunner { private: void EnsureCudaInitialized(); + void ValidateTensor(const TensorView& tensor) const; void ValidateInput(const TensorView& tensor, const InputSpec& spec) const; void ValidateOutput(const TensorView& tensor, const OutputSpec& spec) const; @@ -38,6 +41,7 @@ class WarpApicRunner { std::mutex run_mutex_; void* cuda_context_ = nullptr; APICGraph* graph_ = nullptr; + cudaEvent_t completed_ = nullptr; }; } // namespace leapp::warp_runtime diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/onnx/ort_wrp_runner_op.cc b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/onnx/ort_wrp_runner_op.cc index a4403949..06820711 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/onnx/ort_wrp_runner_op.cc +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/onnx/ort_wrp_runner_op.cc @@ -1,10 +1,14 @@ #include #include #include +#include #include #include +#include #include +#include + #include "onnxruntime_c_api.h" #include "../core/runtime_metadata.h" @@ -13,30 +17,34 @@ namespace leapp::warp_runtime::onnx_adapter { namespace { -const OrtApi* g_ort = nullptr; constexpr const char* kDomain = "com.nvidia.warp"; constexpr const char* kOpName = "WrpRunner"; -void CheckOrt(OrtStatus* status) { +void CheckOrt(const OrtApi& api, OrtStatus* status) { if (status == nullptr) { return; } - std::string message = g_ort->GetErrorMessage(status); - g_ort->ReleaseStatus(status); - throw std::runtime_error(message); + const std::unique_ptr owner(status, api.ReleaseStatus); + throw std::runtime_error(api.GetErrorMessage(status)); +} + +void CheckCuda(cudaError_t status) { + if (status != cudaSuccess) { + throw std::runtime_error(cudaGetErrorString(status)); + } } -std::string GetStringAttribute(const OrtKernelInfo* info, const char* name) { +std::string GetStringAttribute(const OrtApi& api, const OrtKernelInfo* info, const char* name) { std::size_t size = 0; - OrtStatus* status = g_ort->KernelInfoGetAttribute_string(info, name, nullptr, &size); + OrtStatus* status = api.KernelInfoGetAttribute_string(info, name, nullptr, &size); if (status != nullptr) { - g_ort->ReleaseStatus(status); + api.ReleaseStatus(status); } if (size == 0) { throw std::runtime_error(std::string("Missing string attribute: ") + name); } std::string value(size, '\0'); - CheckOrt(g_ort->KernelInfoGetAttribute_string(info, name, value.data(), &size)); + CheckOrt(api, api.KernelInfoGetAttribute_string(info, name, value.data(), &size)); if (!value.empty() && value.back() == '\0') { value.pop_back(); } @@ -75,40 +83,49 @@ ONNXTensorElementDataType ToOrtType(ElementType dtype) { } } -TensorView TensorViewFromOrt(const OrtValue* value, bool is_cuda) { +TensorView TensorViewFromOrt(const OrtApi& api, const OrtValue* value) { OrtTensorTypeAndShapeInfo* shape_info = nullptr; - CheckOrt(g_ort->GetTensorTypeAndShape(value, &shape_info)); + CheckOrt(api, api.GetTensorTypeAndShape(value, &shape_info)); + const std::unique_ptr + owner(shape_info, api.ReleaseTensorTypeAndShapeInfo); ONNXTensorElementDataType ort_dtype = ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; - CheckOrt(g_ort->GetTensorElementType(shape_info, &ort_dtype)); + CheckOrt(api, api.GetTensorElementType(shape_info, &ort_dtype)); std::size_t rank = 0; - CheckOrt(g_ort->GetDimensionsCount(shape_info, &rank)); + CheckOrt(api, api.GetDimensionsCount(shape_info, &rank)); std::vector dims(rank); - CheckOrt(g_ort->GetDimensions(shape_info, dims.data(), rank)); + CheckOrt(api, api.GetDimensions(shape_info, dims.data(), rank)); std::size_t elem_count = 0; - CheckOrt(g_ort->GetTensorShapeElementCount(shape_info, &elem_count)); - g_ort->ReleaseTensorTypeAndShapeInfo(shape_info); + CheckOrt(api, api.GetTensorShapeElementCount(shape_info, &elem_count)); void* data = nullptr; - CheckOrt(g_ort->GetTensorMutableData(const_cast(value), &data)); + CheckOrt(api, api.GetTensorMutableData(const_cast(value), &data)); TensorView view; view.data = data; view.dtype = FromOrtType(ort_dtype); view.shape = std::move(dims); view.num_bytes = elem_count * ElementSize(view.dtype); - view.is_cuda = is_cuda; + const OrtMemoryInfo* memory_info = nullptr; + CheckOrt(api, api.GetTensorMemoryInfo(value, &memory_info)); + OrtMemoryInfoDeviceType device_type; + api.MemoryInfoGetDeviceType(memory_info, &device_type); + view.is_cuda = device_type == OrtMemoryInfoDeviceType_GPU; + CheckOrt(api, api.MemoryInfoGetId(memory_info, &view.device_index)); + if (view.is_cuda && view.device_index != 0) { + throw std::runtime_error("Warp replay currently requires CUDA device 0"); + } return view; } struct Kernel { - explicit Kernel(const OrtKernelInfo* info) - : metadata(ParseRuntimeMetadata(GetStringAttribute(info, "runtime_metadata"))), + Kernel(const OrtApi& api, const OrtKernelInfo* info) + : api(api), metadata(ParseRuntimeMetadata(GetStringAttribute(api, info, "runtime_metadata"))), runner(std::make_unique(metadata)) {} void Compute(OrtKernelContext* ctx) { std::size_t input_count = 0; std::size_t output_count = 0; - CheckOrt(g_ort->KernelContext_GetInputCount(ctx, &input_count)); - CheckOrt(g_ort->KernelContext_GetOutputCount(ctx, &output_count)); + CheckOrt(api, api.KernelContext_GetInputCount(ctx, &input_count)); + CheckOrt(api, api.KernelContext_GetOutputCount(ctx, &output_count)); if (input_count == 0) { throw std::runtime_error("WrpRunner expected at least one bundle input"); } @@ -121,72 +138,88 @@ struct Kernel { } const OrtValue* bundle = nullptr; - CheckOrt(g_ort->KernelContext_GetInput(ctx, bundle_index, &bundle)); - TensorView bundle_view = TensorViewFromOrt(bundle, false); + CheckOrt(api, api.KernelContext_GetInput(ctx, bundle_index, &bundle)); + RuntimeInvocation invocation; + CheckOrt(api, api.KernelContext_GetGPUComputeStream(ctx, &invocation.cuda_stream)); + TensorView bundle_view = TensorViewFromOrt(api, bundle); if (bundle_view.dtype != ElementType::UInt8) { throw std::runtime_error("WrpRunner bundle input must be tensor(uint8)"); } - runner->LoadOnce(static_cast(bundle_view.data), bundle_view.num_bytes); - - RuntimeInvocation invocation; + std::call_once(load_once, [&]() { + // Variadic CUDA inputs include the archive. Only loading needs it on host. + std::vector host_bundle; + const auto* data = static_cast(bundle_view.data); + if (bundle_view.is_cuda) { + host_bundle.resize(bundle_view.num_bytes); + auto stream = static_cast(invocation.cuda_stream); + CheckCuda(cudaMemcpyAsync(host_bundle.data(), data, host_bundle.size(), + cudaMemcpyDeviceToHost, stream)); + CheckCuda(cudaStreamSynchronize(stream)); + data = host_bundle.data(); + } + runner->LoadOnce(data, bundle_view.num_bytes); + }); invocation.inputs.reserve(metadata.inputs.size()); for (std::size_t i = 0; i < metadata.inputs.size(); ++i) { const OrtValue* input = nullptr; - CheckOrt(g_ort->KernelContext_GetInput(ctx, i, &input)); - invocation.inputs.push_back(TensorViewFromOrt(input, false)); + CheckOrt(api, api.KernelContext_GetInput(ctx, i, &input)); + invocation.inputs.push_back(TensorViewFromOrt(api, input)); } for (std::size_t i = 0; i < metadata.outputs.size(); ++i) { const auto& spec = metadata.outputs[i]; OrtValue* output = nullptr; - CheckOrt(g_ort->KernelContext_GetOutput( + CheckOrt(api, api.KernelContext_GetOutput( ctx, i, spec.shape.data(), spec.shape.size(), &output)); - invocation.outputs.push_back(TensorViewFromOrt(output, false)); + invocation.outputs.push_back(TensorViewFromOrt(api, output)); } - invocation.cuda_stream = nullptr; runner->Run(invocation); } + const OrtApi& api; + std::once_flag load_once; RuntimeMetadata metadata; std::unique_ptr runner; }; -OrtStatusPtr CreateKernelV2(const OrtCustomOp*, const OrtApi*, const OrtKernelInfo* info, void** kernel) { +OrtStatusPtr CreateKernelV2(const OrtCustomOp*, const OrtApi* api, const OrtKernelInfo* info, void** kernel) { try { - *kernel = new Kernel(info); + *kernel = new Kernel(*api, info); return nullptr; } catch (const std::exception& e) { - return g_ort->CreateStatus(ORT_FAIL, e.what()); - } -} - -void* CreateKernel(const OrtCustomOp*, const OrtApi*, const OrtKernelInfo* info) { - try { - return new Kernel(info); + return api->CreateStatus(ORT_FAIL, e.what()); } catch (...) { - return nullptr; + return api->CreateStatus(ORT_FAIL, "Unknown exception creating Warp kernel"); } } void DestroyKernel(void* kernel) { delete static_cast(kernel); } -void KernelCompute(void* kernel, OrtKernelContext* context) { static_cast(kernel)->Compute(context); } OrtStatusPtr KernelComputeV2(void* kernel, OrtKernelContext* context) { + auto* value = static_cast(kernel); try { - static_cast(kernel)->Compute(context); + value->Compute(context); return nullptr; } catch (const std::exception& e) { - return g_ort->CreateStatus(ORT_FAIL, e.what()); + return value->api.CreateStatus(ORT_FAIL, e.what()); + } catch (...) { + return value->api.CreateStatus(ORT_FAIL, "Unknown exception in Warp kernel"); } } +struct WarpOp : OrtCustomOp { + const char* provider; +}; + const char* GetName(const OrtCustomOp*) { return kOpName; } -const char* GetExecutionProviderType(const OrtCustomOp*) { return nullptr; } +const char* GetExecutionProviderType(const OrtCustomOp* op) { + return static_cast(op)->provider; +} std::size_t GetInputTypeCount(const OrtCustomOp*) { return 1; } ONNXTensorElementDataType GetInputType(const OrtCustomOp*, std::size_t) { @@ -210,17 +243,16 @@ int GetVariadicOutputHomogeneity(const OrtCustomOp*) { return 0; } int GetStartVersion(const OrtCustomOp*) { return 1; } int GetEndVersion(const OrtCustomOp*) { return 1; } -OrtCustomOp CreateOp() { - OrtCustomOp op{}; +WarpOp CreateOp(const char* provider) { + WarpOp op{}; + op.provider = provider; op.version = ORT_API_VERSION; - op.CreateKernel = CreateKernel; op.GetName = GetName; op.GetExecutionProviderType = GetExecutionProviderType; op.GetInputTypeCount = GetInputTypeCount; op.GetInputType = GetInputType; op.GetOutputTypeCount = GetOutputTypeCount; op.GetOutputType = GetOutputType; - op.KernelCompute = KernelCompute; op.KernelDestroy = DestroyKernel; op.GetInputCharacteristic = GetInputCharacteristic; op.GetOutputCharacteristic = GetOutputCharacteristic; @@ -236,7 +268,27 @@ OrtCustomOp CreateOp() { return op; } -OrtCustomOp g_op = CreateOp(); +OrtCustomOpDomain* GetDomain(const OrtApi& api) { + static const WarpOp cpu_op = CreateOp(nullptr); + static const WarpOp cuda_op = CreateOp("CUDAExecutionProvider"); + using Domain = std::unique_ptr; + // Domains must outlive every session using them; reuse one per runtime API. + static std::vector> domains; + static std::mutex mutex; + std::lock_guard guard(mutex); + for (const auto& entry : domains) { + if (entry.first == &api) { + return entry.second.get(); + } + } + OrtCustomOpDomain* value = nullptr; + CheckOrt(api, api.CreateCustomOpDomain(kDomain, &value)); + Domain domain(value, api.ReleaseCustomOpDomain); + CheckOrt(api, api.CustomOpDomain_Add(value, &cpu_op)); + CheckOrt(api, api.CustomOpDomain_Add(value, &cuda_op)); + domains.emplace_back(&api, std::move(domain)); + return value; +} } // namespace } // namespace leapp::warp_runtime::onnx_adapter @@ -244,15 +296,16 @@ OrtCustomOp g_op = CreateOp(); extern "C" OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api_base) { using namespace leapp::warp_runtime::onnx_adapter; - g_ort = api_base->GetApi(ORT_API_VERSION); - OrtCustomOpDomain* domain = nullptr; - OrtStatus* status = g_ort->CreateCustomOpDomain(kDomain, &domain); - if (status != nullptr) { - return status; + const auto* api = api_base->GetApi(ORT_API_VERSION); + if (api == nullptr) { + return api_base->GetApi(1)->CreateStatus( + ORT_FAIL, "Warp adapter was built for a newer ONNX Runtime API"); } - status = g_ort->CustomOpDomain_Add(domain, &g_op); - if (status != nullptr) { - return status; + try { + return api->AddCustomOpDomain(options, GetDomain(*api)); + } catch (const std::exception& e) { + return api->CreateStatus(ORT_FAIL, e.what()); + } catch (...) { + return api->CreateStatus(ORT_FAIL, "Unknown exception registering Warp ops"); } - return g_ort->AddCustomOpDomain(options, domain); } diff --git a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/torch/torch_warp_runner_op.cc b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/torch/torch_warp_runner_op.cc index 67a55a28..59589f60 100644 --- a/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/torch/torch_warp_runner_op.cc +++ b/leapp/leapp_graph/custom_operator_registry/warp_operator/runtime/torch/torch_warp_runner_op.cc @@ -3,9 +3,13 @@ #include #include #include +#include +#include #include -#include +#include +#include +#include #include #include "../core/runtime_metadata.h" @@ -59,29 +63,41 @@ TensorView FromTensor(const at::Tensor& tensor) { struct RunnerCacheEntry { RuntimeMetadata metadata; - at::Tensor bundle; + c10::weak_intrusive_ptr bundle_storage; std::unique_ptr runner; }; std::mutex g_cache_mutex; std::unordered_map> g_cache; +// Called with g_cache_mutex held. Model constants own their storage. +void PruneRunners() { + for (auto it = g_cache.begin(); it != g_cache.end();) { + if (it->second->bundle_storage.expired()) { + it = g_cache.erase(it); + } else { + ++it; + } + } +} + std::shared_ptr GetRunner(const std::string& runtime_metadata, const at::Tensor& bundle) { const std::string cache_key = runtime_metadata + ":" + std::to_string(reinterpret_cast(bundle.data_ptr())); std::lock_guard guard(g_cache_mutex); + PruneRunners(); auto it = g_cache.find(cache_key); if (it != g_cache.end()) { return it->second; } - auto entry = std::make_shared(); - entry->metadata = ParseRuntimeMetadata(runtime_metadata); - entry->bundle = bundle; - entry->runner = std::make_unique(entry->metadata); + auto metadata = ParseRuntimeMetadata(runtime_metadata); + auto runner = std::make_unique(metadata); const at::Tensor cpu_bundle = bundle.device().is_cpu() ? bundle.contiguous() : bundle.cpu().contiguous(); - entry->runner->LoadOnce(static_cast(cpu_bundle.data_ptr()), cpu_bundle.nbytes()); + runner->LoadOnce(static_cast(cpu_bundle.data_ptr()), cpu_bundle.nbytes()); + auto entry = std::make_shared(RunnerCacheEntry{ + std::move(metadata), bundle.storage().getWeakStorageImpl(), std::move(runner)}); g_cache.emplace(cache_key, entry); return entry; } @@ -89,26 +105,36 @@ std::shared_ptr GetRunner(const std::string& runtime_metadata, std::vector WarpRunner(const std::vector& inputs, const std::string& runtime_metadata, const at::Tensor& bundle) { + TORCH_CHECK(!bundle.is_cuda() || bundle.get_device() == 0, + "Warp replay currently requires CUDA device 0"); + for (const auto& input : inputs) { + TORCH_CHECK(!input.is_cuda() || input.get_device() == 0, + "Warp replay currently requires CUDA device 0"); + } + const c10::cuda::CUDAGuard device_guard(0); auto entry = GetRunner(runtime_metadata, bundle); const auto& metadata = entry->metadata; std::vector outputs; outputs.reserve(metadata.outputs.size()); - const at::Device device = inputs.empty() ? at::Device(at::kCUDA, 0) : inputs.front().device(); + const at::Device device = inputs.empty() ? bundle.device() : inputs.front().device(); for (const auto& spec : metadata.outputs) { std::vector shape = spec.shape; outputs.push_back(at::empty(shape, at::TensorOptions().dtype(ToTorchType(spec.dtype)).device(device))); } + std::vector contiguous_inputs; + contiguous_inputs.reserve(inputs.size()); RuntimeInvocation invocation; invocation.inputs.reserve(inputs.size()); for (const auto& input : inputs) { - invocation.inputs.push_back(FromTensor(input)); + contiguous_inputs.push_back(input.contiguous()); + invocation.inputs.push_back(FromTensor(contiguous_inputs.back())); } invocation.outputs.reserve(outputs.size()); for (const auto& output : outputs) { invocation.outputs.push_back(FromTensor(output)); } - invocation.cuda_stream = at::cuda::getCurrentCUDAStream().stream(); + invocation.cuda_stream = c10::cuda::getCurrentCUDAStream(0).stream(); entry->runner->Run(invocation); return outputs; } @@ -116,6 +142,28 @@ std::vector WarpRunner(const std::vector& inputs, } // namespace } // namespace leapp::warp_runtime::torch_adapter +TORCH_LIBRARY_FRAGMENT(leapp, m) { + // Python may have already defined the schema while exporting a model. + if (!c10::Dispatcher::singleton().findSchema({"leapp::warp_runner", ""})) { + m.def("warp_runner(Tensor[] inputs, str runtime_metadata, Tensor bundle) -> Tensor[]"); + } +} + +TORCH_LIBRARY_IMPL(leapp, CPU, m) { + m.impl("warp_runner", leapp::warp_runtime::torch_adapter::WarpRunner); +} + TORCH_LIBRARY_IMPL(leapp, CUDA, m) { m.impl("warp_runner", leapp::warp_runtime::torch_adapter::WarpRunner); } + +// Call after releasing model constants to reclaim unused graphs immediately. +extern "C" int LEAPP_WARP_ReleaseUnusedRunners() noexcept { + try { + std::lock_guard guard(leapp::warp_runtime::torch_adapter::g_cache_mutex); + leapp::warp_runtime::torch_adapter::PruneRunners(); + return 0; + } catch (...) { + return 1; + } +} diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index cd29d550..1073a28a 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -40,6 +40,10 @@ explicit unsafe borrows only for private initial state and feedback descriptors. The `bindings/common` crate adds one lock per session and explicit closure for Python and C++. The Rust core and both bindings require Linux and the native SDKs. +Embedded Warp operators in ONNX/PT2 bundles use optional native adapters installed +beside their Triton backends. See [Warp operators](src/engines/triton/backends/README.md#warp-operators) +for SDK requirements, installation, and integration tests. + ## Rust ```rust diff --git a/packages/leapp-runtime-rust/src/engines/triton.rs b/packages/leapp-runtime-rust/src/engines/triton.rs index bf9dd514..f82ba667 100644 --- a/packages/leapp-runtime-rust/src/engines/triton.rs +++ b/packages/leapp-runtime-rust/src/engines/triton.rs @@ -50,13 +50,16 @@ pub(super) fn open( for location in options.output_placement.overrides.values() { location.validate()?; } - let pipeline = convert::convert(bundle, options.device, graph) + let installation = Options::default(); + let mut pipeline = convert::convert(bundle, options.device, graph) .map_err(|e| crate::RuntimeError::new("convert", e))?; + convert::configure_plugins(&mut pipeline, &installation.backend_directory) + .map_err(|e| crate::RuntimeError::new("load", e))?; let repository = convert::repository::ModelRepository::create(&pipeline)?; Ok(std::sync::Arc::new(TritonEngine::load( &pipeline, repository, - &Options::default(), + &installation, )?)) } diff --git a/packages/leapp-runtime-rust/src/engines/triton/backends/README.md b/packages/leapp-runtime-rust/src/engines/triton/backends/README.md index 369acde7..8e3f8ae6 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/README.md +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/README.md @@ -56,7 +56,8 @@ loader, for example with `LD_LIBRARY_PATH`. For integration tests, also make the installed ONNX Runtime backend available: ```sh -ln -s "$TRITON_ROOT/backends/onnxruntime" target/plugins/backends/onnxruntime +mkdir -p target/plugins/backends/onnxruntime +ln -s "$TRITON_ROOT/backends/onnxruntime/"* target/plugins/backends/onnxruntime/ export LEAPP_TRITON_BACKENDS="$PWD/target/plugins/backends" python tests/generators/generate_models.py --output target/triton-fixtures python tests/generators/generate_pt2_models.py --output target/pt2-fixtures @@ -70,6 +71,59 @@ dynamic and empty tensors, error recovery, integer tensors, incompatible archive devices, and ONNX/PT2 ensembles. `LEAPP_TRITON_FIXTURES` and `LEAPP_PT2_FIXTURES` override fixture directories. +## Warp operators + +Build LEAPP's native Warp adapters with PyTorch matching NativeRT. Select the +ONNX Runtime header version shipped by Triton's ONNX backend, for example: + +```sh +leapp-build-warp-runtime --build-dir /path/to/warp-build -- \ + -DLEAPP_WARP_ONNXRUNTIME_VERSION=1.24.4 -DCMAKE_INSTALL_LIBDIR=lib +``` + +Use your installed backend's version; omitting the setting uses Python's ONNX +Runtime version. Alternatively, supply `-DORT_INCLUDE_DIR=/path/to/headers`. +`-DWARP_ROOT`, `-DTorch_DIR`, and `-DCUDAToolkit_ROOT` select existing SDKs; +the build does not install them. Keep Warp's export and runtime releases aligned. +Replay requires CUDA even for CPU model instances. +The runner supports CUDA device 0 only. ONNX GPU models use a CUDA custom-op kernel, +keeping replay tensors on the GPU; CPU models retain host staging. The embedded +archive is copied to host once when loading. Warm each model serially before +concurrent inference. +PT2 model unload releases expired Warp cache entries while preserving live models. + +Install the adapters and link them beside the backends under +`LEAPP_TRITON_BACKENDS`. Use local backend directories as above so the links do +not modify the Triton SDK: + +```sh +cmake --install /path/to/warp-build --prefix "$PWD/target/plugins" +ln -s ../../lib/libleapp_wrp_onnx_custom_op.so \ + target/plugins/backends/onnxruntime/libleapp_wrp_onnx_custom_op.so +ln -s ../../lib/libleapp_wrp_torch_custom_op.so \ + target/plugins/backends/leapp_pt2/libleapp_wrp_torch_custom_op.so +``` + +Make the selected Warp `bin/`, LibTorch `lib/`, and CUDA library directories +available to the loader, alongside NativeRT and Triton, through `LD_LIBRARY_PATH` +or an explicitly configured `CMAKE_INSTALL_RPATH`. The installation contains the +LEAPP adapters; the SDK libraries stay in their existing locations. + +Session construction discovers adapters for the backends used by the bundle. +Absent adapters leave ordinary models unchanged; Warp models need their adapter +installed. An invalid installed adapter causes a load error, even for ordinary +models using that backend. PT2 loads the library before creating model instances +and retains its handle until those instances are destroyed. + +Generate and run the optional integration fixtures from this package directory: + +```sh +python tests/generators/generate_warp_models.py --output target/warp-fixtures +export LEAPP_WARP_FIXTURES="$PWD/target/warp-fixtures" +export WARP_APIC_TMPDIR="$(mktemp -d)" +cargo test --locked engines::triton::tests::warp -- --ignored --test-threads=1 +``` + ## TorchScript `torch/jit` executes scripted and traced modules with positional tensor arguments diff --git a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.cc b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.cc index 39a7285d..53070758 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.cc +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.cc @@ -12,7 +12,8 @@ namespace leapp::internal { -Status ReadModelConfig(TRITONBACKEND_Model* model, ModelConfig* state) { +Status ReadModelConfig(TRITONBACKEND_Model* model, ModelConfig* state, + std::filesystem::path* custom_op_library) { constexpr uint32_t kConfigVersion = 1; TRITONSERVER_Message* message = nullptr; @@ -32,6 +33,11 @@ Status ReadModelConfig(TRITONBACKEND_Model* model, ModelConfig* state) { return status; } const auto config = nlohmann::json::parse(json, json + size); + if (custom_op_library != nullptr) { + *custom_op_library = config.value("parameters", nlohmann::json::object()) + .value("custom_op_library", nlohmann::json::object()) + .value("string_value", ""); + } if (config.value("max_batch_size", 0) != 0 || config.value("model_transaction_policy", nlohmann::json::object()) .value("decoupled", false)) { diff --git a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.h b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.h index cf7bcd5e..51fd541b 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.h +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/common/src/model_config.h @@ -26,7 +26,8 @@ struct ModelConfig { std::vector outputs; }; -Status ReadModelConfig(TRITONBACKEND_Model* model, ModelConfig* config); +Status ReadModelConfig(TRITONBACKEND_Model* model, ModelConfig* config, + std::filesystem::path* custom_op_library = nullptr); } // namespace leapp::internal diff --git a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/CMakeLists.txt b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/CMakeLists.txt index 01b68d05..c42cdfa8 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/CMakeLists.txt +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/CMakeLists.txt @@ -43,7 +43,7 @@ add_library(triton_leapp_pt2 SHARED target_include_directories(triton_leapp_pt2 PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}") target_link_libraries(triton_leapp_pt2 PRIVATE - leapp_backend_common "${NATIVERT_LIBRARY}") + leapp_backend_common "${NATIVERT_LIBRARY}" ${CMAKE_DL_LIBS}) target_compile_options(triton_leapp_pt2 PRIVATE -Wall -Wextra -Wpedantic) target_link_options(triton_leapp_pt2 PRIVATE "LINKER:--no-undefined") set_target_properties(triton_leapp_pt2 PROPERTIES diff --git a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/backend.cc b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/backend.cc index 3b3b7eb4..879dde2e 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/backend.cc +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/backend.cc @@ -1,8 +1,11 @@ // Copyright 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +#include + #include #include +#include #include #include #include @@ -29,10 +32,22 @@ TRITONSERVER_Error* TRITONBACKEND_Initialize(TRITONBACKEND_Backend*) { TRITONSERVER_Error* TRITONBACKEND_ModelInitialize(TRITONBACKEND_Model* model) { try { auto state = std::make_unique(); - auto status = leapp::internal::ReadModelConfig(model, state.get()); + std::filesystem::path custom_op_library; + auto status = + leapp::internal::ReadModelConfig(model, state.get(), &custom_op_library); if (!status.ok()) { return leapp::internal::NewBackendError(status); } + if (!custom_op_library.empty()) { + void* library = dlopen(custom_op_library.c_str(), RTLD_NOW | RTLD_LOCAL); + if (library == nullptr) { + const auto message = "cannot load custom-op library '" + + custom_op_library.string() + "': " + dlerror(); + return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_UNAVAILABLE, + message.c_str()); + } + state->custom_op_library = std::shared_ptr(library, dlclose); + } auto* transferred = state.release(); if (auto* error = TRITONBACKEND_ModelSetState(model, transferred)) { state.reset(transferred); @@ -52,7 +67,23 @@ TRITONSERVER_Error* TRITONBACKEND_ModelFinalize(TRITONBACKEND_Model* model) { if (auto* error = TRITONBACKEND_ModelState(model, &state)) { return error; } - delete static_cast(state); + auto* model_state = static_cast(state); + if (model_state == nullptr) { + return nullptr; + } + using ReleaseUnusedRunners = int (*)(); + auto release_unused = model_state->custom_op_library + ? reinterpret_cast(dlsym( + model_state->custom_op_library.get(), "LEAPP_WARP_ReleaseUnusedRunners")) + : nullptr; + // Instance teardown has released model constants; keep the library loaded + // until its optional Warp cache cleanup has finished. + const int cleanup_status = release_unused ? release_unused() : 0; + delete model_state; + if (cleanup_status != 0) { + return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, + "failed to release unused Warp runners"); + } return nullptr; } diff --git a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/native_rt_instance.h b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/native_rt_instance.h index 77691a05..e2be7cec 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/native_rt_instance.h +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/torch/pt2/src/native_rt_instance.h @@ -23,7 +23,10 @@ namespace leapp::internal { -struct NativeRtModel : ModelConfig {}; +struct NativeRtModel : ModelConfig { + // Keep dispatcher registrations loaded until all model instances are gone. + std::shared_ptr custom_op_library; +}; // Triton owns model state until all its instances have been destroyed. class NativeRtInstance { diff --git a/packages/leapp-runtime-rust/src/engines/triton/convert.rs b/packages/leapp-runtime-rust/src/engines/triton/convert.rs index 6d2bb794..a42e7829 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/convert.rs +++ b/packages/leapp-runtime-rust/src/engines/triton/convert.rs @@ -18,7 +18,7 @@ use crate::{ session::graph::{port_ref, require, Graph as ModelPorts}, Bundle, ExecutionDevice, PortRef, PortSpec, }; -use anyhow::Result; +use anyhow::{Context, Result}; use proto::model_config::{model_ensembling::Step, ModelConfig, ModelEnsembling}; use std::{ collections::{HashMap, HashSet}, @@ -27,6 +27,60 @@ use std::{ const MODEL_VERSION: i64 = 1; +pub(super) fn configure_plugins( + pipeline: &mut TritonPipeline, + backend_directory: &std::path::Path, +) -> Result<()> { + for (backend, library) in [ + ("onnxruntime", "libleapp_wrp_onnx_custom_op.so"), + ("leapp_pt2", "libleapp_wrp_torch_custom_op.so"), + ] { + if !pipeline.models.iter().any(|m| m.config.backend == backend) { + continue; + } + let path = backend_directory.join(backend).join(library); + match std::fs::symlink_metadata(&path) { + Err(e) if e.kind() == std::io::ErrorKind::NotFound => continue, + result => result.with_context(|| format!("inspect Warp adapter {}", path.display()))?, + }; + let resolved = path + .canonicalize() + .with_context(|| format!("resolve Warp adapter {}", path.display()))?; + require( + resolved.is_file(), + format!("Warp adapter is not a file: {}", path.display()), + )?; + std::fs::File::open(&resolved) + .with_context(|| format!("read Warp adapter {}", path.display()))?; + let filename = resolved + .to_str() + .context("Warp adapter path must be UTF-8")?; + for model in pipeline + .models + .iter_mut() + .filter(|m| m.config.backend == backend) + { + if backend == "onnxruntime" { + model + .config + .model_operations + .mut_or_insert_default() + .op_library_filename + .push(filename.into()); + } else { + model.config.parameters.insert( + "custom_op_library".into(), + proto::model_config::ModelParameter { + string_value: filename.into(), + ..Default::default() + }, + ); + } + } + } + Ok(()) +} + #[derive(Debug)] pub(crate) struct TritonModelSpec { pub config: ModelConfig, diff --git a/packages/leapp-runtime-rust/tests/functional_tests/triton.rs b/packages/leapp-runtime-rust/tests/functional_tests/triton.rs index aafcea0b..daa21e43 100644 --- a/packages/leapp-runtime-rust/tests/functional_tests/triton.rs +++ b/packages/leapp-runtime-rust/tests/functional_tests/triton.rs @@ -13,4 +13,6 @@ pub(crate) mod support; mod tensorrt; #[path = "triton/torchscript.rs"] mod torchscript; +#[path = "triton/warp.rs"] +mod warp; // state.rs is included under Session to inspect its private feedback state. diff --git a/packages/leapp-runtime-rust/tests/functional_tests/triton/warp.rs b/packages/leapp-runtime-rust/tests/functional_tests/triton/warp.rs new file mode 100644 index 00000000..610d5426 --- /dev/null +++ b/packages/leapp-runtime-rust/tests/functional_tests/triton/warp.rs @@ -0,0 +1,110 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::support::{placed, TestTensor}; +use crate::{Bundle, ExecutionDevice, MemoryLocation, OutputPlacement, Session, SessionOptions}; +use std::path::PathBuf; + +#[test] +#[ignore = "requires Warp adapters, CUDA, and generated Warp fixtures"] +fn warp_plugins_execute_through_sessions() { + let root = std::env::var_os("LEAPP_WARP_FIXTURES") + .map(PathBuf::from) + .expect("set LEAPP_WARP_FIXTURES after running generate_warp_models.py"); + let extract_root = std::env::var_os("WARP_APIC_TMPDIR") + .expect("set WARP_APIC_TMPDIR to an empty directory for lifecycle checks"); + let graph_count = || std::fs::read_dir(&extract_root).unwrap().count(); + for (fixture, device) in [ + ("warp_onnx", ExecutionDevice::Cpu), + ("warp_onnx", ExecutionDevice::Cuda { device_id: 0 }), + ("warp_pt2_cpu", ExecutionDevice::Cpu), + ("warp_pt2_cuda", ExecutionDevice::Cuda { device_id: 0 }), + ] { + let bundle = Bundle::open(root.join(format!("{fixture}.yaml"))).unwrap(); + for location in [ + MemoryLocation::Cpu, + MemoryLocation::PinnedCpu, + MemoryLocation::Cuda { device_id: 0 }, + ] { + let graphs_before = graph_count(); + let mut session = Session::open( + &bundle, + SessionOptions { + device, + output_placement: OutputPlacement { + default_location: location, + ..Default::default() + }, + }, + ) + .unwrap(); + let x = placed( + &(0..12).map(|i| i as f32).collect::>(), + &[4, 3], + location, + ) + .unwrap(); + let gain = placed(&[1_i32, 2, 3, 4], &[4], location).unwrap(); + let run = |session: &mut Session| { + session + .run(&[("model/x", x.view()), ("model/gain", gain.view())]) + .unwrap() + }; + let expected: Vec<_> = (0..12) + .map(|i| i as f32 * (i / 3 + 1) as f32 + 4.0) + .collect(); + assert_eq!( + run(&mut session) + .get("model/result") + .unwrap() + .to_cpu_vec::() + .unwrap(), + expected, + "{fixture} {device:?} {location:?}" + ); + assert_eq!(graph_count(), graphs_before + 1); + let mut peer = session.new_session().unwrap(); + drop(session); + assert_eq!(graph_count(), graphs_before + 1); + let result = run(&mut peer); + drop(peer); + let output = result.get("model/result").unwrap(); + assert_eq!(output.location(), location); + assert_eq!(output.to_cpu_vec::().unwrap(), expected); + drop(result); + assert_eq!(graph_count(), graphs_before); + } + } +} + +#[test] +#[ignore = "requires Warp adapters, CUDA, and generated Warp fixtures"] +fn warp_plugin_load_errors_are_reported() { + use crate::engines::triton::{convert, Options, TritonEngine}; + let root = PathBuf::from(std::env::var_os("LEAPP_WARP_FIXTURES").unwrap()); + for (fixture, backend, filename) in [ + ("warp_onnx", "onnxruntime", "libleapp_wrp_onnx_custom_op.so"), + ( + "warp_pt2_cpu", + "leapp_pt2", + "libleapp_wrp_torch_custom_op.so", + ), + ] { + let bundle = Bundle::open(root.join(format!("{fixture}.yaml"))).unwrap(); + let mut pipeline = convert::convert_bundle(&bundle, ExecutionDevice::Cpu).unwrap(); + let installation = tempfile::tempdir().unwrap(); + std::fs::create_dir(installation.path().join(backend)).unwrap(); + std::fs::write( + installation.path().join(backend).join(filename), + b"invalid library", + ) + .unwrap(); + convert::configure_plugins(&mut pipeline, installation.path()).unwrap(); + let repository = convert::repository::ModelRepository::create(&pipeline).unwrap(); + let error = TritonEngine::load(&pipeline, repository, &Options::default()) + .err() + .unwrap(); + assert_eq!(error.phase(), "load"); + assert!(error.to_string().contains("failed to load"), "{error}"); + } +} diff --git a/packages/leapp-runtime-rust/tests/generators/generate_warp_models.py b/packages/leapp-runtime-rust/tests/generators/generate_warp_models.py new file mode 100644 index 00000000..3ae07b26 --- /dev/null +++ b/packages/leapp-runtime-rust/tests/generators/generate_warp_models.py @@ -0,0 +1,123 @@ +# Copyright 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Capture one Warp graph and export ONNX/PT2 bundles for plugin integration.""" + +import argparse +import json +from pathlib import Path + +import torch +import warp as wp + +from leapp.leapp_graph.custom_operator_registry.warp_operator.bundle import pack_bundle +from leapp.leapp_graph.custom_operator_registry.warp_operator.metadata import ( + build_runtime_metadata, + encode_runtime_metadata, +) +from leapp.leapp_graph.datatypes.warp.warp_segment import WarpTensorRef + + +@wp.kernel +def transform( + x: wp.array(dtype=wp.vec3), + gain: wp.array(dtype=wp.int32), + y: wp.array(dtype=wp.vec3), +): + i = wp.tid() + y[i] = x[i] * float(gain[i]) + wp.vec3(3.0) + + +class WarpModel(torch.nn.Module): + def __init__(self, archive, metadata): + super().__init__() + self.register_buffer("archive", torch.tensor(list(archive), dtype=torch.uint8)) + self.metadata = metadata + + def forward(self, x, gain): + return ( + torch.ops.leapp.warp_runner([x, gain], self.metadata, self.archive)[0] + 1 + ) + + +def write_bundle(root, name, backend, model_path): + def port(name, dtype, shape): + return {"name": name, "dtype": dtype, "shape": shape, "type": "tensor"} + + bundle = { + "system information": {"leapp config version": "1.3"}, + "models": { + "model": { + "inputs": [port("x", "float32", [4, 3]), port("gain", "int32", [4])], + "outputs": [port("result", "float32", [4, 3])], + "parameters": {"backend": backend, "model_path": str(model_path)}, + } + }, + "pipeline": { + "data_flow": {}, + "feedback_flow": {}, + "inputs": {"model": ["x", "gain"]}, + "outputs": {"model": ["result"]}, + "configs": {}, + }, + } + (root / f"{name}.yaml").write_text(json.dumps(bundle)) + + +def generate(root): + root.mkdir(parents=True, exist_ok=True) + wp.init() + x = wp.zeros(4, dtype=wp.vec3, device="cuda:0") + gain = wp.zeros(4, dtype=wp.int32, device="cuda:0") + y = wp.zeros_like(x) + wp.load_module(device="cuda:0") + with wp.ScopedCapture(device="cuda:0", apic=True) as capture: + wp.launch(transform, dim=4, inputs=[x, gain], outputs=[y], device="cuda:0") + wp.capture_save( + capture.graph, + str(root / "segment"), + inputs={"x": x, "gain": gain}, + outputs={"y": y}, + ) + archive = pack_bundle(root / "segment.wrp") + output = WarpTensorRef.from_value("y", y) + metadata = encode_runtime_metadata( + build_runtime_metadata( + segment=None, + input_refs=[ + WarpTensorRef.from_value("x", x), + WarpTensorRef.from_value("gain", gain), + ], + output_refs=[output], + output_shapes=[list(output.storage_shape)], + output_dtypes=[output.storage_dtype], + output_mask=[True], + ) + ) + for device in ("cpu", "cuda"): + args = ( + torch.zeros((4, 3), device=device), + torch.ones(4, dtype=torch.int32, device=device), + ) + program = torch.export.export(WarpModel(archive, metadata).to(device), args) + name = f"warp_pt2_{device}" + torch.export.save(program, root / f"{name}.pt2") + write_bundle(root, name, "pt2", f"{name}.pt2") + if device == "cpu": + torch.onnx.export( + program, + args, + root / "warp.onnx", + input_names=["x", "gain"], + output_names=["result"], + opset_version=18, + dynamo=True, + external_data=False, + ) + write_bundle(root, "warp_onnx", "onnx", "warp.onnx") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True, type=Path) + generate(parser.parse_args().output) diff --git a/packages/leapp-runtime-rust/tests/unit_tests/triton/convert.rs b/packages/leapp-runtime-rust/tests/unit_tests/triton/convert.rs index c580de3d..fb71395a 100644 --- a/packages/leapp-runtime-rust/tests/unit_tests/triton/convert.rs +++ b/packages/leapp-runtime-rust/tests/unit_tests/triton/convert.rs @@ -648,3 +648,74 @@ fn torchscript_ports_use_ordinals_without_changing_public_or_feedback_names() { } } } + +#[test] +fn configures_installed_warp_adapters_only_for_used_backends() { + let mut f = Fixture::new(); + f.value["models"]["a_policy"]["parameters"]["backend"] = yaml("pt2"); + let backends = tempfile::tempdir().unwrap(); + let mut pipeline = f.convert(); + configure_plugins(&mut pipeline, backends.path()).unwrap(); + assert!(pipeline.models[0].config.model_operations.is_none()); + assert!(!pipeline.models[1] + .config + .parameters + .contains_key("custom_op_library")); + for (backend, name) in [ + ("onnxruntime", "libleapp_wrp_onnx_custom_op.so"), + ("leapp_pt2", "libleapp_wrp_torch_custom_op.so"), + ] { + fs::create_dir(backends.path().join(backend)).unwrap(); + fs::write(backends.path().join(backend).join(name), b"library").unwrap(); + } + configure_plugins(&mut pipeline, backends.path()).unwrap(); + assert_eq!( + pipeline.models[0] + .config + .model_operations + .as_ref() + .unwrap() + .op_library_filename, + [backends + .path() + .join("onnxruntime/libleapp_wrp_onnx_custom_op.so") + .canonicalize() + .unwrap() + .to_str() + .unwrap()] + ); + assert_eq!( + pipeline.models[1].config.parameters["custom_op_library"].string_value, + backends + .path() + .join("leapp_pt2/libleapp_wrp_torch_custom_op.so") + .canonicalize() + .unwrap() + .to_str() + .unwrap() + ); +} + +#[test] +fn rejects_broken_warp_installations_but_ignores_unused_backends() { + let f = Fixture::new(); + let backends = tempfile::tempdir().unwrap(); + let pt2 = backends + .path() + .join("leapp_pt2/libleapp_wrp_torch_custom_op.so"); + fs::create_dir_all(&pt2).unwrap(); + configure_plugins(&mut f.convert(), backends.path()).unwrap(); + let onnx = backends + .path() + .join("onnxruntime/libleapp_wrp_onnx_custom_op.so"); + fs::create_dir_all(&onnx).unwrap(); + let error = configure_plugins(&mut f.convert(), backends.path()).unwrap_err(); + assert!(error.to_string().contains("not a file"), "{error}"); + fs::remove_dir(&onnx).unwrap(); + std::os::unix::fs::symlink("missing.so", &onnx).unwrap(); + let error = configure_plugins(&mut f.convert(), backends.path()).unwrap_err(); + assert!( + error.to_string().contains("resolve Warp adapter"), + "{error}" + ); +} diff --git a/scripts/leapp-build-warp-runtime b/scripts/leapp-build-warp-runtime index 03159a6b..020a4055 100755 --- a/scripts/leapp-build-warp-runtime +++ b/scripts/leapp-build-warp-runtime @@ -5,11 +5,11 @@ # import argparse -import importlib import os import shutil import subprocess import sys +import tempfile from pathlib import Path import leapp @@ -29,124 +29,9 @@ def _runtime_source_dir() -> Path: ) -def _require_module(name: str, install_hint: str): - try: - return importlib.import_module(name) - except ImportError as exc: - raise RuntimeError( - f"{name} is required to build LEAPP Warp support but is not " - "installed. " - f"Install it first with {install_hint}. " - "leapp-build-warp-runtime never installs dependencies." - ) from exc - - -def _find_nvcc() -> Path: - candidates = [] - if os.environ.get("CUDACXX"): - candidates.append(Path(os.environ["CUDACXX"])) - nvcc = shutil.which("nvcc") - if nvcc: - candidates.append(Path(nvcc)) - if os.environ.get("CUDA_PATH"): - candidates.append( - Path(os.environ["CUDA_PATH"]) - / "bin" - / ("nvcc.exe" if sys.platform == "win32" else "nvcc") - ) - if sys.platform != "win32": - candidates.append(Path("/usr/local/cuda/bin/nvcc")) - for candidate in candidates: - if candidate.is_file(): - return candidate - raise RuntimeError( - "The CUDA compiler nvcc is required to build LEAPP Warp support but " - "was not found. Install the CUDA Toolkit or set CUDACXX." - ) - - -def _warp_library_name() -> str: - if sys.platform == "win32": - return "warp.dll" - if sys.platform == "darwin": - return "libwarp.dylib" - return "warp.so" - - -def _onnxruntime_library_name(version: str) -> str: - if sys.platform == "win32": - return "onnxruntime.dll" - if sys.platform == "darwin": - return f"libonnxruntime.{version}.dylib" - return f"libonnxruntime.so.{version}" - - -def _validate_build_resources() -> tuple[Path, Path | None]: - if shutil.which("cmake") is None: - raise RuntimeError( - "CMake is required to build LEAPP Warp support but was not found." - ) - - warp = _require_module( - "warp", - 'pip install "leapp[warp-cu12]" (CUDA 12) or ' - 'pip install "leapp[warp-cu13]" (CUDA 13)', - ) - warp_dir = Path(warp.__file__).resolve().parent - warp_library = warp_dir / "bin" / _warp_library_name() - apic_header = warp_dir / "native" / "apic.h" - if not warp_library.is_file(): - raise RuntimeError( - f"Warp is installed, but its native library was not found: {warp_library}" - ) - if not apic_header.is_file(): - raise RuntimeError( - f"Warp is installed, but its APIC header was not found: {apic_header}" - ) - - warp_import_library = None - if sys.platform == "win32": - configured_import_library = os.environ.get("LEAPP_WARP_IMPORT_LIBRARY") - warp_import_library = ( - Path(configured_import_library).expanduser() - if configured_import_library - else warp_dir / "bin" / "warp.lib" - ) - if not warp_import_library.is_file(): - raise RuntimeError( - "The Windows Warp import library was not found: " - f"{warp_import_library}. The warp-lang wheel does not include " - "warp.lib; provide one with LEAPP_WARP_IMPORT_LIBRARY." - ) - - onnxruntime = _require_module("onnxruntime", "pip install onnxruntime") - ort_library = ( - Path(onnxruntime.__file__).resolve().parent - / "capi" - / _onnxruntime_library_name(onnxruntime.__version__) - ) - if not ort_library.is_file(): - raise RuntimeError( - "ONNX Runtime is installed, but its shared library was not found: " - f"{ort_library}" - ) - - torch = _require_module("torch", "pip install torch") - cmake_prefix = getattr(getattr(torch, "utils", None), "cmake_prefix_path", None) - if not cmake_prefix: - raise RuntimeError("PyTorch is installed, but its CMake path is unavailable.") - torch_cmake_dir = Path(cmake_prefix) / "Torch" - if not (torch_cmake_dir / "TorchConfig.cmake").is_file(): - raise RuntimeError( - "PyTorch is installed, but its CMake resources were not found: " - f"{torch_cmake_dir}" - ) - - return _find_nvcc(), warp_import_library - - def build_warp_runtime( build_dir: str | Path | None = None, + cmake_args: list[str] | None = None, ) -> dict[str, Path]: source_dir = _runtime_source_dir() if not (source_dir / "CMakeLists.txt").is_file(): @@ -155,28 +40,33 @@ def build_warp_runtime( "Reinstall LEAPP from a distribution that includes native sources." ) - nvcc, warp_import_library = _validate_build_resources() + if shutil.which("cmake") is None: + raise RuntimeError( + "CMake is required to build LEAPP Warp support but was not found." + ) artifacts = warp_runtime_artifact_paths(build_dir) output_dir = next(iter(artifacts.values())).parent - if output_dir.exists(): - shutil.rmtree(output_dir) - output_dir.mkdir(parents=True, exist_ok=True) configure_command = [ "cmake", "-S", str(source_dir), - "-B", - str(output_dir), f"-DPython3_EXECUTABLE={sys.executable}", - f"-DCMAKE_CUDA_COMPILER={nvcc}", "-DCMAKE_BUILD_TYPE=Release", "-DLEAPP_WARP_BUILD_ONNX=ON", "-DLEAPP_WARP_BUILD_TORCH=ON", ] - if warp_import_library is not None: - configure_command.append(f"-DWARP_IMPORT_LIBRARY={warp_import_library}") - subprocess.run(configure_command, check=True) + configure_command.extend(cmake_args or []) + if output_dir.exists(): + # Let CMake validate the selected SDKs before removing a cached build. + with tempfile.TemporaryDirectory( + prefix="leapp-warp-preflight-", dir=output_dir.parent + ) as preflight_dir: + subprocess.run([*configure_command, "-B", preflight_dir], check=True) + shutil.rmtree(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + subprocess.run([*configure_command, "-B", str(output_dir)], check=True) subprocess.run( ["cmake", "--build", str(output_dir), "--config", "Release", "-j"], check=True, @@ -222,13 +112,20 @@ def main() -> int: type=Path, help="Override the build directory.", ) + parser.add_argument( + "cmake_args", + nargs=argparse.REMAINDER, + help="CMake configure arguments after -- (for example -DWARP_ROOT=/path/to/warp).", + ) args = parser.parse_args() + if args.cmake_args[:1] == ["--"]: + args.cmake_args = args.cmake_args[1:] if args.status: return show_status(args.build_dir) try: - artifacts = build_warp_runtime(args.build_dir) + artifacts = build_warp_runtime(args.build_dir, args.cmake_args) except (OSError, RuntimeError, subprocess.CalledProcessError) as exc: print(f"Failed to build LEAPP Warp runtime: {exc}", file=sys.stderr) return 1 diff --git a/tests/functional_tests/test_warp_tracing.py b/tests/functional_tests/test_warp_tracing.py index 712c7d5e..d068d8a9 100644 --- a/tests/functional_tests/test_warp_tracing.py +++ b/tests/functional_tests/test_warp_tracing.py @@ -4,15 +4,18 @@ # import contextlib +import json import os import sys import tempfile import types import unittest +from pathlib import Path from unittest import mock import leapp import numpy as np +import onnxruntime as ort import torch import warp as wp from leapp.leapp import _MANAGER as annotate @@ -78,7 +81,7 @@ def _run_warp_to_torch_to_warp_node(self, device): leapp.stop() return node - def test_single_warp_op_two_pass_compiles(self): + def _compile_copy_node(self): leapp.start(name=self.TEST_GRAPH_NAME) source = wp.array([1.0, 2.0, 3.0], dtype=wp.float32, device=self.DEVICE) @@ -92,7 +95,10 @@ def test_single_warp_op_two_pass_compiles(self): node = annotate.get_nodes()["node_a"] leapp.stop() leapp.compile_graph(visualize=False) + return node + def test_single_warp_op_two_pass_compiles(self): + node = self._compile_copy_node() self.assertFalse(node.has_pending_warp_segments) self.verify_node_io(node, inputs=1, outputs=1) self.assertEqual(len(node.warp_segments), 1) @@ -100,6 +106,99 @@ def test_single_warp_op_two_pass_compiles(self): self.assertIsNotNone(node.warp_segments[0].apic_graph) self.verify_all_models_exist("node_a") + # GPU I/O must keep the Warp node on CUDA without ORT host-copy nodes. + for stream in (torch.cuda.default_stream(), torch.cuda.Stream()): + with self.subTest(stream=stream.cuda_stream): + options = ort.SessionOptions() + options.enable_profiling = True + options.profile_file_prefix = str(Path(self.TEST_GRAPH_NAME) / "warp") + options.register_custom_ops_library( + os.environ["LEAPP_WARP_ONNX_CUSTOM_OP_LIBRARY"] + ) + session = ort.InferenceSession( + str(Path(self.TEST_GRAPH_NAME) / "node_a.onnx"), + options, + providers=[ + ( + "CUDAExecutionProvider", + {"user_compute_stream": str(stream.cuda_stream)}, + ) + ], + ) + binding = session.io_binding() + binding.bind_output(session.get_outputs()[0].name, "cuda", 0) + for values in ([10.0, 20.0, 30.0], [7.0, -2.0, 5.0]): + expected = np.array(values, dtype=np.float32) + tensor = ort.OrtValue.ortvalue_from_numpy(expected, "cuda", 0) + binding.bind_ortvalue_input(session.get_inputs()[0].name, tensor) + session.run_with_iobinding(binding) + binding.synchronize_outputs() + np.testing.assert_array_equal( + binding.get_outputs()[0].numpy(), expected + ) + events = json.loads(Path(session.end_profiling()).read_text()) + nodes = [event["args"] for event in events if event["cat"] == "Node"] + warp_nodes = [ + node for node in nodes if node.get("op_name") == "WrpRunner" + ] + self.assertTrue(warp_nodes) + self.assertTrue( + all( + node["provider"] == "CUDAExecutionProvider" + for node in warp_nodes + ) + ) + self.assertFalse( + any("Memcpy" in node.get("op_name", "") for node in nodes) + ) + + def test_native_torch_strides_and_stream_ordering(self): + import onnx + + self._compile_copy_node() + model = onnx.load(Path(self.TEST_GRAPH_NAME) / "node_a.onnx") + op = next(node for node in model.graph.node if node.op_type == "WrpRunner") + metadata = next( + a.s.decode() for a in op.attribute if a.name == "runtime_metadata" + ) + archive = next(t for t in model.graph.initializer if t.name == op.input[-1]) + bundle = torch.tensor(onnx.numpy_helper.to_array(archive), dtype=torch.uint8) + torch.ops.load_library(os.environ["LEAPP_WARP_PT2_CUSTOM_OP_LIBRARY"]) + + output_index = next( + i for i, spec in enumerate(json.loads(metadata)["outputs"]) if spec["mask"] + ) + + def replay(values): + return torch.ops.leapp.warp_runner([values], metadata, bundle)[output_index] + + for location in ("cpu", "pinned", "cuda"): + with self.subTest(location=location): + values = torch.arange(6, dtype=torch.float32, device="cpu") + values = ( + values.pin_memory() if location == "pinned" else values.to(location) + ) + values = values[::2] + self.assertFalse(values.is_contiguous()) + torch.testing.assert_close(replay(values), values) + + # Warm before overlapping submissions: first-use graph capture is serial. + replay(torch.zeros(3, device="cuda")) + torch.cuda.synchronize() + streams = [torch.cuda.Stream(), torch.cuda.Stream()] + values = [torch.full((3,), float(i), device="cuda") for i in range(16)] + for stream in streams: + stream.wait_stream(torch.cuda.current_stream()) + outputs = [] + for i, value in enumerate(values): + with torch.cuda.stream(streams[i % 2]): + if i == 0: + torch.cuda._sleep(100_000_000) + outputs.append(replay(value)) + torch.cuda.synchronize() + for actual, expected in zip(outputs, values): + torch.testing.assert_close(actual, expected) + def test_warp_write_to_plain_torch_alias_is_traced(self): leapp.start(name=self.TEST_GRAPH_NAME) source = torch.tensor([1.0, 2.0, 3.0], device=self.DEVICE) diff --git a/tests/unit_tests/test_warp_runtime.py b/tests/unit_tests/test_warp_runtime.py index 9add6b53..cd2bebf7 100644 --- a/tests/unit_tests/test_warp_runtime.py +++ b/tests/unit_tests/test_warp_runtime.py @@ -6,7 +6,6 @@ from importlib.machinery import SourceFileLoader from importlib.util import module_from_spec, spec_from_loader from pathlib import Path -from types import SimpleNamespace import pytest @@ -24,35 +23,6 @@ _BUILD_SCRIPT_LOADER.exec_module(_BUILD_SCRIPT) -def _create_build_resources(tmp_path, windows=False): - warp_dir = tmp_path / "warp" - (warp_dir / "bin").mkdir(parents=True) - (warp_dir / "native").mkdir() - (warp_dir / "bin" / ("warp.dll" if windows else "warp.so")).touch() - (warp_dir / "native" / "apic.h").touch() - - ort_dir = tmp_path / "onnxruntime" - (ort_dir / "capi").mkdir(parents=True) - ort_library = "onnxruntime.dll" if windows else "libonnxruntime.so.1.2.3" - (ort_dir / "capi" / ort_library).touch() - - torch_dir = tmp_path / "torch" - (torch_dir / "share" / "cmake" / "Torch").mkdir(parents=True) - (torch_dir / "share" / "cmake" / "Torch" / "TorchConfig.cmake").touch() - - modules = { - "warp": SimpleNamespace(__file__=str(warp_dir / "__init__.py")), - "onnxruntime": SimpleNamespace( - __file__=str(ort_dir / "__init__.py"), - __version__="1.2.3", - ), - "torch": SimpleNamespace( - utils=SimpleNamespace(cmake_prefix_path=str(torch_dir / "share" / "cmake")) - ), - } - return modules - - def test_cache_directory_respects_xdg_cache_home(tmp_path, monkeypatch): monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) @@ -106,83 +76,19 @@ def test_missing_environment_override_does_not_fall_back(tmp_path, monkeypatch): @pytest.mark.parametrize( - ("missing_resource", "message"), + "cmake_args", [ - ("warp/bin/warp.so", "native library was not found"), - ( - "onnxruntime/capi/libonnxruntime.so.1.2.3", - "shared library was not found", - ), - ( - "torch/share/cmake/Torch/TorchConfig.cmake", - "CMake resources were not found", - ), + [], + [ + "-DWARP_ROOT=/sdk/warp with spaces", + "-DTorch_DIR=/sdk/torch/share/cmake/Torch", + "-DCUDAToolkit_ROOT=/sdk/cuda", + "-DORT_INCLUDE_DIR=/sdk/onnxruntime/include", + "-DCMAKE_BUILD_TYPE=Debug", + ], ], ) -def test_build_preflight_reports_missing_native_resources( - tmp_path, monkeypatch, missing_resource, message -): - modules = _create_build_resources(tmp_path) - (tmp_path / missing_resource).unlink() - monkeypatch.setattr( - _BUILD_SCRIPT.importlib, - "import_module", - lambda name: modules[name], - ) - monkeypatch.setattr( - _BUILD_SCRIPT.shutil, - "which", - lambda name: f"/usr/bin/{name}", - ) - monkeypatch.setattr( - _BUILD_SCRIPT, - "_find_nvcc", - lambda: Path("/usr/local/cuda/bin/nvcc"), - ) - - with pytest.raises(RuntimeError, match=message): - _BUILD_SCRIPT._validate_build_resources() - - -def test_missing_warp_dependency_recommends_cuda_extra(monkeypatch): - def missing_module(name): - raise ImportError(name) - - monkeypatch.setattr(_BUILD_SCRIPT.importlib, "import_module", missing_module) - monkeypatch.setattr( - _BUILD_SCRIPT.shutil, - "which", - lambda name: "/usr/bin/cmake" if name == "cmake" else None, - ) - - with pytest.raises(RuntimeError, match="never installs dependencies") as exc_info: - _BUILD_SCRIPT._validate_build_resources() - - message = str(exc_info.value) - assert "leapp[warp-cu12]" in message - assert "leapp[warp-cu13]" in message - - -def test_windows_preflight_reports_missing_warp_import_library(tmp_path, monkeypatch): - modules = _create_build_resources(tmp_path, windows=True) - monkeypatch.setattr(_BUILD_SCRIPT.sys, "platform", "win32") - monkeypatch.delenv("LEAPP_WARP_IMPORT_LIBRARY", raising=False) - monkeypatch.setattr( - _BUILD_SCRIPT.importlib, - "import_module", - lambda name: modules[name], - ) - monkeypatch.setattr( - _BUILD_SCRIPT.shutil, - "which", - lambda name: f"C:/{name}.exe", - ) - - with pytest.raises(RuntimeError, match="LEAPP_WARP_IMPORT_LIBRARY"): - _BUILD_SCRIPT._validate_build_resources() - - -def test_build_configures_both_adapters(tmp_path, monkeypatch): +def test_build_configures_both_adapters(tmp_path, monkeypatch, cmake_args): source_dir = tmp_path / "source" source_dir.mkdir() (source_dir / "CMakeLists.txt").touch() @@ -200,18 +106,30 @@ def run(command, check): monkeypatch.setattr(_BUILD_SCRIPT, "_runtime_source_dir", lambda: source_dir) monkeypatch.setattr( - _BUILD_SCRIPT, - "_validate_build_resources", - lambda: (Path("/usr/local/cuda/bin/nvcc"), None), + _BUILD_SCRIPT.shutil, + "which", + lambda name: "/usr/bin/cmake", ) monkeypatch.setattr(_BUILD_SCRIPT.subprocess, "run", run) - artifacts = _BUILD_SCRIPT.build_warp_runtime(build_dir) + monkeypatch.setattr( + _BUILD_SCRIPT.sys, + "argv", + ["leapp-build-warp-runtime", "--build-dir", str(build_dir), "--", *cmake_args], + ) - assert len(commands) == 2 + assert _BUILD_SCRIPT.main() == 0 + + assert len(commands) == 3 assert "-DLEAPP_WARP_BUILD_ONNX=ON" in commands[0] assert "-DLEAPP_WARP_BUILD_TORCH=ON" in commands[0] - assert set(artifacts) == {"onnx", "pt2"} + if cmake_args: + assert commands[0][-len(cmake_args) - 2 : -2] == cmake_args + assert "--" not in commands[0] + assert all( + path.is_file() + for path in warp_runtime.warp_runtime_artifact_paths(build_dir).values() + ) assert not stale_file.exists() From 0bda0ed70d98de040de4446df9ab814fc955de9a Mon Sep 17 00:00:00 2001 From: Marc Alwan Date: Fri, 23 Oct 2026 16:27:54 -0700 Subject: [PATCH 2/2] ci: restore GPU ONNX Runtime before Warp tests Signed-off-by: Marc Alwan --- .github/workflows/ci.yml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 277aa8bf..43e67eb4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,6 +57,15 @@ jobs: - name: Install package and test dependencies run: python -m pip install -e ".[dev,test,warp-cu12]" cmake ninja + - name: Select GPU ONNX Runtime + run: | + # The base CPU dependency and Warp GPU extra share package files. + # Keep the CUDA-compatible version resolved by the Warp extra. + ort_gpu_version=$(python -c "from importlib.metadata import version; print(version('onnxruntime-gpu'))") + python -m pip uninstall -y onnxruntime + python -m pip install --force-reinstall --no-deps "onnxruntime-gpu==$ort_gpu_version" + python -c "import onnxruntime as ort; assert 'CUDAExecutionProvider' in ort.get_available_providers(), ort.get_available_providers()" + - name: Build Warp runtime run: leapp-build-warp-runtime