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 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/AGENTS.md b/packages/leapp-runtime-rust/AGENTS.md index dee42c51..441e45f5 100644 --- a/packages/leapp-runtime-rust/AGENTS.md +++ b/packages/leapp-runtime-rust/AGENTS.md @@ -95,6 +95,14 @@ cargo test --locked engines::triton::tests::torchscript -- --ignored --test-thre cargo test --locked engines::triton::tests::tensorrt -- --ignored --test-threads=1 ``` +For Warp, install the adapters as described in `src/engines/triton/backends/README.md`. +Generate `tests/generators/generate_warp_models.py` fixtures, set +`LEAPP_WARP_FIXTURES` and an empty `WARP_APIC_TMPDIR`, then run +`cargo test --locked engines::triton::tests::warp -- --ignored --test-threads=1`. +This covers ONNX/PT2 CPU and GPU instances, CPU/pinned/CUDA output placement, +pool limits 0/1/64, changed inputs, retained results, peer sessions, and unload. +Warp replay requires CUDA even with CPU model instances. + The Python and C++ binding READMEs describe building and testing each adapter. Also verify the installed C++ SDK with `bindings/cxx/tests/consumer` so installed headers and CMake targets work without generated bridge headers or Rust sources. diff --git a/packages/leapp-runtime-rust/README.md b/packages/leapp-runtime-rust/README.md index 90f7ed4a..f9406be9 100644 --- a/packages/leapp-runtime-rust/README.md +++ b/packages/leapp-runtime-rust/README.md @@ -21,6 +21,10 @@ See the [Python bindings](bindings/python/README.md) for `import lrt` and the [C++ bindings](bindings/cxx/README.md) for the `lrt` namespace and `lrt::runtime` CMake target. CPU execution still requires the CUDA runtime SDK. +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 c2581a7c..1a5f3a3e 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 c9e7e16a..43dbaed2 100644 --- a/packages/leapp-runtime-rust/src/engines/triton/backends/README.md +++ b/packages/leapp-runtime-rust/src/engines/triton/backends/README.md @@ -61,7 +61,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 @@ -75,6 +76,62 @@ 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. +Warp bundles use the same session-owned output pool and initialization options +described above. Retained results remain valid across subsequent runs and session +closure; the adapters continue to own their internal replay buffers. + +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..4540076e --- /dev/null +++ b/packages/leapp-runtime-rust/tests/functional_tests/triton/warp.rs @@ -0,0 +1,137 @@ +// 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, OutputPoolOptions, Session, + SessionOptions, +}; +use std::{collections::HashMap, 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 }, + ] { + for max_entries in [0, 1, 64] { + let graphs_before = graph_count(); + let mut session = Session::open( + &bundle, + SessionOptions { + device, + output_placement: OutputPlacement { + default_location: location, + ..Default::default() + }, + output_pool: OutputPoolOptions { max_entries }, + }, + ) + .unwrap(); + let gain = placed(&[1_i32, 2, 3, 4], &[4], location).unwrap(); + let run = |session: &mut Session, offset: f32| { + let values: Vec<_> = (0..12).map(|i| i as f32 + offset).collect(); + let x = placed(&values, &[4, 3], location).unwrap(); + session + .run(&HashMap::from([ + ("model/x", x.view()), + ("model/gain", gain.view()), + ])) + .unwrap() + }; + let expected = |offset: f32| -> Vec { + (0..12) + .map(|i| (i as f32 + offset) * (i / 3 + 1) as f32 + 4.0) + .collect() + }; + let retained = run(&mut session, 0.0); + assert_eq!( + retained + .get("model/result") + .unwrap() + .to_cpu_vec::() + .unwrap(), + expected(0.0), + "{fixture} {device:?} {location:?} pool={max_entries}" + ); + for offset in [1.0, 2.0, 3.0] { + assert_eq!( + run(&mut session, offset) + .get("model/result") + .unwrap() + .to_cpu_vec::() + .unwrap(), + expected(offset) + ); + } + 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, 4.0); + drop(peer); + let output = result.get("model/result").unwrap(); + assert_eq!(output.location(), location); + assert_eq!(output.to_cpu_vec::().unwrap(), expected(4.0)); + assert_eq!( + retained + .get("model/result") + .unwrap() + .to_cpu_vec::() + .unwrap(), + expected(0.0) + ); + drop(result); + drop(retained); + 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()