From 4edc98c55271426c5997a77d01262ae8abff3d8b Mon Sep 17 00:00:00 2001 From: Bhagirath Mehta Date: Fri, 24 Jul 2026 14:55:37 -0500 Subject: [PATCH] Rebuild telemetry hardening stack on core Carry forward the runtime packaging, download, catalog, shutdown, streaming, binding lifetime, and test hardening from PR #880 on top of the latest telemetry core branch. Drop the obsolete telemetry ABI/config/redaction churn so #879 remains the source of truth for non-essential telemetry semantics and backend-owned cleanup. Files changed: - .github/workflows/samples-integration-test.yml - .pipelines/v2/templates/* - sdk_v2/cpp/** - sdk_v2/js/** - sdk_v2/python/** - package metadata under sdk/js and www Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 297bc539-5f7f-4a70-9411-48e91a5bf532 --- .../workflows/samples-integration-test.yml | 6 +- .pipelines/v2/templates/steps-build-js.yml | 8 + .pipelines/v2/templates/steps-build-linux.yml | 47 ++- .pipelines/v2/templates/steps-build-macos.yml | 40 ++- .../v2/templates/steps-build-python.yml | 24 +- .../v2/templates/steps-build-windows.yml | 14 +- .../v2/templates/steps-pack-cpp-sdk.yml | 44 ++- .pipelines/v2/templates/steps-pack-js.yml | 1 + README.md | 2 +- sdk_v2/cpp/CMakeLists.txt | 82 +++-- sdk_v2/cpp/build.py | 25 +- sdk_v2/cpp/cmake/FindOnnxRuntime.cmake | 20 +- sdk_v2/cpp/cmake/FindOnnxRuntimeGenAI.cmake | 16 +- sdk_v2/cpp/cmake/FindWinMLEpCatalog.cmake | 42 +-- sdk_v2/cpp/nuget/pack.py | 21 +- sdk_v2/cpp/src/catalog/azure_model_catalog.cc | 38 ++- sdk_v2/cpp/src/catalog/base_model_catalog.cc | 81 ++--- sdk_v2/cpp/src/catalog/base_model_catalog.h | 2 + sdk_v2/cpp/src/catalog/catalog_client.cc | 15 +- sdk_v2/cpp/src/catalog/catalog_client.h | 9 +- .../cpp/src/download/blob_download_state.cc | 57 +++- sdk_v2/cpp/src/download/blob_download_state.h | 12 +- sdk_v2/cpp/src/download/blob_downloader.cc | 139 ++++++-- sdk_v2/cpp/src/download/blob_downloader.h | 21 +- sdk_v2/cpp/src/download/download_manager.cc | 98 +++--- sdk_v2/cpp/src/download/download_manager.h | 33 +- sdk_v2/cpp/src/download/file_writer.cc | 20 +- .../src/ep_detection/cuda_ep_bootstrapper.cc | 41 ++- sdk_v2/cpp/src/ep_detection/ep_detector.cc | 100 +++--- sdk_v2/cpp/src/ep_detection/ep_detector.h | 27 +- sdk_v2/cpp/src/ep_detection/ep_utils.cc | 25 ++ sdk_v2/cpp/src/ep_detection/ep_utils.h | 8 + .../ep_detection/webgpu_ep_bootstrapper.cc | 16 +- .../src/ep_detection/webgpu_ep_bootstrapper.h | 1 + .../generative/audio/audio_session.cc | 8 +- .../generative/audio/audio_session.h | 3 +- .../generative/chat/chat_session.cc | 7 +- .../generative/chat/chat_session.h | 3 +- .../embeddings/embeddings_session.cc | 31 +- .../embeddings/embeddings_session.h | 7 +- .../cpp/src/inferencing/model_load_manager.cc | 115 +++++-- .../cpp/src/inferencing/model_load_manager.h | 31 ++ sdk_v2/cpp/src/inferencing/session/session.cc | 41 +-- sdk_v2/cpp/src/inferencing/session/session.h | 7 - .../inferencing/session/session_manager.cc | 99 +++--- .../src/inferencing/session/session_manager.h | 30 +- sdk_v2/cpp/src/manager.cc | 4 +- sdk_v2/cpp/src/model.cc | 70 ++-- sdk_v2/cpp/src/model.h | 10 +- .../service/audio_transcriptions_handler.cc | 50 +-- .../service/audio_transcriptions_handler.h | 4 +- .../src/service/chat_completions_handler.cc | 56 ++-- .../src/service/chat_completions_handler.h | 4 +- sdk_v2/cpp/src/service/embeddings_handler.cc | 9 +- sdk_v2/cpp/src/service/handler_utils.h | 58 +++- sdk_v2/cpp/src/service/responses_handler.cc | 177 ++++++----- sdk_v2/cpp/src/service/responses_handler.h | 12 +- sdk_v2/cpp/src/service/web_service.cc | 153 +++++---- sdk_v2/cpp/src/service/web_service.h | 134 +++++++- sdk_v2/cpp/test/CMakeLists.txt | 9 +- .../internal_api/audio/audio_session_test.cc | 2 +- .../test/internal_api/azure_catalog_test.cc | 14 +- .../internal_api/base_model_catalog_test.cc | 103 ++++++ .../internal_api/blob_download_state_test.cc | 24 ++ .../internal_api/chat/chat_session_test.cc | 2 +- sdk_v2/cpp/test/internal_api/download_test.cc | 298 ++++++++++++++++-- .../cpp/test/internal_api/ep_detector_test.cc | 20 +- .../internal_api/model_load_manager_test.cc | 13 + .../test/internal_api/session_manager_test.cc | 27 +- .../cpp/test/internal_api/web_service_test.cc | 18 +- sdk_v2/cpp/test/test_main.cc | 27 ++ sdk_v2/js/native/src/catalog.cc | 73 ++++- sdk_v2/js/native/src/catalog.h | 7 + sdk_v2/js/native/src/manager.cc | 51 ++- sdk_v2/js/native/src/manager.h | 12 +- sdk_v2/js/native/src/model.cc | 136 +++++++- sdk_v2/js/native/src/model.h | 9 + sdk_v2/js/native/src/session.cc | 130 ++++++-- sdk_v2/js/native/src/session.h | 20 +- sdk_v2/js/script/copy-native.mjs | 1 + sdk_v2/js/script/install-native.cjs | 1 - sdk_v2/js/script/pack-prebuilds.mjs | 10 +- sdk_v2/js/src/catalog.ts | 1 + sdk_v2/js/test/manager-dispose.test.ts | 24 ++ .../foundry_local_sdk/_native/lib_loader.py | 25 +- sdk_v2/python/src/foundry_local_sdk/items.py | 2 +- .../python/src/foundry_local_sdk/request.py | 2 +- .../python/src/foundry_local_sdk/response.py | 2 +- .../python/src/foundry_local_sdk/session.py | 2 +- 89 files changed, 2432 insertions(+), 891 deletions(-) create mode 100644 sdk_v2/cpp/test/test_main.cc diff --git a/.github/workflows/samples-integration-test.yml b/.github/workflows/samples-integration-test.yml index 5613f3b67..0aedb0608 100644 --- a/.github/workflows/samples-integration-test.yml +++ b/.github/workflows/samples-integration-test.yml @@ -99,7 +99,7 @@ jobs: clean: true - name: Setup Node.js - uses: actions/setup-node@v7 + uses: actions/setup-node@v6 with: node-version: '20.x' @@ -109,7 +109,7 @@ jobs: python-version: '3.x' - name: Setup .NET SDK for NuGet authentication - uses: actions/setup-dotnet@v6 + uses: actions/setup-dotnet@v5 with: dotnet-version: '10.0.x' @@ -185,7 +185,7 @@ jobs: clean: true - name: Setup .NET SDK - uses: actions/setup-dotnet@v6 + uses: actions/setup-dotnet@v5 with: dotnet-version: | 8.0.x diff --git a/.pipelines/v2/templates/steps-build-js.yml b/.pipelines/v2/templates/steps-build-js.yml index 657704bd9..cf18c80c4 100644 --- a/.pipelines/v2/templates/steps-build-js.yml +++ b/.pipelines/v2/templates/steps-build-js.yml @@ -156,6 +156,14 @@ steps: throw "WinML runtime DLL not found in native artifact: $winmlDll" } + $directMlDll = Join-Path "${{ parameters.nativeArtifactDir }}" 'DirectML.dll' + if (Test-Path $directMlDll) { + Copy-Item $directMlDll -Destination $dst -Force + Write-Host " staged DirectML.dll" + } else { + throw "DirectML.dll not found in native artifact: $directMlDll" + } + Get-ChildItem $dst | ForEach-Object { Write-Host " $($_.Name) $($_.Length) bytes" } displayName: 'Stage prebuild directory' diff --git a/.pipelines/v2/templates/steps-build-linux.yml b/.pipelines/v2/templates/steps-build-linux.yml index b1328cba2..20a2a08a0 100644 --- a/.pipelines/v2/templates/steps-build-linux.yml +++ b/.pipelines/v2/templates/steps-build-linux.yml @@ -75,27 +75,44 @@ steps: displayName: 'Dump vcpkg error logs' condition: failed() -# Stage the redistributable native artifact. Vcpkg uses static linkage on -# Linux (see sdk_v2/cpp/triplets/x64-linux.cmake), so libfoundry_local.so -# carries its transitive deps (azure-*, spdlog, fmt, libcurl, libssl/libcrypto, -# zlib, brotli*) inside itself — the only file we need to forward downstream -# is libfoundry_local.so. ORT/GenAI .so files are also present in bin/ but are -# supplied to consumers separately (pip on the Python side, NuGet on the C# -# side); test/example binaries are not part of the redistributable surface. +# Stage the redistributable native artifacts for standalone C++ consumers. - bash: | set -euo pipefail src='$(Build.SourcesDirectory)/sdk_v2/cpp/build/Linux/${{ parameters.buildConfig }}/bin' dst='$(Build.ArtifactStagingDirectory)/native' mkdir -p "$dst" - primary="$src/libfoundry_local.so" - if [ ! -f "$primary" ]; then - echo "ERROR: libfoundry_local.so not found at $primary" >&2 - exit 1 - fi - cp -P "$primary" "$dst/" - echo " staged libfoundry_local.so" - + required=( + "$src/libfoundry_local.so" + "$src/libonnxruntime-genai.so" + "$src/libonnxruntime.so" + ) + for f in "${required[@]}"; do + if [ ! -e "$f" ]; then + echo "ERROR: required native artifact not found: $f" >&2 + exit 1 + fi + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + done + for f in "$src"/libonnxruntime.so.*; do + if [ -e "$f" ]; then + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + fi + done + for f in "$src"/libonnxruntime-genai*.so; do + if [ -e "$f" ] && [ "$(basename "$f")" != "libonnxruntime-genai.so" ]; then + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + fi + done + for f in "$src"/libonnxruntime_providers_*.so; do + if [ -e "$f" ]; then + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + fi + done displayName: 'Stage native artifacts' # Stage debug symbols into a separate artifact so they can be published diff --git a/.pipelines/v2/templates/steps-build-macos.yml b/.pipelines/v2/templates/steps-build-macos.yml index a5da9fab9..eeb90ab53 100644 --- a/.pipelines/v2/templates/steps-build-macos.yml +++ b/.pipelines/v2/templates/steps-build-macos.yml @@ -81,26 +81,38 @@ steps: displayName: 'Dump vcpkg error logs' condition: failed() -# Stage the redistributable native artifact. macOS has no triplet overlay so -# vcpkg uses its default arm64-osx triplet, which is static — libfoundry_local.dylib -# carries its transitive deps (azure-*, spdlog, fmt, libcurl, libssl/libcrypto, -# zlib, brotli*) inside itself. ORT/GenAI .dylib files are supplied to consumers -# separately (pip on the Python side, NuGet on the C# side); test/example binaries -# are not part of the redistributable surface. +# Stage the redistributable native artifacts for standalone C++ consumers. - bash: | set -euo pipefail src='$(Build.SourcesDirectory)/sdk_v2/cpp/build/macOS/${{ parameters.buildConfig }}/bin' dst='$(Build.ArtifactStagingDirectory)/native' mkdir -p "$dst" - primary="$src/libfoundry_local.dylib" - if [ ! -f "$primary" ]; then - echo "ERROR: libfoundry_local.dylib not found at $primary" >&2 - exit 1 - fi - cp -P "$primary" "$dst/" - echo " staged libfoundry_local.dylib" - + required=( + "$src/libfoundry_local.dylib" + "$src/libonnxruntime-genai.dylib" + "$src/libonnxruntime.dylib" + ) + for f in "${required[@]}"; do + if [ ! -e "$f" ]; then + echo "ERROR: required native artifact not found: $f" >&2 + exit 1 + fi + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + done + for f in "$src"/libonnxruntime.*.dylib; do + if [ -e "$f" ]; then + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + fi + done + for f in "$src"/libonnxruntime_providers_*.dylib; do + if [ -e "$f" ]; then + cp -P "$f" "$dst/" + echo " staged $(basename "$f")" + fi + done displayName: 'Stage native artifacts' # Stage debug symbols into a separate artifact so they can be published diff --git a/.pipelines/v2/templates/steps-build-python.yml b/.pipelines/v2/templates/steps-build-python.yml index 50f86c9da..84201a682 100644 --- a/.pipelines/v2/templates/steps-build-python.yml +++ b/.pipelines/v2/templates/steps-build-python.yml @@ -106,10 +106,32 @@ steps: # deps (onnxruntime-{core,gpu} / onnxruntime-genai-{core,cuda}). Test # and example binaries are also filtered out upstream. We just copy # the curated artifact contents straight in. - $files = Get-ChildItem -Path "${{ parameters.nativeArtifactDir }}" -Recurse -File + $rid = '${{ parameters.rid }}' + if ($rid -like 'win-*') { + $allowed = @('foundry_local.dll', 'foundry_local.lib', 'Microsoft.Windows.AI.MachineLearning.dll', 'DirectML.dll') + } elseif ($rid -like 'linux-*') { + $allowed = @('libfoundry_local.so') + } elseif ($rid -like 'osx-*') { + $allowed = @('libfoundry_local.dylib') + } else { + throw "Unsupported Python wheel RID: $rid" + } + $files = foreach ($name in $allowed) { + Get-ChildItem -Path "${{ parameters.nativeArtifactDir }}" -Filter $name -File -ErrorAction SilentlyContinue + } if ($files.Count -eq 0) { throw "No native artifacts found under ${{ parameters.nativeArtifactDir }}" } + if ($rid -like 'win-*') { + $found = @{} + foreach ($f in $files) { + $found[$f.Name] = $true + } + $missing = @($allowed | Where-Object { -not $found.ContainsKey($_) }) + if ($missing.Count -gt 0) { + throw "Missing required Windows native artifact(s): $($missing -join ', ') under ${{ parameters.nativeArtifactDir }}" + } + } foreach ($f in $files) { Copy-Item $f.FullName -Destination $dest -Force } diff --git a/.pipelines/v2/templates/steps-build-windows.yml b/.pipelines/v2/templates/steps-build-windows.yml index 5f1d0f761..e72fde515 100644 --- a/.pipelines/v2/templates/steps-build-windows.yml +++ b/.pipelines/v2/templates/steps-build-windows.yml @@ -138,7 +138,13 @@ steps: (Join-Path $binDir 'foundry_local.dll'), (Join-Path $binDir 'foundry_local.pdb'), (Join-Path $linkDir 'foundry_local.lib'), - (Join-Path $binDir 'Microsoft.Windows.AI.MachineLearning.dll') + (Join-Path $binDir 'Microsoft.Windows.AI.MachineLearning.dll'), + (Join-Path $binDir 'DirectML.dll'), + (Join-Path $binDir 'onnxruntime.dll'), + (Join-Path $binDir 'onnxruntime-genai.dll') + ) + $optionalPatterns = @( + 'onnxruntime_providers_*.dll' ) foreach ($s in $sources) { @@ -149,6 +155,12 @@ steps: Copy-Item -Path $s -Destination $dst -Force Write-Host " staged $(Split-Path -Leaf $s)" } + foreach ($pattern in $optionalPatterns) { + Get-ChildItem -Path $binDir -Filter $pattern -File -ErrorAction SilentlyContinue | ForEach-Object { + Copy-Item -Path $_.FullName -Destination $dst -Force + Write-Host " staged $($_.Name)" + } + } - ${{ if eq(parameters.stageHeaders, true) }}: - task: CopyFiles@2 diff --git a/.pipelines/v2/templates/steps-pack-cpp-sdk.yml b/.pipelines/v2/templates/steps-pack-cpp-sdk.yml index a878dbf3d..984bc6a33 100644 --- a/.pipelines/v2/templates/steps-pack-cpp-sdk.yml +++ b/.pipelines/v2/templates/steps-pack-cpp-sdk.yml @@ -61,9 +61,8 @@ steps: Copy-Item -Path (Join-Path $includeSrc '*') -Destination $includeDst -Recurse -Force - # The SDK tgz should include only public wrapper headers. + # The SDK tgz should include public wrapper headers and their header-only dependencies. $pathsToPrune = @( - (Join-Path $includeDst 'gsl'), (Join-Path $includeDst '_manifest') ) foreach ($pathToPrune in $pathsToPrune) { @@ -72,13 +71,17 @@ steps: } } - foreach ($f in $Files) { - $src = Join-Path $nativeSrc $f - if (-not (Test-Path $src)) { - throw "Required native file not found: $src" + foreach ($entry in $Files) { + $optional = $entry.StartsWith('?') + $pattern = if ($optional) { $entry.Substring(1) } else { $entry } + $matches = @(Get-ChildItem -Path $nativeSrc -Filter $pattern -File -ErrorAction SilentlyContinue) + if ($matches.Count -eq 0 -and -not $optional) { + throw "Required native file not found: $(Join-Path $nativeSrc $pattern)" + } + foreach ($match in $matches) { + $dst = Join-Path $libDst $match.Name + Copy-Item -Path $match.FullName -Destination $dst -Force } - $dst = Join-Path $libDst (Split-Path -Leaf $src) - Copy-Item -Path $src -Destination $dst -Force } if ($SymbolsArtifact) { @@ -115,23 +118,40 @@ steps: New-SdkArchive -Rid 'win-x64' -NativeArtifact 'cpp-native-win-x64' -Files @( 'foundry_local.dll', 'foundry_local.pdb', - 'foundry_local.lib' + 'foundry_local.lib', + 'Microsoft.Windows.AI.MachineLearning.dll', + 'DirectML.dll', + 'onnxruntime.dll', + 'onnxruntime-genai.dll', + '?onnxruntime_providers_*.dll' ) New-SdkArchive -Rid 'win-arm64' -NativeArtifact 'cpp-native-win-arm64' -Files @( 'foundry_local.dll', 'foundry_local.pdb', - 'foundry_local.lib' + 'foundry_local.lib', + 'Microsoft.Windows.AI.MachineLearning.dll', + 'DirectML.dll', + 'onnxruntime.dll', + 'onnxruntime-genai.dll', + '?onnxruntime_providers_*.dll' ) New-SdkArchive -Rid 'linux-x64' -NativeArtifact 'cpp-native-linux-x64' -Files @( - 'libfoundry_local.so' + 'libfoundry_local.so', + 'libonnxruntime-genai.so', + '?libonnxruntime-genai*.so', + '?libonnxruntime_providers_*.so', + 'libonnxruntime.so*' ) -SymbolsArtifact 'cpp-native-symbols-linux-x64' -SymbolItems @( 'libfoundry_local.so.dbg' ) New-SdkArchive -Rid 'osx-arm64' -NativeArtifact 'cpp-native-osx-arm64' -Files @( - 'libfoundry_local.dylib' + 'libfoundry_local.dylib', + 'libonnxruntime-genai.dylib', + 'libonnxruntime*.dylib', + '?libonnxruntime_providers_*.dylib' ) -SymbolsArtifact 'cpp-native-symbols-osx-arm64' -SymbolItems @( 'libfoundry_local.dylib.dSYM' ) diff --git a/.pipelines/v2/templates/steps-pack-js.yml b/.pipelines/v2/templates/steps-pack-js.yml index 439ab41fd..9d7a4aff1 100644 --- a/.pipelines/v2/templates/steps-pack-js.yml +++ b/.pipelines/v2/templates/steps-pack-js.yml @@ -74,6 +74,7 @@ steps: $dst = '${{ parameters.outputDir }}' New-Item -ItemType Directory -Force -Path $dst | Out-Null Set-Location "$(Build.SourcesDirectory)/sdk_v2/js" + Copy-Item "$(Build.SourcesDirectory)/sdk_v2/deps_versions.json" -Destination "deps_versions.json" -Force npm pack --pack-destination $dst if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE } Get-ChildItem $dst | ForEach-Object { Write-Host " $($_.Name) $($_.Length) bytes" } diff --git a/README.md b/README.md index af8a1a6e9..41885aa8e 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ Foundry Local is an **end-to-end local AI solution** for building applications t User data never leaves the device, responses start immediately with zero network latency, and your app works offline. No per-token costs, no API keys, no backend infrastructure to maintain, and no Azure subscription required. -Foundry Local may collect usage data and send it to Microsoft to help improve our products and services. See the [privacy statement](sdk_v2/cpp/docs/Privacy.md) for more details. +This project may collect usage data and send it to Microsoft to help improve our products and services. See the [privacy statement](sdk_v2/cpp/docs/Privacy.md) for more details. ### Key Features diff --git a/sdk_v2/cpp/CMakeLists.txt b/sdk_v2/cpp/CMakeLists.txt index 43fb436db..bd20c5ebc 100644 --- a/sdk_v2/cpp/CMakeLists.txt +++ b/sdk_v2/cpp/CMakeLists.txt @@ -1,5 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. -cmake_minimum_required(VERSION 3.20) +cmake_minimum_required(VERSION 3.21) # Android: web service not applicable on device. # Detect Android early via the vcpkg triplet name because the ANDROID variable @@ -51,20 +51,18 @@ option(FOUNDRY_LOCAL_BUILD_EXAMPLES "Build example programs" ON) option(FOUNDRY_LOCAL_BUILD_SERVICE "Build web service support (requires oat++)" ON) option(FOUNDRY_LOCAL_ENABLE_ASAN "Enable AddressSanitizer + UndefinedBehaviorSanitizer (Linux only)" OFF) -# Optional 1DS ingestion token override. Override only via environment so it does -# not appear in CMake cache files. -if(DEFINED CACHE{FOUNDRY_LOCAL_TELEMETRY_TOKEN} - AND NOT "$CACHE{FOUNDRY_LOCAL_TELEMETRY_TOKEN}" STREQUAL "") - message(FATAL_ERROR - "Do not pass FOUNDRY_LOCAL_TELEMETRY_TOKEN via -D or CMake cache. " - "Use the FOUNDRY_LOCAL_TELEMETRY_TOKEN environment variable for temporary test-tenant overrides.") -endif() -unset(FOUNDRY_LOCAL_TELEMETRY_TOKEN CACHE) - -set(FOUNDRY_LOCAL_TELEMETRY_TOKEN "") -if(DEFINED ENV{FOUNDRY_LOCAL_TELEMETRY_TOKEN}) - set(FOUNDRY_LOCAL_TELEMETRY_TOKEN "$ENV{FOUNDRY_LOCAL_TELEMETRY_TOKEN}") +# Optional 1DS ingestion-token override. Pass an override via the +# FOUNDRY_LOCAL_TELEMETRY_TOKEN environment variable so it does not land in +# process argv or CMakeCache.txt. +get_property(_FL_TELEMETRY_TOKEN_IN_CACHE CACHE FOUNDRY_LOCAL_TELEMETRY_TOKEN PROPERTY TYPE SET) +if(_FL_TELEMETRY_TOKEN_IN_CACHE) + get_property(_FL_TELEMETRY_TOKEN_CACHE_VALUE CACHE FOUNDRY_LOCAL_TELEMETRY_TOKEN PROPERTY VALUE) + unset(FOUNDRY_LOCAL_TELEMETRY_TOKEN CACHE) + if(_FL_TELEMETRY_TOKEN_CACHE_VALUE) + message(FATAL_ERROR "FOUNDRY_LOCAL_TELEMETRY_TOKEN must be supplied as an environment variable, not via -D.") + endif() endif() +set(FOUNDRY_LOCAL_TELEMETRY_TOKEN "$ENV{FOUNDRY_LOCAL_TELEMETRY_TOKEN}") if(FOUNDRY_LOCAL_TELEMETRY_TOKEN MATCHES "^\\$\\(") set(FOUNDRY_LOCAL_TELEMETRY_TOKEN "") endif() @@ -119,9 +117,6 @@ endif() # 1DS C++ client telemetry — provided by the cpp-client-telemetry vcpkg port. find_package(MSTelemetry CONFIG REQUIRED) -if(NOT WIN32) - find_package(OpenSSL REQUIRED) -endif() message(STATUS "1DS telemetry: enabled (cpp-client-telemetry found)") # -------------------------------------------------------------------------- @@ -305,8 +300,6 @@ function(foundry_local_configure_target TARGET LINK_SCOPE) if(NOT CMAKE_SYSTEM_NAME STREQUAL "WindowsStore") target_compile_definitions(${TARGET} PRIVATE FOUNDRY_LOCAL_USE_WINHTTP_TRANSPORT=1) endif() - else() - target_link_libraries(${TARGET} ${LINK_SCOPE} OpenSSL::Crypto) endif() if(APPLE) @@ -420,6 +413,27 @@ if(WIN32) /DELAYLOAD:Microsoft.Windows.AI.MachineLearning.dll ) endif() + + # Strip unreferenced symbols and fold identical COMDATs in Release builds. + # MSVC's /DEBUG (needed for PDBs) silently disables /OPT:REF and /OPT:ICF + # unless they are explicitly re-enabled. /INCREMENTAL:NO is required because + # incremental linking also prevents dead-code elimination. + target_link_options(foundry_local PRIVATE + $<$:/OPT:REF> + $<$:/OPT:ICF> + $<$:/INCREMENTAL:NO> + ) +endif() + +# Strip unreferenced sections on Linux (GCC/Clang) and macOS (Apple ld). +if(UNIX AND NOT APPLE) + target_link_options(foundry_local PRIVATE + $<$:-Wl,--gc-sections> + ) +elseif(APPLE) + target_link_options(foundry_local PRIVATE + $<$:-Wl,-dead_strip> + ) endif() # -------------------------------------------------------------------------- @@ -469,6 +483,13 @@ if(TARGET OnnxRuntime::OnnxRuntime) $ ) endif() + if(WinMLEpCatalog_FOUND AND EXISTS "${WINML_EP_CATALOG_DLL_DIR}/DirectML.dll") + add_custom_command(TARGET foundry_local POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${WINML_EP_CATALOG_DLL_DIR}/DirectML.dll" + $ + ) + endif() elseif(APPLE) # macOS: copy dylibs so consumers that only link libfoundry_local.dylib (e.g. cache_only_tests) find the correct # ORT version instead of any system-installed ORT, which would cause an Ort::InitApi() version mismatch. @@ -500,6 +521,14 @@ if(TARGET OnnxRuntime::OnnxRuntime) $ ) + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + add_custom_command(TARGET foundry_local POST_BUILD + COMMAND ${CMAKE_COMMAND} -E create_symlink + libonnxruntime.so + $/libonnxruntime.so.1 + ) + endif() + if(EXISTS "${ORT_LIB_DIR}/libonnxruntime_providers_shared.so") add_custom_command(TARGET foundry_local POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy_if_different @@ -507,6 +536,21 @@ if(TARGET OnnxRuntime::OnnxRuntime) $ ) endif() + + if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + file(GLOB _ORT_GENAI_RUNTIME_LIBS CONFIGURE_DEPENDS + "${ORT_GENAI_LIB_DIR}/libonnxruntime-genai*.so") + file(GLOB _ORT_PROVIDER_RUNTIME_LIBS CONFIGURE_DEPENDS + "${ORT_LIB_DIR}/libonnxruntime_providers_*.so") + set(_LINUX_OPTIONAL_RUNTIME_LIBS ${_ORT_GENAI_RUNTIME_LIBS} ${_ORT_PROVIDER_RUNTIME_LIBS}) + if(_LINUX_OPTIONAL_RUNTIME_LIBS) + add_custom_command(TARGET foundry_local POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${_LINUX_OPTIONAL_RUNTIME_LIBS} + $ + ) + endif() + endif() endif() endif() diff --git a/sdk_v2/cpp/build.py b/sdk_v2/cpp/build.py index f7a9c094d..15ca08d4a 100644 --- a/sdk_v2/cpp/build.py +++ b/sdk_v2/cpp/build.py @@ -57,10 +57,25 @@ def _path_from_env_var(var: str) -> Path | None: # Subprocess runner # --------------------------------------------------------------------------- +def _redact_command_arg(arg: str) -> str: + sensitive_words = ("token", "secret", "password", "passwd", "apikey", "api_key", "accesskey", "sig") + if "=" in arg: + key, value = arg.split("=", 1) + normalized_key = key.lstrip("-").lower() + if "://" in value: + return f"{key}=" + if any(word in normalized_key for word in sensitive_words): + return f"{key}=" + return arg + + def run(cmd: list[str], cwd: Path | str | None = None, env: dict[str, str] | None = None) -> None: """Run a command, logging and streaming output. Raises on non-zero exit.""" - log.info("Running: %s", " ".join(str(c) for c in cmd)) - subprocess.run(cmd, cwd=cwd, env=env, check=True) + redacted = " ".join(_redact_command_arg(str(c)) for c in cmd) + log.info("Running: %s", redacted) + result = subprocess.run(cmd, cwd=cwd, env=env, check=False) + if result.returncode != 0: + raise RuntimeError(f"Command failed with exit code {result.returncode}: {redacted}") # --------------------------------------------------------------------------- @@ -329,6 +344,12 @@ def _validate_args(args: argparse.Namespace) -> None: args.cmake_extra_defines = ( [f"-D{d}" for j in args.cmake_extra_defines for d in j] if args.cmake_extra_defines else [] ) + for define in args.cmake_extra_defines: + if define.startswith("-DFOUNDRY_LOCAL_TELEMETRY_TOKEN"): + raise ValueError( + "FOUNDRY_LOCAL_TELEMETRY_TOKEN must be supplied as an environment variable, " + "not via --cmake_extra_defines or command-line arguments." + ) # Cross-compilation if args.arm64: diff --git a/sdk_v2/cpp/cmake/FindOnnxRuntime.cmake b/sdk_v2/cpp/cmake/FindOnnxRuntime.cmake index f11385873..5ed855cf1 100644 --- a/sdk_v2/cpp/cmake/FindOnnxRuntime.cmake +++ b/sdk_v2/cpp/cmake/FindOnnxRuntime.cmake @@ -97,7 +97,7 @@ else() message(STATUS "Downloading ${ORT_PACKAGE_NAME} ${ORT_VERSION} from nuget.org") endif() else() - message(STATUS "Using pre-configured ORT_FETCH_URL: ${ORT_FETCH_URL}") + message(STATUS "Using pre-configured ORT_FETCH_URL") endif() # Normalize backslashes (Windows paths) and handle .nupkg extension @@ -106,11 +106,15 @@ else() set(_ORT_ZIP_PATH "${CMAKE_BINARY_DIR}/_deps/ortlib-download/ort.zip") get_filename_component(_ORT_ZIP_DIR "${_ORT_ZIP_PATH}" DIRECTORY) file(MAKE_DIRECTORY "${_ORT_ZIP_DIR}") - file(COPY_FILE "${ORT_FETCH_URL}" "${_ORT_ZIP_PATH}") + configure_file("${ORT_FETCH_URL}" "${_ORT_ZIP_PATH}" COPYONLY) set(ORT_FETCH_URL "${_ORT_ZIP_PATH}") endif() - FetchContent_Declare(ortlib URL ${ORT_FETCH_URL} DOWNLOAD_EXTRACT_TIMESTAMP TRUE DOWNLOAD_NAME ort.zip) + set(_ORT_FETCH_ARGS URL ${ORT_FETCH_URL} DOWNLOAD_NAME ort.zip) + if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.24") + list(APPEND _ORT_FETCH_ARGS DOWNLOAD_EXTRACT_TIMESTAMP TRUE) + endif() + FetchContent_Declare(ortlib ${_ORT_FETCH_ARGS}) FetchContent_MakeAvailable(ortlib) set(_ORT_HEADER_DIR "${ortlib_SOURCE_DIR}/build/native/include") @@ -149,7 +153,7 @@ else() message(STATUS "Downloading ${_ORT_GPU_LINUX_PACKAGE} ${ORT_VERSION} from nuget.org") endif() else() - message(STATUS "Using pre-configured ORT_GPU_LINUX_FETCH_URL: ${ORT_GPU_LINUX_FETCH_URL}") + message(STATUS "Using pre-configured ORT_GPU_LINUX_FETCH_URL") endif() # Normalize backslashes and handle .nupkg extension @@ -158,11 +162,15 @@ else() set(_ORT_GPU_ZIP_PATH "${CMAKE_BINARY_DIR}/_deps/ort_gpu_linux-download/ort_gpu_linux.zip") get_filename_component(_ORT_GPU_ZIP_DIR "${_ORT_GPU_ZIP_PATH}" DIRECTORY) file(MAKE_DIRECTORY "${_ORT_GPU_ZIP_DIR}") - file(COPY_FILE "${ORT_GPU_LINUX_FETCH_URL}" "${_ORT_GPU_ZIP_PATH}") + configure_file("${ORT_GPU_LINUX_FETCH_URL}" "${_ORT_GPU_ZIP_PATH}" COPYONLY) set(ORT_GPU_LINUX_FETCH_URL "${_ORT_GPU_ZIP_PATH}") endif() - FetchContent_Declare(ort_gpu_linux URL ${ORT_GPU_LINUX_FETCH_URL} DOWNLOAD_EXTRACT_TIMESTAMP TRUE DOWNLOAD_NAME ort_gpu_linux.zip) + set(_ORT_GPU_FETCH_ARGS URL ${ORT_GPU_LINUX_FETCH_URL} DOWNLOAD_NAME ort_gpu_linux.zip) + if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.24") + list(APPEND _ORT_GPU_FETCH_ARGS DOWNLOAD_EXTRACT_TIMESTAMP TRUE) + endif() + FetchContent_Declare(ort_gpu_linux ${_ORT_GPU_FETCH_ARGS}) FetchContent_MakeAvailable(ort_gpu_linux) set(_ORT_LIB_DIR "${ort_gpu_linux_SOURCE_DIR}/runtimes/${_ORT_PLATFORM}/native") diff --git a/sdk_v2/cpp/cmake/FindOnnxRuntimeGenAI.cmake b/sdk_v2/cpp/cmake/FindOnnxRuntimeGenAI.cmake index f83002d64..f8bdc4b10 100644 --- a/sdk_v2/cpp/cmake/FindOnnxRuntimeGenAI.cmake +++ b/sdk_v2/cpp/cmake/FindOnnxRuntimeGenAI.cmake @@ -160,7 +160,7 @@ else() message(STATUS "Downloading ${_GENAI_PACKAGE_NAME} ${ORT_GENAI_VERSION} from nuget.org") endif() else() - message(STATUS "Using caller-provided GENAI_FETCH_URL: ${GENAI_FETCH_URL}") + message(STATUS "Using caller-provided GENAI_FETCH_URL") endif() # Normalize to forward slashes — backslashes from Windows paths cause CMake @@ -173,15 +173,19 @@ else() set(_GENAI_ZIP_PATH "${CMAKE_BINARY_DIR}/_deps/genai-download/genai.zip") get_filename_component(_GENAI_ZIP_DIR "${_GENAI_ZIP_PATH}" DIRECTORY) file(MAKE_DIRECTORY "${_GENAI_ZIP_DIR}") - file(COPY_FILE "${GENAI_FETCH_URL}" "${_GENAI_ZIP_PATH}") + configure_file("${GENAI_FETCH_URL}" "${_GENAI_ZIP_PATH}" COPYONLY) set(GENAI_FETCH_URL "${_GENAI_ZIP_PATH}") - message(STATUS "Copied .nupkg to .zip for CMake extraction: ${GENAI_FETCH_URL}") + message(STATUS "Copied local GenAI .nupkg to .zip for CMake extraction") endif() - FetchContent_Declare(genailib + set(_GENAI_FETCH_ARGS URL ${GENAI_FETCH_URL} - DOWNLOAD_EXTRACT_TIMESTAMP TRUE - DOWNLOAD_NAME genai.zip # .nupkg is a ZIP; force CMake to recognize the format + DOWNLOAD_NAME genai.zip) + if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.24") + list(APPEND _GENAI_FETCH_ARGS DOWNLOAD_EXTRACT_TIMESTAMP TRUE) + endif() + FetchContent_Declare(genailib + ${_GENAI_FETCH_ARGS} ) FetchContent_MakeAvailable(genailib) diff --git a/sdk_v2/cpp/cmake/FindWinMLEpCatalog.cmake b/sdk_v2/cpp/cmake/FindWinMLEpCatalog.cmake index dc732a4ff..4bf4eb0f7 100644 --- a/sdk_v2/cpp/cmake/FindWinMLEpCatalog.cmake +++ b/sdk_v2/cpp/cmake/FindWinMLEpCatalog.cmake @@ -60,30 +60,32 @@ if(NOT _WINML_PLATFORM_UPPER MATCHES "^(AMD64|X64|X86_64|ARM64|AARCH64)$") return() endif() -include(cmake/nuget.cmake) - # WINML_EP_CATALOG_FETCH_URL can be set externally (e.g. for CI where nuget.org is blocked). set(WINML_EP_CATALOG_FETCH_URL "" CACHE STRING "Override URL or local path for the WinML EP Catalog NuGet package") -if(WINML_EP_CATALOG_FETCH_URL) - # Use FetchContent to download/extract the pre-downloaded package - include(FetchContent) - string(REPLACE "\\" "/" WINML_EP_CATALOG_FETCH_URL "${WINML_EP_CATALOG_FETCH_URL}") - if(WINML_EP_CATALOG_FETCH_URL MATCHES "\\.nupkg$" AND NOT WINML_EP_CATALOG_FETCH_URL MATCHES "^https?://") - set(_WINML_ZIP_PATH "${CMAKE_BINARY_DIR}/_deps/winml_ep_catalog-download/winml_ep_catalog.zip") - get_filename_component(_WINML_ZIP_DIR "${_WINML_ZIP_PATH}" DIRECTORY) - file(MAKE_DIRECTORY "${_WINML_ZIP_DIR}") - file(COPY_FILE "${WINML_EP_CATALOG_FETCH_URL}" "${_WINML_ZIP_PATH}") - set(WINML_EP_CATALOG_FETCH_URL "${_WINML_ZIP_PATH}") - endif() - FetchContent_Declare(winml_ep_catalog URL ${WINML_EP_CATALOG_FETCH_URL} DOWNLOAD_EXTRACT_TIMESTAMP TRUE DOWNLOAD_NAME winml_ep_catalog.zip) - FetchContent_MakeAvailable(winml_ep_catalog) - set(_WINML_EP_ROOT "${winml_ep_catalog_SOURCE_DIR}") - message(STATUS "WinML EP Catalog via FetchContent: ${_WINML_EP_ROOT}") -else() - install_nuget_package(Microsoft.Windows.AI.MachineLearning ${WINML_EP_CATALOG_VERSION} _WINML_EP_ROOT - SOURCE https://api.nuget.org/v3/index.json) +include(FetchContent) +set(_WINML_EP_CATALOG_URL "${WINML_EP_CATALOG_FETCH_URL}") +if(NOT _WINML_EP_CATALOG_URL) + set(_WINML_EP_CATALOG_URL + "https://api.nuget.org/v3-flatcontainer/microsoft.windows.ai.machinelearning/${WINML_EP_CATALOG_VERSION}/microsoft.windows.ai.machinelearning.${WINML_EP_CATALOG_VERSION}.nupkg") +endif() + +string(REPLACE "\\" "/" _WINML_EP_CATALOG_URL "${_WINML_EP_CATALOG_URL}") +if(_WINML_EP_CATALOG_URL MATCHES "\\.nupkg$" AND NOT _WINML_EP_CATALOG_URL MATCHES "^https?://") + set(_WINML_ZIP_PATH "${CMAKE_BINARY_DIR}/_deps/winml_ep_catalog-download/winml_ep_catalog.zip") + get_filename_component(_WINML_ZIP_DIR "${_WINML_ZIP_PATH}" DIRECTORY) + file(MAKE_DIRECTORY "${_WINML_ZIP_DIR}") + configure_file("${_WINML_EP_CATALOG_URL}" "${_WINML_ZIP_PATH}" COPYONLY) + set(_WINML_EP_CATALOG_URL "${_WINML_ZIP_PATH}") +endif() +set(_WINML_FETCH_ARGS URL ${_WINML_EP_CATALOG_URL} DOWNLOAD_NAME winml_ep_catalog.zip) +if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.24") + list(APPEND _WINML_FETCH_ARGS DOWNLOAD_EXTRACT_TIMESTAMP TRUE) endif() +FetchContent_Declare(winml_ep_catalog ${_WINML_FETCH_ARGS}) +FetchContent_MakeAvailable(winml_ep_catalog) +set(_WINML_EP_ROOT "${winml_ep_catalog_SOURCE_DIR}") +message(STATUS "WinML EP Catalog via FetchContent: ${_WINML_EP_ROOT}") # Load the package's first-party CMake config for target discovery and layout # resolution. The config lives at build/cmake/-config.cmake diff --git a/sdk_v2/cpp/nuget/pack.py b/sdk_v2/cpp/nuget/pack.py index 902ba288a..ec27eaf3b 100644 --- a/sdk_v2/cpp/nuget/pack.py +++ b/sdk_v2/cpp/nuget/pack.py @@ -55,11 +55,16 @@ "osx_arm64": ("osx-arm64", "libfoundry_local.dylib"), } -# Sibling files copied into runtimes//native/ when present in the upstream -# artifact. Windows builds drop Microsoft.Windows.AI.MachineLearning.dll alongside -# foundry_local.dll; other platforms don't, so presence alone drives inclusion. +# Optional sibling files copied into runtimes//native/ when present in the upstream +# artifact. Non-Windows platforms do not have these files, so presence alone drives inclusion there. OPTIONAL_SIBLINGS: tuple[str, ...] = ( +) + +# Required Windows runtime siblings. The reg-free WinML runtime depends on DirectML, so a Windows RID +# with only Microsoft.Windows.AI.MachineLearning.dll is incomplete and must fail packaging. +WINDOWS_REQUIRED_SIBLINGS: tuple[str, ...] = ( "Microsoft.Windows.AI.MachineLearning.dll", + "DirectML.dll", ) log = logging.getLogger("pack") @@ -159,6 +164,16 @@ def stage(args: argparse.Namespace, staging: Path) -> int: log.info(" %s → runtimes/%s/native/%s", lib_path, rid, lib_path.name) staged = 1 + required_siblings = WINDOWS_REQUIRED_SIBLINGS if rid.startswith("win-") else () + for sibling in required_siblings: + sibling_path = src_dir / sibling + if not sibling_path.is_file(): + log.error("Expected Windows runtime dependency %s at %s but file not found.", sibling, src_dir) + sys.exit(1) + shutil.copy2(sibling_path, native_dir) + log.info(" %s → runtimes/%s/native/%s", sibling_path, rid, sibling) + staged += 1 + for sibling in OPTIONAL_SIBLINGS: sibling_path = src_dir / sibling if sibling_path.is_file(): diff --git a/sdk_v2/cpp/src/catalog/azure_model_catalog.cc b/sdk_v2/cpp/src/catalog/azure_model_catalog.cc index 45c73842c..7d83e59d6 100644 --- a/sdk_v2/cpp/src/catalog/azure_model_catalog.cc +++ b/sdk_v2/cpp/src/catalog/azure_model_catalog.cc @@ -24,8 +24,7 @@ namespace { /// Split a catalog URL into structured telemetry dimensions. The Azure Foundry /// catalog URL looks like "https://ai.azure.com/api//", e.g. /// "https://ai.azure.com/api/eastus/ux/v1.0" -> {ai.azure.com, eastus, ux/v1.0}. -/// Custom URLs that don't follow the "/api//" convention keep an empty -/// region and put the whole path in `format`. The embedded snapshot is "static". +/// Custom/private URLs are bucketed so telemetry does not disclose hostnames or paths. struct ParsedCatalogUrl { std::string endpoint; std::string region; @@ -56,9 +55,6 @@ ParsedCatalogUrl ParseCatalogUrl(const std::string& url) { if (auto at = out.endpoint.rfind('@'); at != std::string::npos) { out.endpoint = out.endpoint.substr(at + 1); } - if (out.endpoint != "ai.azure.com") { - return {"custom", "", ""}; - } if (auto q = path.find_first_of("?#"); q != std::string::npos) { path = path.substr(0, q); @@ -77,12 +73,13 @@ ParsedCatalogUrl ParseCatalogUrl(const std::string& url) { pos = next + 1; } - size_t format_start = 0; - if (segments.size() >= 2 && segments[0] == "api") { - out.region = segments[1]; - format_start = 2; + if (out.endpoint != "ai.azure.com" || segments.size() < 4 || segments[0] != "api" || + segments[2] != "ux" || segments[3].empty() || segments[3][0] != 'v') { + return {"custom", "", ""}; } - for (size_t i = format_start; i < segments.size(); ++i) { + + out.region = segments[1]; + for (size_t i = 2; i < segments.size(); ++i) { if (!out.format.empty()) { out.format += '/'; } @@ -186,9 +183,8 @@ std::vector AzureModelCatalog::FetchModels() const { base_info.region = parsed.region; base_info.format = parsed.format; base_info.correlation_id = correlation_id; - base_info.user_agent = DefaultUserAgent(); - auto model_infos = FetchAllModelInfosWithCachedModels(*client, cached_model_ids, logger_, telemetry_, base_info); + auto model_infos = FetchAllModelInfosWithCachedModels(*client, cached_model_ids, logger_, telemetry_, &base_info); for (const auto& info : model_infos) { // Check if the model is locally cached and pass the path if so. @@ -203,16 +199,24 @@ std::vector AzureModelCatalog::FetchModels() const { } }; + bool fetched_any_catalog = false; for (const auto& [url, filter] : catalog_urls_) { try { fetch_from(url, filter); + fetched_any_catalog = true; } catch (const std::exception& ex) { // One failing URL shouldn't block others — skip and continue. + auto parsed = ParseCatalogUrl(url); logger_.Log(LogLevel::Error, - fmt::format("failed to fetch catalog from {}: {}", url, ex.what())); + fmt::format("failed to fetch catalog from {}: {}", parsed.endpoint, + ScrubStringForTelemetry(ex.what()))); } } + if (!fetched_any_catalog) { + FL_THROW(FOUNDRY_LOCAL_ERROR_NETWORK, "failed to fetch any configured catalog source"); + } + logger_.Log(LogLevel::Information, fmt::format("Populated model info for {} models.", models.size())); @@ -249,8 +253,10 @@ std::vector AzureModelCatalog::FetchModelVersions( out.push_back(model_factory_(std::move(info), /*local_path=*/"")); } } catch (const std::exception& ex) { + auto parsed = ParseCatalogUrl(url); logger_.Log(LogLevel::Error, - fmt::format("FetchModelVersions: failed to query {} — {}", url, ex.what())); + fmt::format("FetchModelVersions: failed to query {} — {}", parsed.endpoint, + ScrubStringForTelemetry(ex.what()))); } } @@ -305,8 +311,10 @@ std::vector AzureModelCatalog::FetchModelsByIds(const std::vector #include "exception.h" +#include "telemetry/telemetry_redaction.h" #include #include @@ -20,6 +21,11 @@ BaseModelCatalog::BaseModelCatalog(std::string name, ILogger& logger) BaseModelCatalog::~BaseModelCatalog() = default; void BaseModelCatalog::PopulateModels(std::vector variants) const { + if (populated_) { + IntegrateVariantsLocked(std::move(variants)); + return; + } + // Group variants by alias into Model containers. // Matches C# Catalog.UpdateModels() pattern: // foreach (modelInfo) { find or create Model by alias, add variant } @@ -47,50 +53,25 @@ void BaseModelCatalog::PopulateModels(std::vector variants) const { model.SelectDefaultVariant(); } - // On refresh: merge new models into stable storage. Existing models keep their addresses. - // New aliases are appended. Existing aliases are left unchanged (their Model* stays valid). - if (populated_) { - // Build a set of existing aliases for fast lookup. - std::unordered_map existing_aliases; - for (auto& m : models_) { - existing_aliases[m->Alias()] = m.get(); - } - - size_t new_count = 0; - for (auto& [alias, model] : alias_to_model) { - if (!existing_aliases.contains(alias)) { - models_.push_back(std::make_unique(std::move(model))); - ++new_count; - } - } - - if (new_count > 0) { - logger_.Log(LogLevel::Information, - fmt::format("Catalog '{}' refresh: {} new model(s) added, {} total.", - name_, new_count, models_.size())); - } else { - logger_.Log(LogLevel::Debug, - fmt::format("Catalog '{}' refresh: no new models. {} total.", - name_, models_.size())); - } - } else { - // Initial population: move all models into stable storage. - models_.reserve(alias_to_model.size()); - for (auto& [alias, model] : alias_to_model) { - models_.push_back(std::make_unique(std::move(model))); - } - - logger_.Log(LogLevel::Debug, - fmt::format("Catalog '{}' populated with {} model(s).", name_, models_.size())); + // Initial population: move all models into stable storage. + models_.reserve(alias_to_model.size()); + for (auto& [alias, model] : alias_to_model) { + models_.push_back(std::make_unique(std::move(model))); } + logger_.Log(LogLevel::Debug, + fmt::format("Catalog '{}' populated with {} model(s).", name_, models_.size())); + RebuildIndex(); populated_ = true; } void BaseModelCatalog::IntegrateVariants(std::vector variants) const { std::lock_guard lock(mutex_); + IntegrateVariantsLocked(std::move(variants)); +} +void BaseModelCatalog::IntegrateVariantsLocked(std::vector variants) const { if (variants.empty()) { return; } @@ -141,6 +122,9 @@ void BaseModelCatalog::IntegrateVariants(std::vector variants) const { it->second->AddVariant(std::move(v)); ++added_variants; } + if (!it->second->HasExplicitVariantSelection()) { + it->second->SelectDefaultVariant(); + } } else { // New alias: build a container and choose default after all variants are added. auto first = std::move(alias_variants.front()); @@ -214,6 +198,7 @@ void BaseModelCatalog::InvalidateCache() { // This is called after EP registration changes — the catalog needs to // re-query with updated device/EP filters. std::lock_guard lock(mutex_); + force_refresh_ = true; next_refresh_at_ = std::chrono::steady_clock::time_point{}; } @@ -224,8 +209,9 @@ void BaseModelCatalog::EnsurePopulated(bool allow_refresh) const { // not worth the complexity to optimise.) std::lock_guard lock(mutex_); - bool needs_refresh = allow_refresh && - std::chrono::steady_clock::now() >= next_refresh_at_; + auto now = std::chrono::steady_clock::now(); + bool needs_refresh = (force_refresh_ && now >= next_refresh_at_) || + (allow_refresh && now >= next_refresh_at_); if (populated_ && !needs_refresh) { return; @@ -236,8 +222,21 @@ void BaseModelCatalog::EnsurePopulated(bool allow_refresh) const { fmt::format("Catalog '{}' refreshing (cache expired).", name_)); } - auto variants = FetchModels(); + std::vector variants; + try { + variants = FetchModels(); + } catch (const std::exception& ex) { + if (populated_) { + logger_.Log(LogLevel::Warning, + fmt::format("Catalog '{}' refresh failed; keeping existing catalog: {}", + name_, ScrubStringForTelemetry(ex.what()))); + next_refresh_at_ = std::chrono::steady_clock::now() + std::chrono::minutes(1); + return; + } + throw; + } PopulateModels(std::move(variants)); + force_refresh_ = false; next_refresh_at_ = std::chrono::steady_clock::now() + kCacheDuration; } @@ -333,8 +332,10 @@ Model* BaseModelCatalog::GetLatestVersion(const Model* model) const { auto alias_it = idx->alias_index.find(std::string(model->Info().alias)); if (alias_it != idx->alias_index.end()) { auto variants = alias_it->second->Variants(); - if (!variants.empty()) { - return variants.front(); + for (auto* variant : variants) { + if (variant != nullptr && variant->Info().name == model->Info().name) { + return variant; + } } } diff --git a/sdk_v2/cpp/src/catalog/base_model_catalog.h b/sdk_v2/cpp/src/catalog/base_model_catalog.h index 97411e395..d7adb5d80 100644 --- a/sdk_v2/cpp/src/catalog/base_model_catalog.h +++ b/sdk_v2/cpp/src/catalog/base_model_catalog.h @@ -98,6 +98,7 @@ class BaseModelCatalog : public ICatalog { std::shared_ptr GetIndex() const; mutable bool populated_ = false; + mutable bool force_refresh_ = false; mutable std::mutex mutex_; mutable std::chrono::steady_clock::time_point next_refresh_at_{}; @@ -111,6 +112,7 @@ class BaseModelCatalog : public ICatalog { /// present. For new aliases, creates a new container. Rebuilds the lookup /// index when the model set actually changed. void IntegrateVariants(std::vector variants) const; + void IntegrateVariantsLocked(std::vector variants) const; /// Build lookup indices from the current models_ collection. /// Builds a complete new ModelIndex locally, then atomically swaps it into index_. diff --git a/sdk_v2/cpp/src/catalog/catalog_client.cc b/sdk_v2/cpp/src/catalog/catalog_client.cc index 7cab4f4f1..acc388b0d 100644 --- a/sdk_v2/cpp/src/catalog/catalog_client.cc +++ b/sdk_v2/cpp/src/catalog/catalog_client.cc @@ -3,6 +3,7 @@ #include "catalog/catalog_client.h" #include "ep_detection/ep_detector.h" #include "telemetry/telemetry.h" +#include "telemetry/telemetry_redaction.h" #include "utils.h" #include @@ -26,7 +27,7 @@ std::vector FetchAllModelInfosWithCachedModels( const std::vector& cached_model_ids, ILogger& logger, ITelemetry& telemetry, - const CatalogFetchInfo& base_info) { + const CatalogFetchInfo* base_info) { auto now = [] { return std::chrono::steady_clock::now(); }; auto elapsed_ms = [](std::chrono::steady_clock::time_point start) { return std::chrono::duration_cast(std::chrono::steady_clock::now() - start) @@ -34,7 +35,10 @@ std::vector FetchAllModelInfosWithCachedModels( }; auto emit = [&](const std::string& operation, ActionStatus status, int64_t duration_ms, int32_t model_count, const std::string& error) { - CatalogFetchInfo info = base_info; + if (base_info == nullptr) { + return; + } + CatalogFetchInfo info = *base_info; info.operation = operation; info.status = status; info.duration_ms = duration_ms; @@ -53,7 +57,9 @@ std::vector FetchAllModelInfosWithCachedModels( const auto start = now(); try { result = client.FetchAllModelInfos(); - } catch (const std::exception&) { + } catch (const std::exception& ex) { + logger.Log(LogLevel::Warning, + fmt::format("catalog: failed to fetch models — {}", ScrubStringForTelemetry(ex.what()))); emit("FetchAll", ActionStatus::kDependencyFailure, elapsed_ms(start), 0, kCatalogFetchFailure); throw; } @@ -89,8 +95,9 @@ std::vector FetchAllModelInfosWithCachedModels( } emit("FetchByIds", ActionStatus::kSuccess, elapsed_ms(start), additional_count, ""); } catch (const std::exception& ex) { + auto error_message = ScrubStringForTelemetry(ex.what()); logger.Log(LogLevel::Warning, - fmt::format("catalog: failed to fetch cached model IDs — {}", ex.what())); + fmt::format("catalog: failed to fetch cached model IDs — {}", error_message)); emit("FetchByIds", ActionStatus::kDependencyFailure, elapsed_ms(start), 0, kCatalogFetchFailure); } catch (...) { logger.Log(LogLevel::Warning, "catalog: failed to fetch cached model IDs — unknown error"); diff --git a/sdk_v2/cpp/src/catalog/catalog_client.h b/sdk_v2/cpp/src/catalog/catalog_client.h index 4438ad812..97ab70a5e 100644 --- a/sdk_v2/cpp/src/catalog/catalog_client.h +++ b/sdk_v2/cpp/src/catalog/catalog_client.h @@ -53,15 +53,16 @@ class ICatalogClient { }; /// Production helper that combines a catalog fetch with locally cached model -/// resolution and BYO synthesis. Emits a CatalogFetch event for the primary fetch -/// and (if it runs) the cached-id lookup, copying the endpoint/region/format/correlation -/// fields from `base_info` and filling in operation/status/duration/model_count/error. +/// resolution and BYO synthesis. When `base_info` is provided, emits a CatalogFetch +/// event for the primary fetch and (if it runs) the cached-id lookup, copying the +/// endpoint/region/format/correlation fields from `base_info` and filling in +/// operation/status/duration/model_count/error. std::vector FetchAllModelInfosWithCachedModels( ICatalogClient& client, const std::vector& cached_model_ids, ILogger& logger, ITelemetry& telemetry, - const CatalogFetchInfo& base_info); + const CatalogFetchInfo* base_info = nullptr); /// Construct a client for the live Azure Foundry catalog. /// - `ep_detector` limits results to models supported by this machine. diff --git a/sdk_v2/cpp/src/download/blob_download_state.cc b/sdk_v2/cpp/src/download/blob_download_state.cc index 07680bd73..b76ff7180 100644 --- a/sdk_v2/cpp/src/download/blob_download_state.cc +++ b/sdk_v2/cpp/src/download/blob_download_state.cc @@ -21,7 +21,7 @@ constexpr const char* kStateFileExtension = ".dlstate"; // bytes | field // -------|-------------------------------------------------------- // 0..3 | magic "FLDS" -// 4 | version (currently 1) +// 4 | version (currently 2) // 5..12 | blob_size (int64) // 13..16 | chunk_size (int32) // 17..20 | total_chunks (int32) @@ -29,14 +29,17 @@ constexpr const char* kStateFileExtension = ".dlstate"; // 25..28 | highest_completed_chunk (int32) // 29..32 | completed_count (int32) // 33..40 | last_modified_unix_ms (int64) -// 41..44 | trunc_bitmap_byte_len (uint32) -// 45.. | trunc_bitmap_byte_len bytes of bitmap data, copied directly out of +// 41..44 | blob_identity_len (uint32) +// 45.. | blob_identity bytes +// ... | trunc_bitmap_byte_len (uint32) +// ... | trunc_bitmap_byte_len bytes of bitmap data, copied directly out of // full_completion_bitmap starting at the byte offset implied by // bitmap_byte_aligned_start. constexpr char kMagic[4] = {'F', 'L', 'D', 'S'}; -constexpr uint8_t kVersion = 1; +constexpr uint8_t kVersion = 2; constexpr int32_t kBitsPerWord = 64; +constexpr uint32_t kMaxBlobIdentityLength = 4096; // Serialize a scalar field in host byte order. Every target we build for // (x64 / arm64) is little-endian, so the on-disk layout is little-endian in @@ -54,6 +57,28 @@ bool ReadNative(std::istream& in, T& out_value) { return static_cast(in); } +void WriteString(std::ostream& out, const std::string& value) { + WriteNative(out, static_cast(value.size())); + if (!value.empty()) { + out.write(value.data(), static_cast(value.size())); + } +} + +bool ReadString(std::istream& in, std::string& value) { + uint32_t size = 0; + if (!ReadNative(in, size)) { + return false; + } + if (size > kMaxBlobIdentityLength) { + return false; + } + value.resize(size); + if (size > 0) { + in.read(value.data(), size); + } + return static_cast(in); +} + int64_t NowUnixMs() { return std::chrono::duration_cast( std::chrono::system_clock::now().time_since_epoch()) @@ -72,13 +97,15 @@ std::unique_ptr BlobDownloadState::CreateNew(std::string blob const std::filesystem::path& local_file_path, int64_t blob_size, int32_t chunk_size, - int32_t total_chunks) { + int32_t total_chunks, + std::string blob_identity) { auto state = std::make_unique(); state->blob_name = std::move(blob_name); state->local_file_path = local_file_path.string(); state->blob_size = blob_size; state->chunk_size = chunk_size; state->total_chunks = total_chunks; + state->blob_identity = std::move(blob_identity); state->bitmap_byte_aligned_start = 0; state->highest_completed_chunk = -1; state->completed_count = 0; @@ -93,6 +120,7 @@ std::unique_ptr BlobDownloadState::LoadState(std::string blob int64_t expected_blob_size, int32_t expected_chunk_size, int32_t expected_total_chunks, + std::string_view expected_blob_identity, ILogger& logger) { auto state_path = GetStateFilePath(local_file_path); std::error_code ec; @@ -122,17 +150,20 @@ std::unique_ptr BlobDownloadState::LoadState(std::string blob int32_t highest_completed_chunk = 0; int32_t completed_count = 0; int64_t last_modified_unix_ms = 0; + std::string blob_identity; uint32_t trunc_len = 0; if (!ReadNative(in, blob_size) || !ReadNative(in, chunk_size) || !ReadNative(in, total_chunks) || !ReadNative(in, bitmap_byte_aligned_start) || !ReadNative(in, highest_completed_chunk) || - !ReadNative(in, completed_count) || !ReadNative(in, last_modified_unix_ms) || !ReadNative(in, trunc_len)) { + !ReadNative(in, completed_count) || !ReadNative(in, last_modified_unix_ms) || + !ReadString(in, blob_identity) || !ReadNative(in, trunc_len)) { logger.Log(LogLevel::Warning, "Download state header truncated: " + state_path.string()); return nullptr; } // Sanity / compatibility checks. if (blob_size != expected_blob_size || chunk_size != expected_chunk_size || - total_chunks != expected_total_chunks) { + total_chunks != expected_total_chunks || + (!expected_blob_identity.empty() && blob_identity != expected_blob_identity)) { logger.Log(LogLevel::Information, "Download state for " + state_path.string() + " is incompatible with current blob layout; starting fresh"); @@ -191,6 +222,7 @@ std::unique_ptr BlobDownloadState::LoadState(std::string blob state->highest_completed_chunk = highest_completed_chunk; state->completed_count = completed_count; state->last_modified_unix_ms = last_modified_unix_ms; + state->blob_identity = std::move(blob_identity); state->full_completion_bitmap = std::move(bitmap); logger.Log(LogLevel::Information, @@ -200,6 +232,16 @@ std::unique_ptr BlobDownloadState::LoadState(std::string blob return state; } +std::unique_ptr BlobDownloadState::LoadState(std::string blob_name, + const std::filesystem::path& local_file_path, + int64_t expected_blob_size, + int32_t expected_chunk_size, + int32_t expected_total_chunks, + ILogger& logger) { + return LoadState(std::move(blob_name), local_file_path, expected_blob_size, expected_chunk_size, + expected_total_chunks, {}, logger); +} + int64_t BlobDownloadState::CalculateDownloadedSize() const noexcept { int64_t bytes = static_cast(completed_count) * chunk_size; // If the final chunk is partial and was completed, adjust the overcount. @@ -321,6 +363,7 @@ bool BlobDownloadState::SaveState(ILogger& logger) { WriteNative(out, highest_completed_chunk); WriteNative(out, completed_count); WriteNative(out, last_modified_unix_ms); + WriteString(out, blob_identity); WriteNative(out, trunc_len); if (trunc_len > 0) { auto* src = reinterpret_cast(full_completion_bitmap.data()) + byte_offset; diff --git a/sdk_v2/cpp/src/download/blob_download_state.h b/sdk_v2/cpp/src/download/blob_download_state.h index 54f114340..acd43be99 100644 --- a/sdk_v2/cpp/src/download/blob_download_state.h +++ b/sdk_v2/cpp/src/download/blob_download_state.h @@ -7,6 +7,7 @@ #include #include #include +#include #include namespace fl { @@ -35,6 +36,7 @@ class BlobDownloadState { int64_t blob_size = 0; int32_t chunk_size = 0; int32_t total_chunks = 0; + std::string blob_identity; /// Serialization marker (always a multiple of 8): chunks below this index are /// complete and dropped from the sidecar's truncated bitmap. The in-memory @@ -64,7 +66,8 @@ class BlobDownloadState { const std::filesystem::path& local_file_path, int64_t blob_size, int32_t chunk_size, - int32_t total_chunks); + int32_t total_chunks, + std::string blob_identity = {}); /// Load existing state from `.dlstate`. Returns nullptr if /// the file does not exist, is corrupted, or has incompatible @@ -73,6 +76,13 @@ class BlobDownloadState { /// and the partial download is no longer valid). /// `logger` receives diagnostics for corrupt/incompatible state files. Required: the /// downloader always has a logger, so there is no optional/null case to handle. + static std::unique_ptr LoadState(std::string blob_name, + const std::filesystem::path& local_file_path, + int64_t expected_blob_size, + int32_t expected_chunk_size, + int32_t expected_total_chunks, + std::string_view expected_blob_identity, + ILogger& logger); static std::unique_ptr LoadState(std::string blob_name, const std::filesystem::path& local_file_path, int64_t expected_blob_size, diff --git a/sdk_v2/cpp/src/download/blob_downloader.cc b/sdk_v2/cpp/src/download/blob_downloader.cc index 8719334d9..af215e704 100644 --- a/sdk_v2/cpp/src/download/blob_downloader.cc +++ b/sdk_v2/cpp/src/download/blob_downloader.cc @@ -49,6 +49,7 @@ constexpr size_t kStreamingBufferBytes = 64 * 1024; struct AzureBlobDownloader::ChunkContext { const Azure::Storage::Blobs::BlobClient& blob_client; const Azure::Core::Context& azure_ctx; + std::string blob_identity; }; AzureBlobDownloader::AzureBlobDownloader(ILogger& logger) : logger_(logger) {} @@ -63,6 +64,11 @@ std::vector AzureBlobDownloader::ListBlobs(const std::string& sas_ BlobItemInfo info; info.name = blob.Name; info.content_length = blob.BlobSize; + if (blob.Details.ETag.HasValue()) { + info.blob_identity = blob.Details.ETag.ToString(); + } else if (blob.VersionId.HasValue()) { + info.blob_identity = blob.VersionId.Value(); + } items.push_back(std::move(info)); } } @@ -74,9 +80,14 @@ std::vector AzureBlobDownloader::ListBlobs(const std::string& sas_ } } -int64_t AzureBlobDownloader::GetBlobSize(ChunkContext& ctx) { +AzureBlobDownloader::BlobProperties AzureBlobDownloader::GetBlobProperties(ChunkContext& ctx) { auto props = ctx.blob_client.GetProperties({}, ctx.azure_ctx).Value; - return props.BlobSize; + BlobProperties result; + result.content_length = props.BlobSize; + if (props.ETag.HasValue()) { + result.blob_identity = props.ETag.ToString(); + } + return result; } bool AzureBlobDownloader::IsCancellationRequested(const ChunkContext& ctx) const { @@ -88,6 +99,9 @@ void AzureBlobDownloader::DownloadChunkStreaming( const std::function& sink) { Azure::Storage::Blobs::DownloadBlobOptions range_opts; range_opts.Range = Azure::Core::Http::HttpRange{offset, size}; + if (!ctx.blob_identity.empty()) { + range_opts.AccessConditions.IfMatch = Azure::ETag(ctx.blob_identity); + } auto result = ctx.blob_client.Download(range_opts, ctx.azure_ctx); auto& body_stream = *result.Value.BodyStream; @@ -155,9 +169,11 @@ void AzureBlobDownloader::DownloadBlob(const std::string& sas_uri, // or by external cancellation; checked by workers between iterations. std::atomic internal_cancel{false}; - ChunkContext chunk_ctx{blob_client, azure_ctx}; + ChunkContext chunk_ctx{blob_client, azure_ctx, ""}; - int64_t blob_size = GetBlobSize(chunk_ctx); + auto blob_properties = GetBlobProperties(chunk_ctx); + int64_t blob_size = blob_properties.content_length; + chunk_ctx.blob_identity = std::move(blob_properties.blob_identity); if (blob_size == 0) { EnsureEmptyBlobFile(local_path); @@ -172,7 +188,7 @@ void AzureBlobDownloader::DownloadBlob(const std::string& sas_uri, // Resume from existing sidecar if it matches the current blob layout. auto state = BlobDownloadState::LoadState(blob_name, local_path, blob_size, static_cast(kChunkSize), - num_chunks, logger_); + num_chunks, chunk_ctx.blob_identity, logger_); if (state) { // Only trust the sidecar if the data file it describes is actually on disk // at full size. If the data file was truncated or removed (e.g. an external @@ -192,7 +208,8 @@ void AzureBlobDownloader::DownloadBlob(const std::string& sas_uri, if (!state) { state = BlobDownloadState::CreateNew(blob_name, local_path, blob_size, - static_cast(kChunkSize), num_chunks); + static_cast(kChunkSize), num_chunks, + chunk_ctx.blob_identity); // Persist the sidecar now, before Open() pre-allocates the data file. // IsDownloadNeeded treats "data file at full size + no sidecar" as a // completed download and skips it. The periodic save below does not run @@ -209,6 +226,28 @@ void AzureBlobDownloader::DownloadBlob(const std::string& sas_uri, } } + auto pending = state->GetPendingChunks(); + if (pending.empty()) { + // A complete sidecar means a previous attempt finished all chunks but did not reach finalization. Do not trust + // that state: a late close error may have left a full-size corrupt file. Start a fresh all-chunks pass. + logger_.Log(LogLevel::Information, + "Resume sidecar for '" + local_path + "' is complete but unfinalized; starting fresh"); + std::error_code remove_ec; + std::filesystem::remove(local_path, remove_ec); + if (remove_ec && std::filesystem::exists(local_path)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, + "failed to remove suspect completed blob before restart: " + local_path); + } + state = BlobDownloadState::CreateNew(blob_name, local_path, blob_size, + static_cast(kChunkSize), num_chunks, + chunk_ctx.blob_identity); + if (!state->SaveState(logger_)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, + "failed to persist reset download state for '" + local_path + "'"); + } + pending = state->GetPendingChunks(); + } + // Track cumulative bytes for progress reporting; seed with bytes already // present on disk so percent stays monotonic across resume. std::atomic bytes_completed{state->CalculateDownloadedSize()}; @@ -216,16 +255,6 @@ void AzureBlobDownloader::DownloadBlob(const std::string& sas_uri, bytes_written_cb(bytes_completed.load()); } - auto pending = state->GetPendingChunks(); - if (pending.empty()) { - // Already complete on disk — drop the sidecar. - BlobDownloadState::DeleteState(local_path, logger_); - if (bytes_written_cb) { - bytes_written_cb(blob_size); - } - return; - } - // Open the file writer once for the whole download. Open() pre-allocates // the file to blob_size if needed, preserving any existing bytes from a // resume. Concurrent WriteAt calls to disjoint ranges are thread-safe — the @@ -360,7 +389,54 @@ void AzureBlobDownloader::DownloadBlob(const std::string& sas_uri, // Release the OS handle before persisting / deleting the sidecar so any // observer that watches the data file sees a fully-closed handle. - writer.Close(); + try { + writer.Close(); + } catch (...) { + bool safe_to_delete_state = false; + std::error_code remove_ec; + if (std::filesystem::remove(local_path, remove_ec) || !std::filesystem::exists(local_path)) { + safe_to_delete_state = true; + } else { + logger_.Log(LogLevel::Warning, + "failed to remove blob file after close failure: " + local_path + + " (" + remove_ec.message() + ")"); + + const std::filesystem::path invalid_path = local_path + ".invalid"; + std::error_code cleanup_ec; + std::filesystem::remove(invalid_path, cleanup_ec); + std::error_code rename_ec; + std::filesystem::rename(local_path, invalid_path, rename_ec); + if (!rename_ec) { + safe_to_delete_state = true; + } else { + logger_.Log(LogLevel::Warning, + "failed to quarantine blob file after close failure: " + local_path + + " (" + rename_ec.message() + ")"); + + std::error_code resize_ec; + std::filesystem::resize_file(local_path, 0, resize_ec); + if (!resize_ec) { + safe_to_delete_state = true; + } else { + logger_.Log(LogLevel::Warning, + "failed to truncate blob file after close failure: " + local_path + + " (" + resize_ec.message() + ")"); + auto fresh_state = BlobDownloadState::CreateNew(blob_name, local_path, blob_size, + static_cast(kChunkSize), num_chunks, + chunk_ctx.blob_identity); + if (!fresh_state->SaveState(logger_)) { + logger_.Log(LogLevel::Error, + "failed to reset download state after close failure for '" + local_path + "'"); + } + } + } + } + + if (safe_to_delete_state) { + BlobDownloadState::DeleteState(local_path, logger_); + } + throw; + } const bool was_cancelled = cancelled && cancelled->load(std::memory_order_relaxed); if (first_error || was_cancelled) { @@ -410,12 +486,11 @@ std::string ComputeRelativePath(const std::string& prefix, const std::string& bl return blob_name.substr(trim); } -/// Returns false if a file at `local_path` already matches the blob's expected -/// `content_length` exactly AND has no `.dlstate` sidecar — in which case the -/// caller can skip the download. Returns true (download needed) for any of: -/// missing file, size mismatch, sidecar present (file may be pre-allocated -/// with holes), or filesystem-stat errors (treat as "redownload to be safe"). -bool IsDownloadNeeded(const BlobItemInfo& blob, const std::string& local_path) { +/// Returns false if a file at `local_path` already matches the blob's expected `content_length` exactly AND has no +/// `.dlstate` sidecar — in which case the caller can skip the download. Returns true (download needed) for any of: +/// missing file, size mismatch, sidecar present (file may be pre-allocated with holes), identity revalidation +/// required, or filesystem-stat errors (treat as "redownload to be safe"). +bool IsDownloadNeeded(const BlobItemInfo& blob, const std::string& local_path, bool require_completed_file_identity) { std::error_code ec; auto status = std::filesystem::status(local_path, ec); if (ec || !std::filesystem::exists(status) || !std::filesystem::is_regular_file(status)) { @@ -435,6 +510,9 @@ bool IsDownloadNeeded(const BlobItemInfo& blob, const std::string& local_path) { if (std::filesystem::exists(sidecar, ec)) { return true; } + if (require_completed_file_identity && !blob.blob_identity.empty()) { + return true; + } return false; } @@ -523,13 +601,14 @@ void DownloadBlobsToDirectory(IBlobDownloader& downloader, int32_t skipped_file_count = 0; blobs_to_download.erase( std::remove_if(blobs_to_download.begin(), blobs_to_download.end(), - [&skipped_bytes, &skipped_file_count](const auto& pair) { - if (IsDownloadNeeded(pair.first, pair.second)) { - return false; - } - skipped_bytes += pair.first.content_length; - ++skipped_file_count; - return true; + [&skipped_bytes, &skipped_file_count, &options](const auto& pair) { + if (!options.skip_completed_files || + IsDownloadNeeded(pair.first, pair.second, options.require_completed_file_identity)) { + return false; + } + skipped_bytes += pair.first.content_length; + ++skipped_file_count; + return true; }), blobs_to_download.end()); diff --git a/sdk_v2/cpp/src/download/blob_downloader.h b/sdk_v2/cpp/src/download/blob_downloader.h index ed2c60228..4824370b3 100644 --- a/sdk_v2/cpp/src/download/blob_downloader.h +++ b/sdk_v2/cpp/src/download/blob_downloader.h @@ -28,6 +28,16 @@ struct BlobDownloadOptions { /// Maximum concurrent chunk downloads per blob. Default matches C# desktop. int max_concurrency = 64; + /// When false, full-size files without a sidecar are redownloaded instead of + /// skipped. Use this while recovering an incomplete model directory. + bool skip_completed_files = false; + + /// When true, same-size files without a sidecar are redownloaded because their + /// remote blob identity cannot be verified. Used when repairing an incomplete + /// model directory, where previously-completed files may belong to an older + /// remote blob version. + bool require_completed_file_identity = false; + /// Progress callback (optional). Return non-zero to cancel the download. DownloadProgressFn progress; }; @@ -36,6 +46,7 @@ struct BlobDownloadOptions { struct BlobItemInfo { std::string name; int64_t content_length = 0; + std::string blob_identity{}; }; /// Interface for blob download operations. Enables testing via mock implementations. @@ -90,8 +101,14 @@ class AzureBlobDownloader : public IBlobDownloader { /// SDK BlobClient + Context pointers used by the production virtuals. struct ChunkContext; - /// Return the blob size in bytes. Production calls `BlobClient::GetProperties`. - virtual int64_t GetBlobSize(ChunkContext& ctx); + struct BlobProperties { + int64_t content_length = 0; + std::string blob_identity{}; + }; + + /// Return the blob size and identity from one properties read so chunk planning + /// and conditional range reads refer to the same remote blob version. + virtual BlobProperties GetBlobProperties(ChunkContext& ctx); /// Read `size` bytes starting at `offset` from the blob and forward them /// piecewise to `sink`. Pulls from the blob client referenced by `ctx`. diff --git a/sdk_v2/cpp/src/download/download_manager.cc b/sdk_v2/cpp/src/download/download_manager.cc index e0512a450..14b0a3377 100644 --- a/sdk_v2/cpp/src/download/download_manager.cc +++ b/sdk_v2/cpp/src/download/download_manager.cc @@ -28,6 +28,8 @@ namespace fl { namespace { const char* kDownloadSignalFileName = "download.tmp"; +const char* kDownloadSignalSidecarSafeMarker = "foundry-local-sidecar-safe-v1\n"; +const char* kDownloadSignalLegacyRepairMarker = "foundry-local-legacy-repair\n"; const char* kGenAIConfigFileName = "genai_config.json"; const char* kInferenceModelFileName = "inference_model.json"; const char* kDefaultRegistryRegion = "centralus"; @@ -45,7 +47,6 @@ bool HasInferenceModelJson(const std::string& model_path) { if (ec) { return false; } - for (const auto& entry : it) { if (entry.is_directory(ec)) { if (std::filesystem::exists(entry.path() / kInferenceModelFileName)) { @@ -57,6 +58,17 @@ bool HasInferenceModelJson(const std::string& model_path) { return false; } +bool HasSidecarSafeDownloadSignal(const std::filesystem::path& signal_path) { + std::ifstream signal(signal_path); + if (!signal.is_open()) { + return false; + } + + std::string marker; + std::getline(signal, marker); + return marker == "foundry-local-sidecar-safe-v1"; +} + /// Resolve the effective model path — the directory containing genai_config.json. /// For single-variant models this is model_path itself. /// For multi-variant models it's the first subdirectory containing genai_config.json. @@ -180,22 +192,30 @@ std::string ResolveRegion(const std::string& config_region, const ModelInfo& inf } // anonymous namespace DownloadManager::DownloadManager(std::string cache_directory, std::string_view catalog_region, int max_concurrency, - ILogger& logger, ITelemetry& telemetry, bool disable_region_fallback, - std::unique_ptr registry_client, - std::unique_ptr blob_downloader) + ILogger& logger, ITelemetry& telemetry) + : DownloadManager(std::move(cache_directory), catalog_region, max_concurrency, logger, false, telemetry) {} + +DownloadManager::DownloadManager(std::string cache_directory, std::string_view catalog_region, int max_concurrency, + ILogger& logger, bool disable_region_fallback, ITelemetry& telemetry) : cache_directory_(std::move(cache_directory)), config_region_(NormalizeConfiguredRegion(catalog_region)), max_concurrency_(max_concurrency), logger_(logger), - registry_client_(registry_client ? std::move(registry_client) - : std::make_unique( - kDefaultRegistryRegion, logger, - std::make_unique(logger, !disable_region_fallback))), - blob_downloader_(blob_downloader ? std::move(blob_downloader) : std::make_unique(logger)), + registry_client_(std::make_unique( + kDefaultRegistryRegion, logger, std::make_unique(logger, !disable_region_fallback))), + blob_downloader_(std::make_unique(logger)), telemetry_(telemetry) {} DownloadManager::~DownloadManager() = default; +void DownloadManager::SetModelRegistryClient(std::unique_ptr client) { + registry_client_ = std::move(client); +} + +void DownloadManager::SetBlobDownloader(std::unique_ptr downloader) { + blob_downloader_ = std::move(downloader); +} + std::string DownloadManager::ComputeModelPath(const ModelInfo& info) const { // Get publisher from string properties std::string publisher; @@ -218,7 +238,7 @@ std::string DownloadManager::ComputeModelPath(const ModelInfo& info) const { auto version = std::string_view(info.model_id).substr(last_colon + 1); SanitizeForPathSegment(bare_id); SanitizeForPathSegment(version); - sanitized_model_dir = FixVersionSuffix(info.model_id); + sanitized_model_dir = SanitizeForPathSegment(FixVersionSuffix(info.model_id)); } std::filesystem::path full_path(cache_directory_); @@ -254,25 +274,11 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, // are recorded as failures; the happy path explicitly sets kSuccess / kSkipped. DownloadTracker tracker(info.model_id, user_agent, telemetry_); tracker.SetMaxConcurrency(static_cast(max_concurrency_)); - auto download_start = clock::now(); auto record_lock_wait = [&]() { tracker.SetLockWaitMs(std::chrono::duration_cast( clock::now() - lock_wait_start) .count()); }; - auto record_download_elapsed = [&]() { - tracker.SetDownloadMs(std::chrono::duration_cast( - clock::now() - download_start) - .count()); - }; - auto record_stats = [&](const BlobDownloadStats& stats) { - tracker.SetEnumerationMs(stats.enumeration_ms); - tracker.SetDownloadMs(stats.download_ms); - tracker.SetFileCount(stats.file_count); - tracker.SetTotalSizeBytes(stats.total_size_bytes); - tracker.SetAlreadyCachedBytes(stats.already_cached_bytes); - tracker.SetSkippedFileCount(stats.skipped_file_count); - }; auto model_path = ComputeModelPath(info); @@ -280,6 +286,7 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, // A valid cache hit requires: directory exists, no in-progress signal file, and // inference_model.json is present (written by DownloadModel on successful completion). auto signal_path = std::filesystem::path(model_path) / kDownloadSignalFileName; + const bool model_dir_existed_before_download = std::filesystem::exists(model_path); if (std::filesystem::exists(model_path) && !std::filesystem::exists(signal_path) && HasInferenceModelJson(model_path)) { // Already cached and download was complete — cancellation request is @@ -314,7 +321,9 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, return progress_cb && progress_cb(0.0f) != 0; }; auto lock = CrossProcessFileLock::TryAcquireForDirectory(model_path, logger_); + bool waited_for_other_process = false; if (!lock) { + waited_for_other_process = true; logger_.Log(LogLevel::Information, "Model download is being performed by another process. Waiting on lock at '" + model_path + "'..."); @@ -335,14 +344,7 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, // (WaitForDirectoryLock) never runs while the mutex is held — the in-line acquire // at the top is the non-blocking TryAcquireForDirectory. Keep it that way. download_guard.unlock(); - try { - lock = CrossProcessFileLock::WaitForDirectoryLock(model_path, cancel_pred, logger_); - } catch (const std::exception& ex) { - record_lock_wait(); - tracker.SetDownloadWaitResult("Failed"); - tracker.RecordException(ex); - throw; - } + lock = CrossProcessFileLock::WaitForDirectoryLock(model_path, cancel_pred, logger_); download_guard.lock(); } record_lock_wait(); @@ -359,20 +361,24 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, return ResolveEffectiveModelPath(model_path); } + const bool had_download_signal = std::filesystem::exists(signal_path); + const bool can_skip_completed_files = + HasSidecarSafeDownloadSignal(signal_path) || + (!model_dir_existed_before_download && !waited_for_other_process && !had_download_signal); + // Create download signal file { std::ofstream signal(signal_path); - // Empty file — its presence indicates download is in progress + signal << (can_skip_completed_files ? kDownloadSignalSidecarSafeMarker : kDownloadSignalLegacyRepairMarker); } - BlobDownloadStats stats; - try { - // Emit 0% immediately so callers know the download process has started. - // This provides a heartbeat during the silent container resolution phase. - if (progress_cb && progress_cb(0.0f) != 0) { - FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "download cancelled by user progress callback"); - } + // Emit 0% immediately so callers know the download process has started. + // This provides a heartbeat during the silent container resolution phase. + if (progress_cb && progress_cb(0.0f) != 0) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "download cancelled by user progress callback"); + } + try { // Step 1: Resolve SAS URI from the region that served this model's catalog entry // (or the explicit override / default registry-region fallback). auto container = registry_client_->ResolveModelContainer(info.uri, ResolveRegion(config_region_, info)); @@ -387,6 +393,8 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, // and the blob container has all files at the root or in variant subdirectories. download_opts.path_prefix = ""; download_opts.max_concurrency = max_concurrency_; + download_opts.skip_completed_files = can_skip_completed_files; + download_opts.require_completed_file_identity = had_download_signal && !can_skip_completed_files; if (progress_cb) { download_opts.progress = [&progress_cb](float percent) { @@ -394,10 +402,16 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, }; } + BlobDownloadStats stats; DownloadBlobsToDirectory(*blob_downloader_, container.blob_sas_uri, model_path, download_opts, stats); - record_stats(stats); + tracker.SetEnumerationMs(stats.enumeration_ms); + tracker.SetDownloadMs(stats.download_ms); + tracker.SetFileCount(stats.file_count); + tracker.SetTotalSizeBytes(stats.total_size_bytes); + tracker.SetAlreadyCachedBytes(stats.already_cached_bytes); + tracker.SetSkippedFileCount(stats.skipped_file_count); tracker.SetDownloadWaitResult("Completed"); // Step 3: Write inference_model.json — use model_id (includes version) so the @@ -413,8 +427,6 @@ std::string DownloadManager::DownloadModel(const ModelInfo& info, tracker.SetStatus(ActionStatus::kSuccess); return ResolveEffectiveModelPath(model_path); } catch (const std::exception& e) { - record_stats(stats); - record_download_elapsed(); tracker.SetDownloadWaitResult("Failed"); tracker.RecordException(e); // Leave the signal file in place so the incomplete download is detected diff --git a/sdk_v2/cpp/src/download/download_manager.h b/sdk_v2/cpp/src/download/download_manager.h index 80605adef..b77ea7345 100644 --- a/sdk_v2/cpp/src/download/download_manager.h +++ b/sdk_v2/cpp/src/download/download_manager.h @@ -5,8 +5,8 @@ #include "download/blob_downloader.h" #include "download/model_registry_client.h" #include "model_info.h" -#include "util/region_fallback.h" +#include #include #include #include @@ -39,12 +39,21 @@ class DownloadManager { std::string_view catalog_region, int max_concurrency, ILogger& logger, - ITelemetry& telemetry, - bool disable_region_fallback = false, - std::unique_ptr registry_client = nullptr, - std::unique_ptr blob_downloader = nullptr); + ITelemetry& telemetry); + DownloadManager(std::string cache_directory, + std::string_view catalog_region, + int max_concurrency, + ILogger& logger, + bool disable_region_fallback, + ITelemetry& telemetry); ~DownloadManager(); + /// Override the model registry client (for testing). + void SetModelRegistryClient(std::unique_ptr client); + + /// Override the blob downloader (for testing). + void SetBlobDownloader(std::unique_ptr downloader); + /// Download a model to the local cache. /// progress_cb reports 0.0 to 100.0 percentage. /// user_agent identifies the calling API surface for telemetry attribution; pass an @@ -55,6 +64,20 @@ class DownloadManager { std::function progress_cb = nullptr, const std::string& user_agent = ""); + /// Preserve the existing `DownloadModel(info, nullptr)` call shape now that + /// a user-agent-first overload also accepts string-like second arguments. + std::string DownloadModel(const ModelInfo& info, std::nullptr_t progress_cb) { + return DownloadModel(info, std::function(progress_cb), ""); + } + + /// Download a model while passing telemetry user-agent attribution first. + /// This avoids forcing callers that only have a user-agent to pass a dummy progress callback. + std::string DownloadModel(const ModelInfo& info, + const std::string& user_agent, + std::function progress_cb = nullptr) { + return DownloadModel(info, progress_cb, user_agent); + } + /// Check if a model is cached locally (directory exists and download is complete). bool IsModelCached(const ModelInfo& info) const; diff --git a/sdk_v2/cpp/src/download/file_writer.cc b/sdk_v2/cpp/src/download/file_writer.cc index cdde85d24..eacf87488 100644 --- a/sdk_v2/cpp/src/download/file_writer.cc +++ b/sdk_v2/cpp/src/download/file_writer.cc @@ -17,6 +17,13 @@ namespace fs = std::filesystem; namespace { +void LogBestEffort(ILogger& logger, LogLevel level, std::string_view message) noexcept { + try { + logger.Log(level, message); + } catch (...) { + } +} + /// Ensure the data file exists at exactly `expected_size`, recreating it at the /// new size if it currently differs (larger or smaller). An existing file that /// is already the right size is left intact — the resume path relies on this. @@ -58,7 +65,15 @@ void EnsureFileExistsAtSize(const fs::path& path, int64_t expected_size) { FileWriter::FileWriter(ILogger& logger) : logger_(logger) {} -FileWriter::~FileWriter() { Close(); } +FileWriter::~FileWriter() { + try { + Close(); + } catch (const std::exception& ex) { + LogBestEffort(logger_, LogLevel::Warning, std::string("FileWriter: close failed in destructor: ") + ex.what()); + } catch (...) { + LogBestEffort(logger_, LogLevel::Warning, "FileWriter: close failed in destructor with unknown error"); + } +} void FileWriter::Open(const fs::path& path, int64_t expected_size) { EnsureFileExistsAtSize(path, expected_size); @@ -81,7 +96,8 @@ void FileWriter::Close() { if (handle_ != platform::kInvalidFileHandle) { std::string error; if (!platform::CloseFile(handle_, error)) { - logger_.Log(LogLevel::Warning, "FileWriter: " + error); + handle_ = platform::kInvalidFileHandle; + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, "FileWriter " + error); } handle_ = platform::kInvalidFileHandle; } diff --git a/sdk_v2/cpp/src/ep_detection/cuda_ep_bootstrapper.cc b/sdk_v2/cpp/src/ep_detection/cuda_ep_bootstrapper.cc index 7857655c0..b7ea216be 100644 --- a/sdk_v2/cpp/src/ep_detection/cuda_ep_bootstrapper.cc +++ b/sdk_v2/cpp/src/ep_detection/cuda_ep_bootstrapper.cc @@ -8,6 +8,7 @@ #include "utils.h" #include "http/http_download.h" #include "util/zip_extract.h" +#include "util/sha256.h" #include @@ -16,7 +17,10 @@ #include #include #include +#include #include +#include +#include namespace { @@ -28,6 +32,8 @@ constexpr int kMaxInstallAttempts = 5; // CUDA EP package is built against the ONNX Runtime version we link against. constexpr const char* kDownloadUrl = "https://foundrypackages-ffhrdhbxb7gpdreh.b02.azurefd.net/cuda-ep-20260501-062935.zip"; +constexpr const char* kPackageSha256 = + "D0DBCE52D121954F2D86E01E6D3418690A5BBA787028CCEE448E4BCDDD4F01E1"; struct ExpectedBinary { const char* filename; @@ -35,6 +41,17 @@ struct ExpectedBinary { }; constexpr ExpectedBinary kExpectedBinaries[] = { + {"cublas64_12.dll", "9513540E4EC4C51EE9E7304138C2CC255C29A8C181F9E80C38EFA25738BECD99"}, + {"cublasLt64_12.dll", "B199D1FF892A81B7FD3D57BA1781549609B41500B36008FEF326038393AD46C7"}, + {"cudart64_12.dll", "C2C9A9C22A9BCBA90E261825968836787B331038047A26770CFFB7A583C28344"}, + {"cudnn64_9.dll", "0D1D71325EB5E91570AB8BA8E399E07BF717FFD76511B2407229A8F45E0B1305"}, + {"cudnn_adv64_9.dll", "6D66BCE22502C2582A9C0E5398EE8CC38ADDCE2C837EB6DB8786ABC650E48DD8"}, + {"cudnn_engines_precompiled64_9.dll", "B410C3B42921AFC6E668FF994FCE1BF12C5A8A9B1A9445EBEE61958BF49B1E0A"}, + {"cudnn_engines_runtime_compiled64_9.dll", "8E62214495C96B93C6333C084FEC49B43F272B7E1977A12FE62275E9070647EB"}, + {"cudnn_graph64_9.dll", "82F710B01D15D20C311009721C771B76360A4954EBF7B5F4A407B0F96587F568"}, + {"cudnn_heuristic64_9.dll", "50719EEFB6692074096BF83C87E9CD186F7CE5B953201DA33669C1277A61949B"}, + {"cudnn_ops64_9.dll", "49487537744256A3D4365C4792B03BF31130AD1FAEA0A13EAFA219620941D837"}, + {"cufft64_11.dll", "F4FEA9227B14843894AD5436725F9638B172171142C95291FC6AE7A493248221"}, {"onnxruntime_providers_cuda.dll", "DD540FCFECFBC68B4675C9ADF09C2858CF6B054563859D79598AA2524406A76F"}, {"onnxruntime-genai-cuda.dll", "BC953F8E2AAFC6219B2D723B65AB8F1A9426A6B7724D6A01ED756FAE8C3DE6AE"}, }; @@ -43,6 +60,15 @@ constexpr const char* kRegistrationName = "Foundry.CUDA"; constexpr const char* kCudaProviderDll = "onnxruntime_providers_cuda.dll"; constexpr const char* kCudaProviderOverrideEnv = "FOUNDRY_LOCAL_CUDA_EP_LIBRARY"; +std::vector> ExpectedBinaryHashes() { + std::vector> expected; + expected.reserve(std::size(kExpectedBinaries)); + for (const auto& binary : kExpectedBinaries) { + expected.emplace_back(binary.filename, binary.sha256); + } + return expected; +} + } // anonymous namespace namespace fl { @@ -122,10 +148,7 @@ bool CudaEpBootstrapper::DownloadAndRegister(bool force, FileLock lock(lock_path); // Check if package already exists and is valid - if (fl::VerifyEpBinaries(ep_dir, - {{kExpectedBinaries[0].filename, kExpectedBinaries[0].sha256}, - {kExpectedBinaries[1].filename, kExpectedBinaries[1].sha256}}, - "CUDA EP", logger)) { + if (fl::VerifyEpBinaries(ep_dir, ExpectedBinaryHashes(), "CUDA EP", logger)) { logger.Log(LogLevel::Information, "CUDA EP: package already valid, skipping download"); } else { // Clean up any partial install @@ -156,6 +179,11 @@ bool CudaEpBootstrapper::DownloadAndRegister(bool force, return false; } + if (!VerifyEpArchive(zip_path, kPackageSha256, "CUDA EP", logger)) { + logger.Log(LogLevel::Warning, "CUDA EP: downloaded archive verification failed"); + return false; + } + // Extract logger.Log(LogLevel::Information, "CUDA EP: extracting..."); @@ -168,10 +196,7 @@ bool CudaEpBootstrapper::DownloadAndRegister(bool force, std::filesystem::remove(zip_path); // Verify - if (!fl::VerifyEpBinaries(ep_dir, - {{kExpectedBinaries[0].filename, kExpectedBinaries[0].sha256}, - {kExpectedBinaries[1].filename, kExpectedBinaries[1].sha256}}, - "CUDA EP", logger)) { + if (!fl::VerifyEpBinaries(ep_dir, ExpectedBinaryHashes(), "CUDA EP", logger)) { logger.Log(LogLevel::Warning, "CUDA EP: verification failed after download"); return false; } diff --git a/sdk_v2/cpp/src/ep_detection/ep_detector.cc b/sdk_v2/cpp/src/ep_detection/ep_detector.cc index 3b5273c9b..46f1f088f 100644 --- a/sdk_v2/cpp/src/ep_detection/ep_detector.cc +++ b/sdk_v2/cpp/src/ep_detection/ep_detector.cc @@ -24,13 +24,37 @@ EpDetector::EpDetector(const OrtApi& ort_api, OrtEnv& ort_env, bootstrappers_(std::move(bootstrappers)), logger_(logger), telemetry_(telemetry) { - // Populate the cache from bootstrappers_. Reads return snapshots so callers do - // not observe concurrent writes to is_registered. + // Populate the C++ cache once from bootstrappers_. Name string storage is stable + // because cached_eps_ size is fixed for the detector's lifetime. cached_eps_.reserve(bootstrappers_.size()); for (const auto& bs : bootstrappers_) { cached_eps_.push_back(EpInfo{bs->Name(), bs->IsRegistered()}); } + + PublishEpSnapshotLocked(); + PublishCApiSnapshotLocked(); +} + +void EpDetector::PublishEpSnapshotLocked() { + auto& snapshot = cached_eps_snapshots_.emplace_back(cached_eps_); + current_cached_eps_ = &snapshot; +} + +void EpDetector::PublishCApiSnapshotLocked() { + auto& snapshot = cached_eps_c_snapshots_.emplace_back(); + if (current_cached_eps_ == nullptr) { + current_cached_eps_ = &cached_eps_; + } + snapshot.reserve(current_cached_eps_->size()); + for (const auto& ep : *current_cached_eps_) { + snapshot.push_back(flEpInfo{ + FOUNDRY_LOCAL_API_VERSION, + ep.name.c_str(), + ep.is_registered, + }); + } + current_cached_eps_c_ = &snapshot; } std::map> EpDetector::GetAvailableDevicesToEPs() const { @@ -102,27 +126,20 @@ std::map> EpDetector::GetAvailableDevicesT } const std::vector& EpDetector::GetDiscoverableEps() const { - thread_local std::vector snapshot; std::lock_guard lock(cache_mutex_); - snapshot = cached_eps_; - return snapshot; + if (current_cached_eps_ == nullptr) { + static const std::vector empty; + return empty; + } + return *current_cached_eps_; } std::span EpDetector::GetDiscoverableEpsCApi() const { - thread_local std::vector snapshot_eps; - thread_local std::vector snapshot_c; std::lock_guard lock(cache_mutex_); - snapshot_eps = cached_eps_; - snapshot_c.clear(); - snapshot_c.reserve(snapshot_eps.size()); - for (const auto& ep : snapshot_eps) { - snapshot_c.push_back(flEpInfo{ - FOUNDRY_LOCAL_API_VERSION, - ep.name.c_str(), - ep.is_registered, - }); + if (current_cached_eps_c_ == nullptr) { + return {}; } - return snapshot_c; + return *current_cached_eps_c_; } EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vector* names, @@ -150,7 +167,6 @@ EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vector( std::chrono::steady_clock::now() - attempt_start) .count(); @@ -183,6 +199,7 @@ EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vector unmatched_requested_names = names ? *names : std::vector{}; if (progress_cb) { wrapped_cb = [&progress_cb, &cancelled](const std::string& ep_name, float percent) -> bool { @@ -212,15 +229,17 @@ EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vectorend()) { continue; } + unmatched_requested_names.erase(std::remove(unmatched_requested_names.begin(), unmatched_requested_names.end(), + bs->Name()), + unmatched_requested_names.end()); } logger_.Log(LogLevel::Information, "Downloading and registering EP: " + bs->Name()); // Per-provider EPDownloadAndRegister event via EpDownloadTracker. const bool was_registered_before = bs->IsRegistered(); - EpDownloadTracker tracker(bs->Name(), /*user_agent=*/std::string{}, telemetry_correlation_id, telemetry_); const auto initial_ready_state = was_registered_before ? EpReadyState::kRegistered : EpReadyState::kNotPresent; - const auto unresolved_ready_state = was_registered_before ? EpReadyState::kRegistered : EpReadyState::kUnknown; + EpDownloadTracker tracker(bs->Name(), /*user_agent=*/std::string{}, telemetry_correlation_id, telemetry_); tracker.RecordInitialState(initial_ready_state); ++telemetry_attempts; @@ -234,11 +253,6 @@ EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vectorName()); result.success = false; - telemetry_status = ActionStatusFromException(ex); - if (telemetry_status == ActionStatus::kCanceled) { - cancelled = true; - result.cancelled = true; - } result.status = "Some EPs failed to register"; tracker.RecordException(ex); record_attempt(); @@ -249,8 +263,8 @@ EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vector cache_lock(cache_mutex_); + const bool registration_changed = !cached_eps_[i].is_registered; cached_eps_[i].is_registered = true; + if (registration_changed) { + PublishEpSnapshotLocked(); + PublishCApiSnapshotLocked(); + } - tracker.RecordDownloadComplete(was_registered_before ? ActionStatus::kSkipped : ActionStatus::kSuccess, - unresolved_ready_state); + tracker.RecordDownloadComplete(ActionStatus::kSkipped, + was_registered_before ? EpReadyState::kRegistered : EpReadyState::kUnknown); tracker.RecordRegisterComplete(ActionStatus::kSuccess, EpReadyState::kRegistered); } else { ++telemetry_failed; result.failed_eps.push_back(bs->Name()); result.success = false; - telemetry_status = ActionStatus::kFailure; // The bootstrapper conflated download + register and returned false. // Record the combined operation as the register phase rather than fabricating a download/register split. - tracker.RecordDownloadComplete(ActionStatus::kSkipped, unresolved_ready_state); - tracker.RecordRegisterComplete(ActionStatus::kFailure, unresolved_ready_state); + tracker.RecordDownloadComplete(ActionStatus::kSkipped, + was_registered_before ? EpReadyState::kRegistered : EpReadyState::kUnknown); + tracker.RecordRegisterComplete(ActionStatus::kFailure, + was_registered_before ? EpReadyState::kRegistered : EpReadyState::kUnknown); } } @@ -280,15 +300,17 @@ EpDownloadResult EpDetector::DownloadAndRegisterEps(const std::vector(unmatched_requested_names.size()); + for (auto& name : unmatched_requested_names) { + result.failed_eps.push_back(std::move(name)); + } + result.status = "Some requested EPs are not available"; } else if (result.failed_eps.empty()) { result.status = "All requested EPs registered successfully"; - telemetry_status = ActionStatus::kSuccess; } else { result.status = "Some EPs failed to register"; - if (telemetry_status == ActionStatus::kSuccess) { - telemetry_status = ActionStatus::kFailure; - } } record_attempt(); diff --git a/sdk_v2/cpp/src/ep_detection/ep_detector.h b/sdk_v2/cpp/src/ep_detection/ep_detector.h index bdf5f6896..add864836 100644 --- a/sdk_v2/cpp/src/ep_detection/ep_detector.h +++ b/sdk_v2/cpp/src/ep_detection/ep_detector.h @@ -3,6 +3,7 @@ #pragma once #include +#include #include #include #include @@ -34,16 +35,17 @@ class IEpDetector { /// Example: { "CPU" → ["CPUExecutionProvider"], "GPU" → ["CudaExecutionProvider"] } virtual std::map> GetAvailableDevicesToEPs() const = 0; - /// Returns a recent snapshot of all discoverable EPs (those with bootstrappers). - /// The returned reference is valid until this method is next called on the same thread. + /// Returns metadata for all discoverable EPs (those with bootstrappers). /// Default: empty — no discoverable EPs beyond what's already registered. virtual const std::vector& GetDiscoverableEps() const { static const std::vector empty; return empty; } - /// C ABI view of a recent discoverable EPs snapshot. Pointers are valid until - /// this method is next called on the same thread. + /// C ABI view of the discoverable EPs. Pointers, struct addresses, and name + /// strings are stable for the detector's lifetime. is_registered values + /// reflect a recent snapshot — concurrent DownloadAndRegisterEps may have + /// updated them since the call returned. /// Default: empty span. virtual std::span GetDiscoverableEpsCApi() const { return {}; @@ -71,8 +73,9 @@ class EpDetector : public IEpDetector { /// @param ort_env The ORT environment singleton. /// @param bootstrappers EP bootstrappers for download/registration. /// @param logger Logger instance. - /// @param telemetry Telemetry sink. DownloadAndRegisterEps emits one EPDownloadAndRegister event per provider via - /// EpDownloadTracker, plus one aggregate EPDownloadAttempt event for the whole call. + /// @param telemetry Telemetry sink. DownloadAndRegisterEps emits one + /// EPDownloadAndRegister event per provider via EpDownloadTracker, + /// plus one aggregate EPDownloadAttempt event for the whole call. EpDetector(const OrtApi& ort_api, OrtEnv& ort_env, std::vector> bootstrappers, ILogger& logger, @@ -99,9 +102,17 @@ class EpDetector : public IEpDetector { std::mutex download_mutex_; std::atomic download_in_progress_{false}; mutable std::mutex cache_mutex_; - // Populated once in the constructor; mutated only under cache_mutex_ and copied - // into per-thread snapshots for callers. + // cached_eps_ is updated under cache_mutex_. Each C ABI view is an immutable + // snapshot retained for the detector lifetime so previously returned spans are + // never mutated or freed while callers may still read them. std::vector cached_eps_; + std::deque> cached_eps_snapshots_; + const std::vector* current_cached_eps_ = nullptr; + std::deque> cached_eps_c_snapshots_; + const std::vector* current_cached_eps_c_ = nullptr; + + void PublishEpSnapshotLocked(); + void PublishCApiSnapshotLocked(); }; } // namespace fl diff --git a/sdk_v2/cpp/src/ep_detection/ep_utils.cc b/sdk_v2/cpp/src/ep_detection/ep_utils.cc index 5cd025b1b..2f8c53fc3 100644 --- a/sdk_v2/cpp/src/ep_detection/ep_utils.cc +++ b/sdk_v2/cpp/src/ep_detection/ep_utils.cc @@ -9,6 +9,7 @@ #include #include +#include #ifdef _WIN32 #define WIN32_LEAN_AND_MEAN @@ -75,6 +76,30 @@ bool VerifyEpBinaries( return true; } +bool VerifyEpBinaries( + const std::filesystem::path& dir, + const std::vector>& expected, + std::string_view ep_name, + ILogger& logger) { + for (const auto& [filename, expected_hash] : expected) { + auto file_path = dir / filename; + + if (!std::filesystem::exists(file_path)) { + return false; + } + + auto hash = Sha256File(file_path); + if (CompareCaseInsensitive(hash, expected_hash) != 0) { + logger.Log(LogLevel::Warning, + fmt::format("{}: hash mismatch for {}: got {}, expected {}", + ep_name, filename, hash, expected_hash)); + return false; + } + } + + return true; +} + void PrependDirToProcessPath([[maybe_unused]] const std::filesystem::path& dir) { #ifdef _WIN32 DWORD len = GetEnvironmentVariableW(L"PATH", nullptr, 0); diff --git a/sdk_v2/cpp/src/ep_detection/ep_utils.h b/sdk_v2/cpp/src/ep_detection/ep_utils.h index ba4690e65..985248ab3 100644 --- a/sdk_v2/cpp/src/ep_detection/ep_utils.h +++ b/sdk_v2/cpp/src/ep_detection/ep_utils.h @@ -4,8 +4,10 @@ #include #include +#include #include #include +#include namespace fl { @@ -37,6 +39,12 @@ bool VerifyEpBinaries( std::string_view ep_name, ILogger& logger); +bool VerifyEpBinaries( + const std::filesystem::path& dir, + const std::vector>& expected, + std::string_view ep_name, + ILogger& logger); + /// Prepend @p dir to the process `PATH` environment variable for the lifetime of the process. /// /// EP provider libraries (CUDA, WebGPU) delay-load sibling dependency DLLs from their own directory, diff --git a/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.cc b/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.cc index d8b63db52..2f5d38aab 100644 --- a/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.cc +++ b/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.cc @@ -56,7 +56,7 @@ const std::unordered_map kPackageMetada "0767448ABEADE590821E88E56C471E0F2DFE89EC9A7642D08ADC3DC94F14AB92"} }, }; -#elif defined(__APPLE__) +#elif defined(__APPLE__) && defined(__aarch64__) const std::unordered_map kPackageMetadata = { {"macos-arm64", {"https://foundrypackages-ffhrdhbxb7gpdreh.b02.azurefd.net/webgpu_ep_0.1.0_macos-arm64.zip", @@ -64,7 +64,7 @@ const std::unordered_map kPackageMetada "A08BCEBE097B555E23938FCC71A5FAAD461F586CAB0B63DC9D21E970F6CA4C87"} }, }; -#else +#elif defined(__linux__) && defined(__x86_64__) && !defined(__ANDROID__) const std::unordered_map kPackageMetadata = { {"linux-x64", {"https://foundrypackages-ffhrdhbxb7gpdreh.b02.azurefd.net/webgpu_ep_0.1.0_linux-x64.zip", @@ -72,6 +72,8 @@ const std::unordered_map kPackageMetada "CBDFF74E6569E3CF66B46F0D194D87CD3CF49E83B7AA46552C39B0218D58B215"} }, }; +#else +const std::unordered_map kPackageMetadata = {}; #endif // Platform-specific package metadata is baked into the binary to keep @@ -80,10 +82,12 @@ const std::unordered_map kPackageMetada constexpr const char* kPlatformKey = "win-arm64"; #elif defined(_WIN32) constexpr const char* kPlatformKey = "win-x64"; -#elif defined(__APPLE__) +#elif defined(__APPLE__) && defined(__aarch64__) constexpr const char* kPlatformKey = "macos-arm64"; -#else +#elif defined(__linux__) && defined(__x86_64__) && !defined(__ANDROID__) constexpr const char* kPlatformKey = "linux-x64"; +#else +constexpr const char* kPlatformKey = ""; #endif const WebGpuPackageMetadata& GetPackageMetadata() { @@ -124,6 +128,10 @@ bool WebGpuEpBootstrapper::IsRegistered() const { return registered_; } +bool WebGpuEpBootstrapper::IsSupported() { + return kPackageMetadata.find(kPlatformKey) != kPackageMetadata.end(); +} + bool WebGpuEpBootstrapper::DownloadAndRegister(bool force, const ProgressCallback& progress_cb, ILogger& logger) { diff --git a/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.h b/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.h index e80d4651d..21a147d32 100644 --- a/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.h +++ b/sdk_v2/cpp/src/ep_detection/webgpu_ep_bootstrapper.h @@ -30,6 +30,7 @@ class WebGpuEpBootstrapper : public IEpBootstrapper { WebGpuEpBootstrapper& operator=(const WebGpuEpBootstrapper&) = delete; const std::string& Name() const override; + static bool IsSupported(); bool IsRegistered() const override; bool DownloadAndRegister(bool force, const ProgressCallback& progress_cb, diff --git a/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc b/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc index 889bf6ae3..321a5397a 100644 --- a/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.cc @@ -89,11 +89,13 @@ int64_t AudioDurationMsFromPcmBytes(int64_t bytes, int32_t sample_rate, int32_t } // namespace AudioSession::AudioSession(const fl::Model& catalog_model, GenAIModelInstance& model, - ILogger& logger, ITelemetry& telemetry) + ILogger& logger, ITelemetry& telemetry, bool session_ref_acquired) : Session(catalog_model, logger, telemetry), logger_(logger), model_(model) { logger_.Log(LogLevel::Debug, fmt::format("Creating AudioSession for model: {}", model.ModelId())); // Last so a throw above does not leak a refcount; nothing below can throw. - model_.AcquireSession(); + if (!session_ref_acquired) { + model_.AcquireSession(); + } } AudioSession::~AudioSession() { @@ -455,7 +457,7 @@ void AudioSession::ProcessAudioTranscriptionJson(const std::string& request_json // Validate file exists namespace fs = std::filesystem; if (!fs::exists(req.filename)) { - FL_LOG_AND_THROW(logger_, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Audio file not found"); + FL_LOG_AND_THROW(logger_, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, fmt::format("Audio file not found: '{}'", req.filename)); } // Build generation options from session defaults diff --git a/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.h b/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.h index e0378ba37..0a5f593a7 100644 --- a/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.h +++ b/sdk_v2/cpp/src/inferencing/generative/audio/audio_session.h @@ -36,7 +36,8 @@ struct SpeechSegmentItem; /// AudioTranscriptionResponse payload. class AudioSession : public Session { public: - AudioSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, ITelemetry& telemetry); + AudioSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, ITelemetry& telemetry, + bool session_ref_acquired = false); ~AudioSession(); // Movable: transfers session refcount ownership to the moved-to instance. diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc index f46ef5b71..7e5025bb1 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.cc @@ -49,11 +49,14 @@ void ApplyToolChoiceToContext(std::optional tool_choice, ToolCallC } // namespace -ChatSession::ChatSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, ITelemetry& telemetry) +ChatSession::ChatSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, + ITelemetry& telemetry, bool session_ref_acquired) : Session(catalog_model, logger, telemetry), logger_(logger), model_(model) { logger_.Log(LogLevel::Debug, fmt::format("Creating ChatSession for model: {}", model.ModelId())); // Last so a throw above does not leak a refcount; nothing below can throw. - model_.AcquireSession(); + if (!session_ref_acquired) { + model_.AcquireSession(); + } } ChatSession::~ChatSession() { diff --git a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h index b6c38f63f..b403aecca 100644 --- a/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h +++ b/sdk_v2/cpp/src/inferencing/generative/chat/chat_session.h @@ -39,7 +39,8 @@ class ChatSession : public Session { // The assistant reply is at history_[history_start + input_count] }; - ChatSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, ITelemetry& telemetry); + ChatSession(const fl::Model& catalog_model, GenAIModelInstance& model, ILogger& logger, ITelemetry& telemetry, + bool session_ref_acquired = false); ~ChatSession(); // Movable: transfers session refcount ownership to the moved-to instance. diff --git a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc index 131c83032..67f6ed63a 100644 --- a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc +++ b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.cc @@ -23,13 +23,15 @@ namespace fl { EmbeddingsSession::EmbeddingsSession(const fl::Model& catalog_model, GenAIModelInstance& model, - ILogger& logger, ITelemetry& telemetry) + ILogger& logger, ITelemetry& telemetry, bool session_ref_acquired) : Session(catalog_model, logger, telemetry, /*allow_concurrent_requests=*/true), logger_(logger), model_(model) { logger_.Log(LogLevel::Debug, fmt::format("Creating EmbeddingsSession for model: {}", model.ModelId())); // Last so a throw above does not leak a refcount; nothing below can throw. - model_.AcquireSession(); + if (!session_ref_acquired) { + model_.AcquireSession(); + } } EmbeddingsSession::~EmbeddingsSession() { @@ -78,7 +80,7 @@ void EmbeddingsSession::ProcessRequestImpl(const Request& request, Response& res } // Single batched forward pass for all inputs. - auto embeddings = GenerateEmbeddingsBatch(inputs, request); + auto embeddings = GenerateEmbeddingsBatch(inputs); // Wrap each embedding as a TensorItem in the response. The vector is heap-allocated // and ownership is transferred to the TensorItem deleter via deleter_user_data_, so @@ -100,7 +102,7 @@ void EmbeddingsSession::ProcessRequestImpl(const Request& request, Response& res } void EmbeddingsSession::ProcessEmbeddingsJson(const std::string& request_json, - const Request& original_request, + const Request& /*original_request*/, Response& response) { // Parse the OpenAI embeddings request. Let nlohmann::json::parse_error propagate — // matches AudioSession::ProcessAudioTranscriptionJson behavior. @@ -122,7 +124,7 @@ void EmbeddingsSession::ProcessEmbeddingsJson(const std::string& request_json, // Reuse GenerateEmbeddingsBatch so the typed and JSON paths stay bit-for-bit equal // under future refactors (parity test relies on this). if (!inputs.empty()) { - auto embeddings = GenerateEmbeddingsBatch(inputs, original_request); + auto embeddings = GenerateEmbeddingsBatch(inputs); output.data.reserve(embeddings.size()); for (size_t i = 0; i < embeddings.size(); ++i) { @@ -147,7 +149,7 @@ void EmbeddingsSession::ProcessEmbeddingsJson(const std::string& request_json, } std::vector> EmbeddingsSession::GenerateEmbeddingsBatch( - const std::vector& inputs, const Request& request) { + const std::vector& inputs) { // Process each input independently (batch_size=1 per forward pass). // // Embedding models like Qwen3-Embedding use bidirectional attention — @@ -162,10 +164,7 @@ std::vector> EmbeddingsSession::GenerateEmbeddingsBatch( results.reserve(inputs.size()); for (const auto& input : inputs) { - if (request.canceled.load(std::memory_order_relaxed)) { - FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "embeddings request canceled"); - } - results.push_back(GenerateSingleEmbedding(input, request)); + results.push_back(GenerateSingleEmbedding(input)); } logger_.Log(LogLevel::Verbose, @@ -174,11 +173,7 @@ std::vector> EmbeddingsSession::GenerateEmbeddingsBatch( return results; } -std::vector EmbeddingsSession::GenerateSingleEmbedding(const std::string& input, const Request& request) { - if (request.canceled.load(std::memory_order_relaxed)) { - FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "embeddings request canceled"); - } - +std::vector EmbeddingsSession::GenerateSingleEmbedding(const std::string& input) { auto& oga_model = model_.GetOgaModel(); auto& tokenizer = model_.GetOgaTokenizer(); @@ -201,13 +196,7 @@ std::vector EmbeddingsSession::GenerateSingleEmbedding(const std::string& generator->AppendTokenSequences(*sequences); // 3. Single forward pass. - if (request.canceled.load(std::memory_order_relaxed)) { - FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "embeddings request canceled"); - } generator->GenerateNextToken(); - if (request.canceled.load(std::memory_order_relaxed)) { - FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "embeddings request canceled"); - } // 4. Extract hidden_states. Shape: [1, token_count, hidden_size] auto hidden_states = generator->GetOutput("hidden_states"); diff --git a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h index 4dfe29a4a..b65f51bed 100644 --- a/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h +++ b/sdk_v2/cpp/src/inferencing/generative/embeddings/embeddings_session.h @@ -22,7 +22,7 @@ class GenAIModelInstance; class EmbeddingsSession : public Session { public: EmbeddingsSession(const fl::Model& catalog_model, GenAIModelInstance& model, - ILogger& logger, ITelemetry& telemetry); + ILogger& logger, ITelemetry& telemetry, bool session_ref_acquired = false); ~EmbeddingsSession(); EmbeddingsSession(EmbeddingsSession&&) = delete; @@ -37,11 +37,10 @@ class EmbeddingsSession : public Session { /// Generate L2-normalized embedding vectors for a list of inputs. /// Each input is processed independently (batch_size=1) to avoid /// padding artifacts with bidirectional-attention embedding models. - std::vector> GenerateEmbeddingsBatch(const std::vector& inputs, - const Request& request); + std::vector> GenerateEmbeddingsBatch(const std::vector& inputs); /// Generate a single L2-normalized embedding vector for one input string. - std::vector GenerateSingleEmbedding(const std::string& input, const Request& request); + std::vector GenerateSingleEmbedding(const std::string& input); /// Process a request whose first item is a TEXT item tagged OPENAI_JSON containing an /// OpenAI EmbeddingCreateRequest payload. Parses the JSON, runs generation via the diff --git a/sdk_v2/cpp/src/inferencing/model_load_manager.cc b/sdk_v2/cpp/src/inferencing/model_load_manager.cc index 0bc321b9b..f10e09e91 100644 --- a/sdk_v2/cpp/src/inferencing/model_load_manager.cc +++ b/sdk_v2/cpp/src/inferencing/model_load_manager.cc @@ -57,11 +57,52 @@ ModelLoadManager::ModelLoadManager(IEpDetector& ep_detector, ILogger& logger) : ep_detector_(ep_detector), logger_(logger) {} ModelLoadManager::~ModelLoadManager() { - // Destroy all loaded models under the lock. + // Destroy all loaded models under the lock. If a caller violates the public + // lifetime contract and keeps direct sessions alive past Manager destruction, + // keep those model instances alive rather than leaving sessions with dangling + // GenAIModelInstance references. std::lock_guard lock(mutex_); + for (auto it = loaded_models_.begin(); it != loaded_models_.end();) { + auto live_sessions = it->second->SessionRefCount(); + if (live_sessions > 0) { + logger_.Log(LogLevel::Warning, + fmt::format("ModelLoadManager destroyed while model '{}' still has {} live session(s); " + "leaving model instance allocated to avoid dangling session references", + it->first, live_sessions)); + it->second.release(); + it = loaded_models_.erase(it); + } else { + ++it; + } + } loaded_models_.clear(); } +ModelLoadManager::LoadedModelLease::~LoadedModelLease() { + Reset(); +} + +ModelLoadManager::LoadedModelLease::LoadedModelLease(LoadedModelLease&& other) noexcept + : model_(other.model_) { + other.model_ = nullptr; +} + +ModelLoadManager::LoadedModelLease& ModelLoadManager::LoadedModelLease::operator=(LoadedModelLease&& other) noexcept { + if (this != &other) { + Reset(); + model_ = other.model_; + other.model_ = nullptr; + } + return *this; +} + +void ModelLoadManager::LoadedModelLease::Reset() noexcept { + if (model_ != nullptr) { + model_->ReleaseSession(); + model_ = nullptr; + } +} + bool ModelLoadManager::HasEP(const std::string& ep_name) const { const auto& device_map = ep_detector_.GetAvailableDevicesToEPs(); for (const auto& [device, eps] : device_map) { @@ -89,6 +130,10 @@ ModelLoadManager::LoadResult ModelLoadManager::LoadModel(std::string_view model_ std::string path_str(model_path); std::string id_str(model_id); std::lock_guard lock(mutex_); + if (shutdown_.load()) { + FL_LOG_AND_THROW(logger_, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, + "cannot load model during shutdown"); + } // Check if model is already loaded auto it = loaded_models_.find(id_str); @@ -202,24 +247,21 @@ void ModelLoadManager::RejectNewLoads() { } void ModelLoadManager::UnloadAll(std::chrono::milliseconds timeout) { - // Snapshot ids+pointers under the lock; the GenAIModelInstance pointers stay valid - // because (a) only this method or UnloadModel can erase entries, and (b) we serialize - // the per-id erase below. - std::vector> snapshot; + std::vector ids; { std::lock_guard lock(mutex_); - snapshot.reserve(loaded_models_.size()); + ids.reserve(loaded_models_.size()); for (auto& [id, instance] : loaded_models_) { - snapshot.emplace_back(id, instance.get()); + ids.emplace_back(id); } } - if (snapshot.empty()) { + if (ids.empty()) { return; } logger_.Log(LogLevel::Information, - fmt::format("Shutdown: unloading {} model(s)", snapshot.size())); + fmt::format("Shutdown: unloading {} model(s)", ids.size())); using clock = std::chrono::steady_clock; constexpr auto kPollInterval = std::chrono::milliseconds(50); @@ -229,27 +271,39 @@ void ModelLoadManager::UnloadAll(std::chrono::milliseconds timeout) { // time linearly with model count. auto deadline = clock::now() + timeout; - for (auto& [id, instance] : snapshot) { - while (instance->SessionRefCount() > 0 && clock::now() < deadline) { + for (const auto& id : ids) { + int remaining = 0; + bool unloaded = false; + while (true) { + { + std::lock_guard lock(mutex_); + auto it = loaded_models_.find(id); + if (it == loaded_models_.end()) { + remaining = 0; + unloaded = true; + break; + } + + remaining = it->second->SessionRefCount(); + if (remaining == 0) { + logger_.Log(LogLevel::Information, fmt::format("unloading model: {}", id)); + loaded_models_.erase(it); + unloaded = true; + break; + } + } + + if (clock::now() >= deadline) { + break; + } + std::this_thread::sleep_for(kPollInterval); } - auto remaining = instance->SessionRefCount(); - if (remaining > 0) { + if (!unloaded && remaining > 0) { logger_.Log(LogLevel::Warning, fmt::format("Shutdown: model '{}' still has {} session(s) after overall {}ms deadline; leaving loaded", id, remaining, timeout.count())); - continue; - } - - try { - UnloadModel(id); - } catch (const std::exception& ex) { - // A new session attached between our refcount poll and the lock acquisition inside - // UnloadModel. Log and move on — IsShutdownRequested-gated callers shouldn't be - // creating new sessions, but we don't crash shutdown over it. - logger_.Log(LogLevel::Warning, - fmt::format("Shutdown: failed to unload '{}': {}", id, ex.what())); } } } @@ -270,4 +324,17 @@ GenAIModelInstance* ModelLoadManager::GetLoadedModel(std::string_view model_id) return nullptr; } +ModelLoadManager::LoadedModelLease ModelLoadManager::AcquireLoadedModel(std::string_view model_id) { + std::lock_guard lock(mutex_); + + std::string id_str(model_id); + auto it = loaded_models_.find(id_str); + if (it == loaded_models_.end()) { + return {}; + } + + it->second->AcquireSession(); + return LoadedModelLease(it->second.get()); +} + } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/model_load_manager.h b/sdk_v2/cpp/src/inferencing/model_load_manager.h index f582217ad..9082c3bfa 100644 --- a/sdk_v2/cpp/src/inferencing/model_load_manager.h +++ b/sdk_v2/cpp/src/inferencing/model_load_manager.h @@ -39,6 +39,33 @@ class ModelLoadManager { GenAIModelInstance* model = nullptr; // non-owning pointer; lifetime managed by this class }; + class LoadedModelLease { + public: + LoadedModelLease() = default; + ~LoadedModelLease(); + + LoadedModelLease(const LoadedModelLease&) = delete; + LoadedModelLease& operator=(const LoadedModelLease&) = delete; + + LoadedModelLease(LoadedModelLease&& other) noexcept; + LoadedModelLease& operator=(LoadedModelLease&& other) noexcept; + + explicit operator bool() const { return model_ != nullptr; } + GenAIModelInstance* get() const { return model_; } + GenAIModelInstance& operator*() const { return *model_; } + GenAIModelInstance* operator->() const { return model_; } + + void Reset() noexcept; + void Release() noexcept { model_ = nullptr; } + + private: + explicit LoadedModelLease(GenAIModelInstance* model) : model_(model) {} + + GenAIModelInstance* model_ = nullptr; + + friend class ModelLoadManager; + }; + ModelLoadManager(IEpDetector& ep_detector, ILogger& logger); ~ModelLoadManager(); @@ -67,6 +94,10 @@ class ModelLoadManager { /// The returned pointer is valid until UnloadModel is called for this model_id. GenAIModelInstance* GetLoadedModel(std::string_view model_id); + /// Get a loaded model by ID and increment its live-session refcount while + /// holding the load-manager lock. Returns an empty lease if not loaded. + LoadedModelLease AcquireLoadedModel(std::string_view model_id); + /// Reject all future LoadModel calls. Called by Manager::Shutdown(). /// Idempotent and thread-safe. void RejectNewLoads(); diff --git a/sdk_v2/cpp/src/inferencing/session/session.cc b/sdk_v2/cpp/src/inferencing/session/session.cc index da667506c..eb2612153 100644 --- a/sdk_v2/cpp/src/inferencing/session/session.cc +++ b/sdk_v2/cpp/src/inferencing/session/session.cc @@ -109,7 +109,7 @@ std::unique_ptr Session::Create(const fl::Model& model) { } auto& lm = mgr.GetModelLoadManager(); - auto* loaded = lm.GetLoadedModel(model.Id()); + auto loaded = lm.AcquireLoadedModel(model.Id()); if (!loaded) { FL_LOG_AND_THROW(logger, FOUNDRY_LOCAL_ERROR_INTERNAL, "loaded model not found in load manager"); } @@ -118,19 +118,22 @@ std::unique_ptr Session::Create(const fl::Model& model) { const auto& info = model.Info(); if (info.task == "chat-completion" || info.task == "vision-language-chat") { - auto session = std::make_unique(model, *loaded, logger, telemetry); + auto session = std::make_unique(model, *loaded, logger, telemetry, true); + loaded.Release(); tracker.SetStatus(ActionStatus::kSuccess); return session; } if (info.task == "automatic-speech-recognition") { - auto session = std::make_unique(model, *loaded, logger, telemetry); + auto session = std::make_unique(model, *loaded, logger, telemetry, true); + loaded.Release(); tracker.SetStatus(ActionStatus::kSuccess); return session; } if (info.task == "embeddings") { - auto session = std::make_unique(model, *loaded, logger, telemetry); + auto session = std::make_unique(model, *loaded, logger, telemetry, true); + loaded.Release(); tracker.SetStatus(ActionStatus::kSuccess); return session; } @@ -162,22 +165,6 @@ void Session::ProcessRequest(const Request& request, Response& response) { lock.lock(); } - struct ActiveRequestGuard { - Session& session; - const Request& request; - ActiveRequestGuard(Session& session, const Request& request) : session(session), request(request) { - std::lock_guard lock(session.active_requests_mutex_); - session.active_requests_.insert(&request); - if (session.session_canceled_) { - request.canceled.store(true, std::memory_order_relaxed); - } - } - ~ActiveRequestGuard() { - std::lock_guard lock(session.active_requests_mutex_); - session.active_requests_.erase(&request); - } - } active_request_guard(*this, request); - // Use the context the caller staged (an HTTP route stages an indirect child // with the route's correlation id); otherwise mint a direct context per call // for direct SDK use. @@ -196,9 +183,11 @@ void Session::ProcessRequest(const Request& request, Response& response) { const auto start = std::chrono::steady_clock::now(); try { ProcessRequestImpl(request, response); + if (request.canceled.load(std::memory_order_relaxed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "request cancelled"); + } - tracker.SetStatus(request.canceled.load(std::memory_order_relaxed) ? ActionStatus::kCanceled - : ActionStatus::kSuccess); + tracker.SetStatus(ActionStatus::kSuccess); } catch (const std::exception& ex) { tracker.RecordException(ex); throw; @@ -231,12 +220,4 @@ void Session::ProcessRequest(const Request& request, Response& response) { } } -void Session::Cancel() { - std::lock_guard lock(active_requests_mutex_); - session_canceled_ = true; - for (const Request* request : active_requests_) { - request->canceled.store(true, std::memory_order_relaxed); - } -} - } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/session.h b/sdk_v2/cpp/src/inferencing/session/session.h index 5358e8b73..11fcf5481 100644 --- a/sdk_v2/cpp/src/inferencing/session/session.h +++ b/sdk_v2/cpp/src/inferencing/session/session.h @@ -8,7 +8,6 @@ #include #include #include -#include #include #include @@ -55,9 +54,6 @@ class Session { /// in-flight callbacks and ensures the Response is fully populated on return. void ProcessRequest(const Request& request, Response& response); - /// Signal all active requests in this session to stop. - void Cancel(); - /// Add a tool definition to this session. /// @throws fl::Exception if tool_def.json_schema is not valid JSON. void AddToolDefinition(ToolDefinition tool_def); @@ -182,9 +178,6 @@ class Session { std::optional request_context_; const bool allow_concurrent_requests_; mutable std::unique_ptr request_mutex_ = std::make_unique(); - mutable std::mutex active_requests_mutex_; - bool session_canceled_ = false; - std::unordered_set active_requests_; }; } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/session_manager.cc b/sdk_v2/cpp/src/inferencing/session/session_manager.cc index e1d7e368e..59f1a86d7 100644 --- a/sdk_v2/cpp/src/inferencing/session/session_manager.cc +++ b/sdk_v2/cpp/src/inferencing/session/session_manager.cc @@ -22,19 +22,19 @@ SessionManager::~SessionManager() { } void SessionManager::Register(Session& session) { - std::lock_guard lock(mutex_); - if (shutting_down_) { + if (shutting_down_.load()) { FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "cannot create session during shutdown"); } - ++sessions_[&session]; + std::lock_guard lock(mutex_); + sessions_.insert(&session); } void SessionManager::Deregister(Session& session) { std::lock_guard lock(mutex_); - auto it = sessions_.find(&session); + auto erased = sessions_.erase(&session); - if (it == sessions_.end()) { + if (erased == 0) { // Bug: session was not registered. Log loudly but don't throw — this may be // called from a destructor where throwing would call std::terminate(). logger_.Log(LogLevel::Error, "SessionManager::Deregister called for unregistered session"); @@ -42,29 +42,22 @@ void SessionManager::Deregister(Session& session) { return; } - if (--it->second == 0) { - sessions_.erase(it); - } - if (sessions_.empty()) { drain_cv_.notify_all(); } } void SessionManager::CancelAll() { - std::vector> cached_sessions; + shutting_down_.store(true); - { - std::lock_guard lock(mutex_); - shutting_down_ = true; - cached_sessions = MoveCacheToDestroyLocked(); - logger_.Log(LogLevel::Information, - fmt::format("SessionManager: cancelling all sessions ({} active)", sessions_.size())); + // Clear cache — frees idle cached sessions so they don't block drain. + ClearCache(); - for (const auto& session_entry : sessions_) { - session_entry.first->Cancel(); - } - } + std::lock_guard lock(mutex_); + logger_.Log(LogLevel::Information, + fmt::format("SessionManager: cancelling all sessions ({} active)", sessions_.size())); + + // Future (Phase 3): iterate sessions_ and call Cancel() on each } void SessionManager::WaitForDrain(std::chrono::milliseconds timeout) { @@ -92,44 +85,61 @@ size_t SessionManager::ActiveCount() const { // --- Session cache --- -std::unique_ptr SessionManager::CheckOut(const std::string& key) { +std::unique_ptr SessionManager::CheckOut(const std::string& key, const std::string& expected_model_id) { std::lock_guard lock(mutex_); auto it = cache_.find(key); if (it == cache_.end()) { return nullptr; } + if (!expected_model_id.empty() && it->second.model_id != expected_model_id) { + logger_.Log(LogLevel::Information, + fmt::format("SessionManager: cached session for '{}' uses model '{}', not requested model '{}'", + key, it->second.model_id, expected_model_id)); + return nullptr; + } - auto session = RemoveCachedLocked(it); + auto session = std::move(it->second.session); + lru_order_.erase(it->second.lru_iter); + cache_.erase(it); logger_.Log(LogLevel::Debug, fmt::format("SessionManager: checked out cached session for '{}'", key)); return session; } -void SessionManager::CheckIn(const std::string& key, std::unique_ptr session) { +void SessionManager::CheckIn(const std::string& key, std::unique_ptr session, std::string model_id) { // Collect evicted sessions to destroy outside the lock. std::vector> evicted; { std::lock_guard lock(mutex_); - if (shutting_down_ || cache_capacity_ == 0) { - evicted.push_back(std::move(session)); + + if (shutting_down_.load()) { + logger_.Log(LogLevel::Debug, + fmt::format("SessionManager: dropping check-in for '{}' during shutdown", key)); return; } + // Replace existing entry for this key (if any) auto existing = cache_.find(key); if (existing != cache_.end()) { - evicted.push_back(RemoveCachedLocked(existing)); + evicted.push_back(std::move(existing->second.session)); + lru_order_.erase(existing->second.lru_iter); + cache_.erase(existing); } + // Evict LRU if at capacity while (cache_.size() >= cache_capacity_) { const auto& lru_key = lru_order_.back(); auto lru_it = cache_.find(lru_key); - evicted.push_back(RemoveCachedLocked(lru_it)); + evicted.push_back(std::move(lru_it->second.session)); + cache_.erase(lru_it); + lru_order_.pop_back(); } + // Insert new entry lru_order_.push_front(key); - cache_[key] = CacheEntry{std::move(session), lru_order_.begin()}; + cache_[key] = CacheEntry{std::move(session), std::move(model_id), lru_order_.begin()}; logger_.Log(LogLevel::Debug, fmt::format("SessionManager: checked in session under '{}' (cache size: {})", key, cache_.size())); @@ -159,39 +169,30 @@ bool SessionManager::EvictCached(const std::string& key) { return false; } - evicted = RemoveCachedLocked(it); + evicted = std::move(it->second.session); + lru_order_.erase(it->second.lru_iter); + cache_.erase(it); } logger_.Log(LogLevel::Debug, fmt::format("SessionManager: evicted cached session for '{}'", key)); return true; } -std::unique_ptr SessionManager::RemoveCachedLocked(CacheMap::iterator it) { - auto session = std::move(it->second.session); - lru_order_.erase(it->second.lru_iter); - cache_.erase(it); - return session; -} - -std::vector> SessionManager::MoveCacheToDestroyLocked() { - std::vector> to_destroy; - - for (auto& cache_entry : cache_) { - to_destroy.push_back(std::move(cache_entry.second.session)); - } - - cache_.clear(); - lru_order_.clear(); - return to_destroy; -} - void SessionManager::ClearCache() { std::vector> to_destroy; { std::lock_guard lock(mutex_); - to_destroy = MoveCacheToDestroyLocked(); + + for (auto& [key, entry] : cache_) { + to_destroy.push_back(std::move(entry.session)); + } + + cache_.clear(); + lru_order_.clear(); } + + // Destroy outside lock } } // namespace fl diff --git a/sdk_v2/cpp/src/inferencing/session/session_manager.h b/sdk_v2/cpp/src/inferencing/session/session_manager.h index c0384bed7..266315314 100644 --- a/sdk_v2/cpp/src/inferencing/session/session_manager.h +++ b/sdk_v2/cpp/src/inferencing/session/session_manager.h @@ -4,6 +4,7 @@ #include "logger.h" +#include #include #include #include @@ -11,7 +12,7 @@ #include #include #include -#include +#include namespace fl { @@ -35,12 +36,14 @@ class ISessionManager { /// /// Cache semantics (check-out / check-in): /// - CheckOut(key) removes the session from the cache and transfers ownership to the caller. +/// If expected_model_id is non-empty and the cached session was created for +/// another model, returns nullptr and leaves the cached session untouched. /// While checked out, the session cannot be evicted or accessed by concurrent requests. /// - CheckIn(key, session) inserts the session into the cache under the given key. /// May evict the least-recently-used entry if at capacity. /// -/// Cached sessions are NOT registered — they are idle. Only sessions actively processing requests should be registered -/// via SessionRegistration. +/// Cached sessions are NOT registered — they are idle. Only sessions actively +/// processing requests should be registered via SessionRegistration. class SessionManager : public ISessionManager { public: static constexpr size_t kDefaultCacheCapacity = 5; @@ -66,18 +69,19 @@ class SessionManager : public ISessionManager { /// Block until all sessions have deregistered, with timeout. void WaitForDrain(std::chrono::milliseconds timeout = std::chrono::seconds(10)); - /// Number of currently active sessions. + /// Number of currently tracked sessions (active + cached). size_t ActiveCount() const; // --- Session cache (Responses API) --- /// Remove a session from the cache by key and return it. - /// Returns nullptr on cache miss. - std::unique_ptr CheckOut(const std::string& key); + /// Returns nullptr on cache miss. The session is still tracked (registered). + std::unique_ptr CheckOut(const std::string& key, const std::string& expected_model_id = {}); /// Insert a session into the cache under the given key. /// May evict the least-recently-used entry if at capacity. - void CheckIn(const std::string& key, std::unique_ptr session); + /// The session remains tracked (registered) while cached. + void CheckIn(const std::string& key, std::unique_ptr session, std::string model_id = {}); /// Remove and destroy a cached session by key. Returns true if an entry was removed. /// @@ -92,24 +96,24 @@ class SessionManager : public ISessionManager { private: struct CacheEntry { std::unique_ptr session; + std::string model_id; std::list::iterator lru_iter; }; - using CacheMap = std::unordered_map; - std::unique_ptr RemoveCachedLocked(CacheMap::iterator it); - std::vector> MoveCacheToDestroyLocked(); + /// Clear all cache entries. Acquires mutex_ to move entries out, then destroys + /// them outside the lock (destroyed sessions call Deregister, which acquires mutex_). void ClearCache(); ILogger& logger_; size_t cache_capacity_; - bool shutting_down_ = false; + std::atomic shutting_down_{false}; mutable std::mutex mutex_; std::condition_variable drain_cv_; - std::unordered_map sessions_; + std::unordered_set sessions_; // LRU cache: front of lru_order_ = most recently used std::list lru_order_; - CacheMap cache_; + std::unordered_map cache_; }; } // namespace fl diff --git a/sdk_v2/cpp/src/manager.cc b/sdk_v2/cpp/src/manager.cc index 55b0502e1..465c1cf00 100644 --- a/sdk_v2/cpp/src/manager.cc +++ b/sdk_v2/cpp/src/manager.cc @@ -442,8 +442,8 @@ Manager::Manager(const Configuration& config) config_.catalog_region.value_or("auto"), download_concurrency, *logger_, - *telemetry_, - disable_region_fallback); + disable_region_fallback, + *telemetry_); model_load_manager_ = std::make_unique(*ep_detector_, *logger_); session_manager_ = std::make_unique(*logger_); catalog_ = std::make_unique( diff --git a/sdk_v2/cpp/src/model.cc b/sdk_v2/cpp/src/model.cc index 1f06201ce..e036ed468 100644 --- a/sdk_v2/cpp/src/model.cc +++ b/sdk_v2/cpp/src/model.cc @@ -111,11 +111,13 @@ Model::Model(Model&& other) noexcept download_manager_(other.download_manager_), model_load_manager_(other.model_load_manager_), variants_(std::move(other.variants_)), - selected_variant_(other.selected_variant_) { + selected_variant_(other.selected_variant_.load(std::memory_order_acquire)), + explicit_variant_selected_(other.explicit_variant_selected_.load(std::memory_order_acquire)) { // After vector move, selected_variant_ still points into the transferred buffer. other.download_manager_ = nullptr; other.model_load_manager_ = nullptr; - other.selected_variant_ = nullptr; + other.selected_variant_.store(nullptr, std::memory_order_release); + other.explicit_variant_selected_.store(false, std::memory_order_release); } Model& Model::operator=(Model&& other) noexcept { @@ -126,10 +128,13 @@ Model& Model::operator=(Model&& other) noexcept { download_manager_ = other.download_manager_; model_load_manager_ = other.model_load_manager_; variants_ = std::move(other.variants_); - selected_variant_ = other.selected_variant_; + selected_variant_.store(other.selected_variant_.load(std::memory_order_acquire), std::memory_order_release); + explicit_variant_selected_.store(other.explicit_variant_selected_.load(std::memory_order_acquire), + std::memory_order_release); other.download_manager_ = nullptr; other.model_load_manager_ = nullptr; - other.selected_variant_ = nullptr; + other.selected_variant_.store(nullptr, std::memory_order_release); + other.explicit_variant_selected_.store(false, std::memory_order_release); } return *this; @@ -163,7 +168,7 @@ Model Model::FromModelInfo(ModelInfo info, Model Model::MakeContainer(Model first_variant) { Model container; container.variants_.push_back(std::make_unique(std::move(first_variant))); - container.selected_variant_ = container.variants_.back().get(); + container.selected_variant_.store(container.variants_.back().get(), std::memory_order_release); return container; } @@ -196,12 +201,14 @@ void Model::SelectDefaultVariant() { for (auto& v : variants_) { if (v->IsCached()) { - selected_variant_ = v.get(); + selected_variant_.store(v.get(), std::memory_order_release); + explicit_variant_selected_.store(false, std::memory_order_release); return; } } - selected_variant_ = variants_.front().get(); + selected_variant_.store(variants_.front().get(), std::memory_order_release); + explicit_variant_selected_.store(false, std::memory_order_release); } // --------------------------------------------------------------------------- @@ -209,24 +216,24 @@ void Model::SelectDefaultVariant() { // --------------------------------------------------------------------------- const std::string& Model::Id() const { - if (selected_variant_) { - return selected_variant_->Id(); + if (auto* selected = SelectedVariant()) { + return selected->Id(); } return info_.model_id; } const std::string& Model::Alias() const { - if (selected_variant_) { - return selected_variant_->Alias(); + if (auto* selected = SelectedVariant()) { + return selected->Alias(); } return info_.alias; } const ModelInfo& Model::Info() const { - if (selected_variant_) { - return selected_variant_->Info(); + if (auto* selected = SelectedVariant()) { + return selected->Info(); } return info_; @@ -251,16 +258,16 @@ std::vector Model::Variants() const { } bool Model::IsCached() const { - if (selected_variant_) { - return selected_variant_->IsCached(); + if (auto* selected = SelectedVariant()) { + return selected->IsCached(); } return cached_; } bool Model::IsLoaded() const { - if (selected_variant_) { - return selected_variant_->IsLoaded(); + if (auto* selected = SelectedVariant()) { + return selected->IsLoaded(); } // ModelLoadManager owns the authoritative loaded-instance map. The pointer is set at @@ -274,8 +281,8 @@ bool Model::IsLoaded() const { // --------------------------------------------------------------------------- void Model::Download(std::function progress_cb) { - if (selected_variant_) { - selected_variant_->Download(std::move(progress_cb)); + if (auto* selected = SelectedVariant()) { + selected->Download(std::move(progress_cb)); return; } @@ -296,16 +303,16 @@ void Model::Download(std::function progress_cb) { } const std::string& Model::GetPath() const { - if (selected_variant_) { - return selected_variant_->GetPath(); + if (auto* selected = SelectedVariant()) { + return selected->GetPath(); } return local_path_; } void Model::Load(ExecutionProvider ep) { - if (selected_variant_) { - selected_variant_->Load(ep); + if (auto* selected = SelectedVariant()) { + selected->Load(ep); return; } @@ -319,8 +326,8 @@ void Model::Load(ExecutionProvider ep) { } void Model::Unload() { - if (selected_variant_) { - selected_variant_->Unload(); + if (auto* selected = SelectedVariant()) { + selected->Unload(); return; } @@ -329,8 +336,8 @@ void Model::Unload() { } void Model::RemoveFromCache() { - if (selected_variant_) { - selected_variant_->RemoveFromCache(); + if (auto* selected = SelectedVariant()) { + selected->RemoveFromCache(); return; } @@ -357,9 +364,12 @@ void Model::SelectVariant(const Model& variant) { "with all variants available."); } + std::lock_guard lock(state_mutex_); + for (auto& v : variants_) { if (v.get() == &variant) { - selected_variant_ = v.get(); + selected_variant_.store(v.get(), std::memory_order_release); + explicit_variant_selected_.store(true, std::memory_order_release); return; } } @@ -434,8 +444,8 @@ Model::IOInfo IOInfoFromCache(const StaticIOCache& cache) { } // namespace Model::IOInfo Model::GetInputOutputInfo() const { - if (selected_variant_) { - return selected_variant_->GetInputOutputInfo(); + if (auto* selected = SelectedVariant()) { + return selected->GetInputOutputInfo(); } const auto& task = Info().task; diff --git a/sdk_v2/cpp/src/model.h b/sdk_v2/cpp/src/model.h index 90710768a..0dcb329c1 100644 --- a/sdk_v2/cpp/src/model.h +++ b/sdk_v2/cpp/src/model.h @@ -70,6 +70,9 @@ class Model { /// first cached variant if any, else the best variant. /// Requires IsContainer() to be true. void SelectDefaultVariant(); + bool HasExplicitVariantSelection() const { + return explicit_variant_selected_.load(std::memory_order_acquire); + } /// Best-first comparator: device priority asc, version desc, created-at desc, model_id asc. /// Exposed so callers that return Model* lists can produce consistent ordering with Variants(). @@ -116,7 +119,7 @@ class Model { IOInfo GetInputOutputInfo() const; /// True if this is a multi-variant container. - bool IsContainer() const { return selected_variant_ != nullptr; } + bool IsContainer() const { return SelectedVariant() != nullptr; } // --- Mutation methods --- @@ -169,11 +172,14 @@ class Model { // Container data (empty/null for leaves). unique_ptr keeps Model addresses // stable across vector growth/reordering. std::vector> variants_; - Model* selected_variant_ = nullptr; // non-null = this is a container + std::atomic selected_variant_{nullptr}; // non-null = this is a container + std::atomic explicit_variant_selected_{false}; // Guards variants_ across reader/writer threads (catalog refresh adding variants // while another thread enumerates via Variants()). mutable std::mutex state_mutex_; + + Model* SelectedVariant() const { return selected_variant_.load(std::memory_order_acquire); } }; } // namespace fl diff --git a/sdk_v2/cpp/src/service/audio_transcriptions_handler.cc b/sdk_v2/cpp/src/service/audio_transcriptions_handler.cc index a084a016f..17a29850c 100644 --- a/sdk_v2/cpp/src/service/audio_transcriptions_handler.cc +++ b/sdk_v2/cpp/src/service/audio_transcriptions_handler.cc @@ -7,6 +7,7 @@ #include "c_api_types.h" #include "catalog.h" #include "contracts/audio_transcriptions.h" +#include "exception.h" #include "inferencing/generative/audio/audio_session.h" #include "inferencing/model_load_manager.h" #include "inferencing/session/session_manager.h" @@ -17,6 +18,7 @@ #include #include +#include #include namespace fl { @@ -57,13 +59,13 @@ std::shared_ptr AudioTranscriptionsHandler } std::shared_ptr AudioTranscriptionsHandler::ResolveModel( - const std::string& model_name, Model*& model, GenAIModelInstance*& loaded) { + const std::string& model_name, Model*& model, ModelLoadManager::LoadedModelLease& loaded) { model = ctx_.catalog.GetModelVariant(model_name); if (!model) { return ErrorResponse(Status::CODE_404, "Model not found", "No model matching '" + model_name + "'"); } - loaded = ctx_.model_load_manager.GetLoadedModel(model->Id()); + loaded = ctx_.model_load_manager.AcquireLoadedModel(model->Id()); if (!loaded) { return ErrorResponse(Status::CODE_400, "Model not loaded", "Model '" + model_name + "' must be loaded before inference"); @@ -122,7 +124,7 @@ std::shared_ptr AudioTranscriptionsHandler // 3. Resolve model std::string model_name = req.model; Model* model = nullptr; - GenAIModelInstance* loaded = nullptr; + ModelLoadManager::LoadedModelLease loaded; if (auto err = ResolveModel(model_name, model, loaded)) { tracker->SetStatus(ActionStatus::kClientError); return err; @@ -139,22 +141,22 @@ std::shared_ptr AudioTranscriptionsHandler // 5. Dispatch to streaming or non-streaming try { - std::unique_ptr session; + std::optional session; { ActionTracker create_tracker(Action::kSessionCreate, ctx_.telemetry, session_ctx); create_tracker.SetModelId(model_name); - session = std::make_unique(*model, *loaded, ctx_.logger, ctx_.telemetry); + session.emplace(*model, *loaded, ctx_.logger, ctx_.telemetry, true); + loaded.Release(); create_tracker.SetStatus(ActionStatus::kSuccess); } - AudioSession& session_ref = *session; - session_ref.SetRequestContext(session_ctx); + session->SetRequestContext(session_ctx); if (stream) { // The route action is recorded by the streaming thread on completion. return HandleStreaming(std::move(*session), std::move(session_request), std::move(tracker)); } else { - SessionRegistration reg(ctx_.session_manager, session_ref); - auto response = HandleNonStreaming(session_ref, session_request); + SessionRegistration reg(ctx_.session_manager, *session); + auto response = HandleNonStreaming(*session, session_request); tracker->SetStatus(ResponseToActionStatus(response)); return response; } @@ -201,11 +203,14 @@ std::shared_ptr AudioTranscriptionsHandler auto body_ptr = body; auto& logger = ctx_.logger; auto& thread_tracker = ctx_.thread_tracker; + auto stream_done = std::make_shared>(false); + auto stream_request = std::make_shared(std::move(session_request)); std::thread streaming_thread([bg_session = std::move(session), body_ptr, &logger, - req = std::move(session_request), &thread_tracker, + req = stream_request, &thread_tracker, route_tracker = std::move(route_tracker), - &session_manager = ctx_.session_manager]() mutable { + &session_manager = ctx_.session_manager, + stream_done]() mutable { try { SessionRegistration reg(session_manager, bg_session); fl::Response bg_response; @@ -221,7 +226,6 @@ std::shared_ptr AudioTranscriptionsHandler if (item->type == FOUNDRY_LOCAL_ITEM_TEXT) { auto& text_item = static_cast(*item); if (!body_ptr->Push("data: " + text_item.text + "\n\n")) { - req.canceled.store(true, std::memory_order_relaxed); return 1; } } else { @@ -234,13 +238,18 @@ std::shared_ptr AudioTranscriptionsHandler }; bg_session.SetStreamingCallback(callback_fn); - bg_session.ProcessRequest(req, bg_response); + bg_session.ProcessRequest(*req, bg_response); + if (req->canceled.load(std::memory_order_relaxed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "audio transcription stream cancelled"); + } - body_ptr->Push("data: [DONE]\n\n"); + // Send terminal event + if (!body_ptr->Push("data: [DONE]\n\n")) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "audio transcription stream cancelled"); + } if (route_tracker) { - route_tracker->SetStatus(req.canceled.load(std::memory_order_relaxed) ? ActionStatus::kCanceled - : ActionStatus::kSuccess); + route_tracker->SetStatus(ActionStatus::kSuccess); } } catch (const std::exception& ex) { logger.Log(LogLevel::Error, fmt::format("Audio streaming transcription failed: {}", ex.what())); @@ -255,11 +264,14 @@ std::shared_ptr AudioTranscriptionsHandler } body_ptr->Finish(); - route_tracker.reset(); - thread_tracker.Remove(std::this_thread::get_id()); + stream_done->store(true, std::memory_order_release); + thread_tracker.NotifyCompleted(); }); - thread_tracker.Track(std::move(streaming_thread)); + thread_tracker.Track(std::move(streaming_thread), stream_done, [body, stream_request] { + body->Abort(); + stream_request->canceled.store(true, std::memory_order_relaxed); + }); auto response = oatpp::web::protocol::http::outgoing::Response::createShared(Status::CODE_200, body); response->putHeader("Content-Type", "text/event-stream"); diff --git a/sdk_v2/cpp/src/service/audio_transcriptions_handler.h b/sdk_v2/cpp/src/service/audio_transcriptions_handler.h index 1f68c058c..9f25345fa 100644 --- a/sdk_v2/cpp/src/service/audio_transcriptions_handler.h +++ b/sdk_v2/cpp/src/service/audio_transcriptions_handler.h @@ -4,6 +4,7 @@ #ifdef FOUNDRY_LOCAL_HAS_WEB_SERVICE +#include "inferencing/model_load_manager.h" #include "service/handler_utils.h" #include @@ -15,7 +16,6 @@ struct ServiceContext; struct AudioTranscriptionRequest; class AudioSession; class Model; -class GenAIModelInstance; class ActionTracker; struct Request; @@ -36,7 +36,7 @@ class AudioTranscriptionsHandler : public HttpRequestHandler { /// Look up model in catalog and verify it's loaded and supports audio. std::shared_ptr ResolveModel(const std::string& model_name, - Model*& model, GenAIModelInstance*& loaded); + Model*& model, ModelLoadManager::LoadedModelLease& loaded); /// Build a Request with an OPENAI_JSON-tagged TEXT item from the original body string. void BuildOpenAIJsonRequest(const std::string& body, Request& session_request); diff --git a/sdk_v2/cpp/src/service/chat_completions_handler.cc b/sdk_v2/cpp/src/service/chat_completions_handler.cc index a351eaf7f..61b086767 100644 --- a/sdk_v2/cpp/src/service/chat_completions_handler.cc +++ b/sdk_v2/cpp/src/service/chat_completions_handler.cc @@ -7,6 +7,7 @@ #include "c_api_types.h" #include "catalog.h" #include "contracts/chat_completions.h" +#include "exception.h" #include "inferencing/generative/chat/chat_session.h" #include "inferencing/model_load_manager.h" #include "inferencing/session/session.h" @@ -20,6 +21,7 @@ #include #include +#include #include namespace fl { @@ -60,13 +62,13 @@ std::shared_ptr ChatCompletionsHandler::Pa } std::shared_ptr ChatCompletionsHandler::ResolveModel( - const std::string& model_name, Model*& model, GenAIModelInstance*& loaded) { + const std::string& model_name, Model*& model, ModelLoadManager::LoadedModelLease& loaded) { model = ctx_.catalog.GetModelVariant(model_name); if (!model) { return ErrorResponse(Status::CODE_404, "Model not found", "No model matching '" + model_name + "'"); } - loaded = ctx_.model_load_manager.GetLoadedModel(model->Id()); + loaded = ctx_.model_load_manager.AcquireLoadedModel(model->Id()); if (!loaded) { return ErrorResponse(Status::CODE_400, "Model not loaded", "Model '" + model_name + "' must be loaded before inference"); @@ -131,7 +133,7 @@ std::shared_ptr ChatCompletionsHandler::ha // 2. Resolve model std::string model_name = req.model; Model* model = nullptr; - GenAIModelInstance* loaded = nullptr; + ModelLoadManager::LoadedModelLease loaded; if (auto err = ResolveModel(model_name, model, loaded)) { tracker->SetStatus(ActionStatus::kClientError); return err; @@ -155,15 +157,15 @@ std::shared_ptr ChatCompletionsHandler::ha // 6. Run inference via ChatSession try { - std::unique_ptr session; + std::optional session; { ActionTracker create_tracker(Action::kSessionCreate, ctx_.telemetry, session_ctx); create_tracker.SetModelId(model_name); - session = std::make_unique(*model, *loaded, ctx_.logger, ctx_.telemetry); + session.emplace(*model, *loaded, ctx_.logger, ctx_.telemetry, true); + loaded.Release(); create_tracker.SetStatus(ActionStatus::kSuccess); } - ChatSession& session_ref = *session; - session_ref.SetRequestContext(session_ctx); + session->SetRequestContext(session_ctx); if (stream) { // The route action is recorded by the streaming thread when the stream @@ -171,8 +173,8 @@ std::shared_ptr ChatCompletionsHandler::ha return HandleStreaming(std::move(*session), std::move(session_request), include_usage_in_stream, std::move(tracker)); } else { - SessionRegistration reg(ctx_.session_manager, session_ref); - auto response = HandleNonStreaming(session_ref, session_request); + SessionRegistration reg(ctx_.session_manager, *session); + auto response = HandleNonStreaming(*session, session_request); tracker->SetStatus(ResponseToActionStatus(response)); return response; } @@ -215,12 +217,15 @@ std::shared_ptr ChatCompletionsHandler::Ha auto body_ptr = body; auto& logger = ctx_.logger; auto& thread_tracker = ctx_.thread_tracker; + auto stream_done = std::make_shared>(false); + auto stream_request = std::make_shared(std::move(session_request)); std::thread streaming_thread([bg_session = std::move(session), body_ptr, &logger, - req = std::move(session_request), + req = stream_request, include_usage, &thread_tracker, route_tracker = std::move(route_tracker), - &session_manager = ctx_.session_manager]() mutable { + &session_manager = ctx_.session_manager, + stream_done]() mutable { try { SessionRegistration reg(session_manager, bg_session); fl::Response bg_response; @@ -236,7 +241,6 @@ std::shared_ptr ChatCompletionsHandler::Ha if (item->type == FOUNDRY_LOCAL_ITEM_TEXT) { auto& text_item = static_cast(*item); if (!body_ptr->Push("data: " + text_item.text + "\n\n")) { - req.canceled.store(true, std::memory_order_relaxed); return 1; } } else { @@ -248,7 +252,10 @@ std::shared_ptr ChatCompletionsHandler::Ha }; bg_session.SetStreamingCallback(callback_fn); - bg_session.ProcessRequest(req, bg_response); + bg_session.ProcessRequest(*req, bg_response); + if (req->canceled.load(std::memory_order_relaxed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "chat completion stream cancelled"); + } // Usage chunk — only if stream_options.include_usage was true if (include_usage) { @@ -269,16 +276,17 @@ std::shared_ptr ChatCompletionsHandler::Ha usage_chunk.usage = std::move(usage); if (!body_ptr->Push("data: " + nlohmann::json(usage_chunk).dump() + "\n\n")) { - req.canceled.store(true, std::memory_order_relaxed); + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "chat completion stream cancelled"); } } - body_ptr->Push("data: [DONE]\n\n"); + if (!body_ptr->Push("data: [DONE]\n\n")) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "chat completion stream cancelled"); + } - // Record final route status after streaming completes. + // Inference streamed to completion — record the route action as a success. if (route_tracker) { - route_tracker->SetStatus(req.canceled.load(std::memory_order_relaxed) ? ActionStatus::kCanceled - : ActionStatus::kSuccess); + route_tracker->SetStatus(ActionStatus::kSuccess); } } catch (const std::exception& ex) { nlohmann::json err = { @@ -286,7 +294,8 @@ std::shared_ptr ChatCompletionsHandler::Ha }; body_ptr->Push("data: " + err.dump() + "\n\n"); - // Mid-stream failure: record the exception; the route action keeps kFailure. + // Mid-stream failure: record the exception. The route action keeps its + // default kFailure status. if (route_tracker) { route_tracker->RecordException(ex); } @@ -295,11 +304,14 @@ std::shared_ptr ChatCompletionsHandler::Ha body_ptr->Finish(); // route_tracker is destroyed with this closure once the thread completes, // recording the route action with the full streaming duration and final status. - route_tracker.reset(); - thread_tracker.Remove(std::this_thread::get_id()); + stream_done->store(true, std::memory_order_release); + thread_tracker.NotifyCompleted(); }); - thread_tracker.Track(std::move(streaming_thread)); + thread_tracker.Track(std::move(streaming_thread), stream_done, [body, stream_request] { + body->Abort(); + stream_request->canceled.store(true, std::memory_order_relaxed); + }); auto response = oatpp::web::protocol::http::outgoing::Response::createShared( Status::CODE_200, body); diff --git a/sdk_v2/cpp/src/service/chat_completions_handler.h b/sdk_v2/cpp/src/service/chat_completions_handler.h index c7225bf77..bb6a010af 100644 --- a/sdk_v2/cpp/src/service/chat_completions_handler.h +++ b/sdk_v2/cpp/src/service/chat_completions_handler.h @@ -4,6 +4,7 @@ #ifdef FOUNDRY_LOCAL_HAS_WEB_SERVICE +#include "inferencing/model_load_manager.h" #include "service/handler_utils.h" #include @@ -15,7 +16,6 @@ struct ServiceContext; struct ChatCompletionRequest; class ChatSession; class Model; -class GenAIModelInstance; class ActionTracker; struct Request; @@ -38,7 +38,7 @@ class ChatCompletionsHandler : public HttpRequestHandler { /// Look up model in catalog and verify it's loaded. Sets output pointers. /// Returns an error response on failure, nullptr on success. std::shared_ptr ResolveModel(const std::string& model_name, - Model*& model, GenAIModelInstance*& loaded); + Model*& model, ModelLoadManager::LoadedModelLease& loaded); /// Build a Request with an OPENAI_JSON-tagged TEXT item from the original body string. /// Populates catalog defaults and model-specific options (tool_call_start/end). diff --git a/sdk_v2/cpp/src/service/embeddings_handler.cc b/sdk_v2/cpp/src/service/embeddings_handler.cc index b70c46f5c..6bf6d7997 100644 --- a/sdk_v2/cpp/src/service/embeddings_handler.cc +++ b/sdk_v2/cpp/src/service/embeddings_handler.cc @@ -8,6 +8,8 @@ #include +#include + #include "catalog.h" #include "contracts/embeddings.h" #include "exception.h" @@ -68,7 +70,7 @@ class EmbeddingsHandler : public HttpRequestHandler { return ErrorResponse(Status::CODE_404, "Model not found", model_name); } - auto* loaded = ctx_.model_load_manager.GetLoadedModel(model->Id()); + auto loaded = ctx_.model_load_manager.AcquireLoadedModel(model->Id()); if (!loaded) { tracker.SetStatus(ActionStatus::kClientError); return ErrorResponse(Status::CODE_400, "Model not loaded", model_name); @@ -81,11 +83,12 @@ class EmbeddingsHandler : public HttpRequestHandler { // 4. Create session and process each input try { - std::unique_ptr session; + std::optional session; { ActionTracker create_tracker(Action::kSessionCreate, ctx_.telemetry, session_ctx); create_tracker.SetModelId(model_name); - session = std::make_unique(*model, *loaded, ctx_.logger, ctx_.telemetry); + session.emplace(*model, *loaded, ctx_.logger, ctx_.telemetry, true); + loaded.Release(); create_tracker.SetStatus(ActionStatus::kSuccess); } session->SetRequestContext(session_ctx); diff --git a/sdk_v2/cpp/src/service/handler_utils.h b/sdk_v2/cpp/src/service/handler_utils.h index b7f3a01d1..1a37953c4 100644 --- a/sdk_v2/cpp/src/service/handler_utils.h +++ b/sdk_v2/cpp/src/service/handler_utils.h @@ -10,8 +10,11 @@ #include #include +#include "telemetry/telemetry_redaction.h" #include "telemetry/telemetry.h" +#include +#include #include #include #include @@ -88,7 +91,38 @@ inline std::string GetUserAgent(const std::shared_ptrgetHeader("User-Agent"); - return ua ? *ua : std::string{}; + if (!ua) { + return {}; + } + std::string scrubbed; + scrubbed.reserve(std::min(ua->size(), 256)); + size_t run_length = 0; + for (char ch : *ua) { + unsigned char c = static_cast(ch); + bool safe = std::isalnum(c) || ch == ' ' || ch == '.' || ch == '-' || ch == '_' || ch == '/' || + ch == '(' || ch == ')' || ch == ';'; + if (!safe) { + scrubbed.push_back('?'); + run_length = 0; + continue; + } + + if (std::isalnum(c)) { + ++run_length; + if (run_length > 32) { + scrubbed.push_back('?'); + continue; + } + } else { + run_length = 0; + } + scrubbed.push_back(ch); + } + constexpr size_t kMaxUserAgentLength = 256; + if (scrubbed.size() > kMaxUserAgentLength) { + scrubbed.resize(kMaxUserAgentLength); + } + return scrubbed; } // ======================================================================== @@ -99,14 +133,12 @@ inline std::string GetUserAgent(const std::shared_ptr lock(mutex_); - if (done_ || queue_.size() >= kMaxQueuedChunks) { - done_ = true; - cv_.notify_one(); + if (aborted_) { return false; } queue_.push(std::move(chunk)); @@ -121,6 +153,18 @@ class SseStreamBody : public oatpp::web::protocol::http::outgoing::Body { cv_.notify_one(); } + void Abort() { + std::lock_guard lock(mutex_); + aborted_ = true; + done_ = true; + cv_.notify_all(); + } + + bool IsAborted() const { + std::lock_guard lock(mutex_); + return aborted_; + } + // -- Body interface -- oatpp::v_io_size read(void* buffer, v_buff_size count, oatpp::async::Action& /*action*/) override { @@ -168,11 +212,11 @@ class SseStreamBody : public oatpp::web::protocol::http::outgoing::Body { v_int64 getKnownSize() override { return -1; } // unknown → chunked transfer private: - std::mutex mutex_; + mutable std::mutex mutex_; std::condition_variable cv_; std::queue queue_; bool done_; - static constexpr size_t kMaxQueuedChunks = 1024; + bool aborted_; }; } // namespace fl diff --git a/sdk_v2/cpp/src/service/responses_handler.cc b/sdk_v2/cpp/src/service/responses_handler.cc index 1ced57750..cd9730227 100644 --- a/sdk_v2/cpp/src/service/responses_handler.cc +++ b/sdk_v2/cpp/src/service/responses_handler.cc @@ -6,6 +6,7 @@ #include "c_api_types.h" #include "catalog.h" +#include "exception.h" #include "inferencing/generative/chat/chat_session.h" #include "inferencing/generative/openresponses/response_converter.h" #include "inferencing/generative/openresponses/response_store.h" @@ -79,13 +80,13 @@ std::shared_ptr ResponsesHandler::ParseAnd } std::shared_ptr ResponsesHandler::ResolveModel( - const std::string& model_name, Model*& model, GenAIModelInstance*& loaded) { + const std::string& model_name, Model*& model, ModelLoadManager::LoadedModelLease& loaded) { model = ctx_.catalog.GetModelVariant(model_name); if (!model) { return ErrorResponse(Status::CODE_404, "Model not found", "No model matching '" + model_name + "'"); } - loaded = ctx_.model_load_manager.GetLoadedModel(model->Id()); + loaded = ctx_.model_load_manager.AcquireLoadedModel(model->Id()); if (!loaded) { return ErrorResponse(Status::CODE_400, "Model not loaded", "Model '" + model_name + "' must be loaded before inference"); @@ -156,7 +157,7 @@ std::shared_ptr ResponsesHandler::handle( // 2. Resolve model std::string model_name = params.model; Model* model = nullptr; - GenAIModelInstance* loaded = nullptr; + ModelLoadManager::LoadedModelLease loaded; if (auto err = ResolveModel(model_name, model, loaded)) { tracker->SetStatus(ActionStatus::kClientError); return err; @@ -177,9 +178,10 @@ std::shared_ptr ResponsesHandler::handle( std::string response_id = ResponseConverter::GenerateId("resp"); std::unique_ptr session; + const std::string resolved_model_id = model->Id(); if (params.previous_response_id) { - session = ctx_.session_manager.CheckOut(*params.previous_response_id); + session = ctx_.session_manager.CheckOut(*params.previous_response_id, resolved_model_id); if (session) { ctx_.logger.Log(LogLevel::Information, @@ -205,7 +207,8 @@ std::shared_ptr ResponsesHandler::handle( if (!session) { ActionTracker create_tracker(Action::kSessionCreate, ctx_.telemetry, session_ctx); create_tracker.SetModelId(model_name); - session = std::make_unique(*model, *loaded, ctx_.logger, ctx_.telemetry); + session = std::make_unique(*model, *loaded, ctx_.logger, ctx_.telemetry, true); + loaded.Release(); create_tracker.SetStatus(ActionStatus::kSuccess); } session->SetRequestContext(session_ctx); @@ -223,22 +226,20 @@ std::shared_ptr ResponsesHandler::handle( // The route action is recorded by the streaming thread when the stream // finishes — move the tracker in rather than marking success now. - return HandleStreaming(std::move(session), std::move(session_request), model_name, + return HandleStreaming(std::move(session), std::move(session_request), model_name, resolved_model_id, response_id, created_at, params, req_json, std::move(tracker)); } else { ctx_.logger.Log(LogLevel::Debug, fmt::format("Creating response {} for model {}", response_id, model_name)); - auto response = HandleNonStreaming(std::move(session), session_request, model_name, + auto response = HandleNonStreaming(std::move(session), session_request, model_name, resolved_model_id, response_id, created_at, params, req_json); tracker->SetStatus(ActionStatus::kSuccess); return response; } } catch (const std::exception& ex) { - if (tracker) { - tracker->RecordException(ex); - } + tracker->RecordException(ex); ctx_.logger.Log(LogLevel::Error, fmt::format("Response {} failed: {}", response_id, ex.what())); @@ -253,7 +254,7 @@ std::shared_ptr ResponsesHandler::handle( std::shared_ptr ResponsesHandler::HandleNonStreaming( std::unique_ptr session, Request& session_request, - const std::string& model_name, const std::string& response_id, + const std::string& model_name, const std::string& model_id, const std::string& response_id, int64_t created_at, const ResponseCreateParams& params, const nlohmann::json& req_json) { SessionRegistration reg(ctx_.session_manager, *session); @@ -285,7 +286,7 @@ std::shared_ptr ResponsesHandler::HandleNo session->SetStreamingCallback(nullptr); // Cache the session for potential reuse on the next turn - ctx_.session_manager.CheckIn(response_id, std::move(session)); + ctx_.session_manager.CheckIn(response_id, std::move(session), model_id); } return JsonResponse(Status::CODE_200, response_json); @@ -293,7 +294,7 @@ std::shared_ptr ResponsesHandler::HandleNo std::shared_ptr ResponsesHandler::HandleStreaming( std::unique_ptr session, Request session_request, - const std::string& model_name, const std::string& response_id, + const std::string& model_name, const std::string& model_id, const std::string& response_id, int64_t created_at, const ResponseCreateParams& params, const nlohmann::json& req_json, std::unique_ptr route_tracker) { auto body = std::make_shared(); @@ -331,60 +332,61 @@ std::shared_ptr ResponsesHandler::HandleSt auto& logger = ctx_.logger; auto& session_manager = ctx_.session_manager; auto& tracker = ctx_.thread_tracker; + auto stream_done = std::make_shared>(false); + auto stream_request = std::make_shared(std::move(session_request)); // Background thread is required: oatpp needs the Response returned immediately so it can // start writing SSE events. ProcessRequest blocks until generation completes. std::thread streaming_thread([body_ptr, &logger, &session_manager, session = std::move(session), - req = std::move(session_request), + req = stream_request, model_name, response_id, created_at, + model_id, should_store, &store, req_copy = std::move(req_copy), params_copy = std::move(params_copy), route_tracker = std::move(route_tracker), - &tracker]() mutable { + &tracker, + stream_done]() mutable { + bool terminal_sent = false; int seq = 2; - try { - SessionRegistration reg(session_manager, *session); - - std::string full_text; // concatenation of all visible runs, used for output_text in completed_response - - // Per-item state for the currently-open item. `current_kind == nullopt` means no item is open. On every type - // transition we close the open item (emitting its done events) and open a new one with a fresh id at the next - // output_index. Adjacent same-typed segments accumulate into the same item naturally because we don't close - // until the type changes. - enum class ItemKind { Reasoning, - Message }; - std::optional current_kind; - std::string current_id; - std::string current_text; - int current_output_index = -1; - int next_output_index = 0; - - // Items that have been *closed* (or, for the currently-open item at end-of-stream, finalized in place). - // Used to construct the final `output[]` array for the response.completed event. - std::vector closed_items; - - auto push_event = [&](const std::string& event_name, const StreamEvent& ev) { - if (!body_ptr->Push("event: " + event_name + "\ndata: " + nlohmann::json(ev).dump() + "\n\n")) { - req.canceled.store(true, std::memory_order_relaxed); - } - }; + std::string full_text; // concatenation of all visible runs, used for output_text in completed_response + + // Per-item state for the currently-open item. `current_kind == nullopt` means no item is open. On every type + // transition we close the open item (emitting its done events) and open a new one with a fresh id at the next + // output_index. Adjacent same-typed segments accumulate into the same item naturally because we don't close + // until the type changes. + enum class ItemKind { Reasoning, + Message }; + std::optional current_kind; + std::string current_id; + std::string current_text; + int current_output_index = -1; + int next_output_index = 0; + + // Items that have been *closed* (or, for the currently-open item at end-of-stream, finalized in place). + // Used to construct the final `output[]` array for the response.completed event. + std::vector closed_items; + std::optional reg; + + auto push_event = [&](const std::string& event_name, const StreamEvent& ev) { + return body_ptr->Push("event: " + event_name + "\ndata: " + nlohmann::json(ev).dump() + "\n\n"); + }; - auto close_current = [&]() { - if (!current_kind.has_value()) { - return; - } + auto close_current = [&]() { + if (!current_kind.has_value()) { + return; + } - if (*current_kind == ItemKind::Reasoning) { - // Emit: response.reasoning.done - StreamEvent done_ev; - done_ev.type = StreamEventType::kReasoningDone; - done_ev.sequence_number = seq++; - done_ev.output_index = current_output_index; - done_ev.item_id = current_id; - done_ev.text = current_text; - push_event("response.reasoning.done", done_ev); + if (*current_kind == ItemKind::Reasoning) { + // Emit: response.reasoning.done + StreamEvent done_ev; + done_ev.type = StreamEventType::kReasoningDone; + done_ev.sequence_number = seq++; + done_ev.output_index = current_output_index; + done_ev.item_id = current_id; + done_ev.text = current_text; + push_event("response.reasoning.done", done_ev); // Emit: response.output_item.done ReasoningOutputItem rs; @@ -492,6 +494,8 @@ std::shared_ptr ResponsesHandler::HandleSt push_event("response.content_part.added", part_added); }; + try { + reg.emplace(session_manager, *session); fl::Response bg_response; fl::Session::StreamingCallbackFn callback_fn = [&](flStreamingCallbackData event, void* /*user_data*/) -> int { fl::ItemQueue* queue = reinterpret_cast(event.item_queue); @@ -529,7 +533,9 @@ std::shared_ptr ResponsesHandler::HandleSt delta.output_index = current_output_index; delta.item_id = current_id; delta.delta = text_item->text; - push_event("response.reasoning.delta", delta); + if (!push_event("response.reasoning.delta", delta)) { + return 1; + } } else { full_text += text_item->text; @@ -540,7 +546,9 @@ std::shared_ptr ResponsesHandler::HandleSt text_delta.content_index = 0; text_delta.item_id = current_id; text_delta.delta = text_item->text; - push_event("response.output_text.delta", text_delta); + if (!push_event("response.output_text.delta", text_delta)) { + return 1; + } } return 0; @@ -548,28 +556,15 @@ std::shared_ptr ResponsesHandler::HandleSt session->SetStreamingCallback(callback_fn); - session->ProcessRequest(req, bg_response); + session->ProcessRequest(*req, bg_response); + if (req->canceled.load(std::memory_order_relaxed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "response stream cancelled"); + } // Close whatever item is still open at end-of-generation so the SSE stream is well-formed. close_current(); - - if (req.canceled.load(std::memory_order_relaxed)) { - auto canceled_response = ResponseConverter::BuildFailedResponseObject( - response_id, created_at, model_name, params_copy, "canceled", "Response generation was canceled"); - - StreamEvent failed; - failed.type = StreamEventType::kResponseFailed; - failed.sequence_number = seq++; - failed.response = canceled_response; - push_event("response.failed", failed); - - if (route_tracker) { - route_tracker->SetStatus(ActionStatus::kCanceled); - } - body_ptr->Finish(); - route_tracker.reset(); - tracker.Remove(std::this_thread::get_id()); - return; + if (req->canceled.load(std::memory_order_relaxed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "response stream cancelled"); } auto completed_response = ResponseConverter::BuildResponseObject( @@ -579,7 +574,9 @@ std::shared_ptr ResponsesHandler::HandleSt completed.type = StreamEventType::kResponseCompleted; completed.sequence_number = seq++; completed.response = completed_response; - push_event("response.completed", completed); + if (!push_event("response.completed", completed)) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "response stream cancelled"); + } // Store if requested if (should_store) { @@ -588,19 +585,23 @@ std::shared_ptr ResponsesHandler::HandleSt store.Store(response_id, response_json, std::move(input_items)); // Deregister before caching — see non-streaming path comment. - reg.Release(); + reg->Release(); // Clear per-request streaming callback before caching — see non-streaming path comment. session->SetStreamingCallback(nullptr); // Cache the session for potential reuse on the next turn - session_manager.CheckIn(response_id, std::move(session)); + session_manager.CheckIn(response_id, std::move(session), model_id); + } + + if (!body_ptr->Push("data: [DONE]\n\n")) { + FL_THROW(FOUNDRY_LOCAL_ERROR_OPERATION_CANCELLED, "response stream cancelled"); } + terminal_sent = true; - // Record final route status after streaming completes. + // Streamed to completion — record the route action as a success. if (route_tracker) { - route_tracker->SetStatus(req.canceled.load(std::memory_order_relaxed) ? ActionStatus::kCanceled - : ActionStatus::kSuccess); + route_tracker->SetStatus(ActionStatus::kSuccess); } } catch (const std::exception& ex) { @@ -622,15 +623,19 @@ std::shared_ptr ResponsesHandler::HandleSt } } - // Terminal event per spec - body_ptr->Push("data: [DONE]\n\n"); + if (!terminal_sent) { + body_ptr->Push("data: [DONE]\n\n"); + } body_ptr->Finish(); + stream_done->store(true, std::memory_order_release); + tracker.NotifyCompleted(); - route_tracker.reset(); - tracker.Remove(std::this_thread::get_id()); }); - tracker.Track(std::move(streaming_thread)); + tracker.Track(std::move(streaming_thread), stream_done, [body, stream_request] { + body->Abort(); + stream_request->canceled.store(true, std::memory_order_relaxed); + }); auto response = oatpp::web::protocol::http::outgoing::Response::createShared( Status::CODE_200, body); diff --git a/sdk_v2/cpp/src/service/responses_handler.h b/sdk_v2/cpp/src/service/responses_handler.h index e455d6203..7ce6af59c 100644 --- a/sdk_v2/cpp/src/service/responses_handler.h +++ b/sdk_v2/cpp/src/service/responses_handler.h @@ -4,6 +4,7 @@ #ifdef FOUNDRY_LOCAL_HAS_WEB_SERVICE +#include "inferencing/model_load_manager.h" #include "service/handler_utils.h" #include @@ -15,7 +16,6 @@ struct ServiceContext; struct Request; class ChatSession; class Model; -class GenAIModelInstance; class ActionTracker; namespace responses { @@ -44,7 +44,7 @@ class ResponsesHandler : public HttpRequestHandler { /// Look up model in catalog and verify it's loaded. Sets output pointers. /// Returns an error response on failure, nullptr on success. std::shared_ptr ResolveModel(const std::string& model_name, - Model*& model, GenAIModelInstance*& loaded); + Model*& model, ModelLoadManager::LoadedModelLease& loaded); /// Load previous response context when chaining via previous_response_id. /// The json storage objects are passed by reference because the output pointers alias into them. @@ -57,14 +57,14 @@ class ResponsesHandler : public HttpRequestHandler { // --- Inference dispatch --- std::shared_ptr HandleNonStreaming(std::unique_ptr session, Request& session_request, - const std::string& model_name, const std::string& response_id, - int64_t created_at, + const std::string& model_name, const std::string& model_id, + const std::string& response_id, int64_t created_at, const responses::ResponseCreateParams& params, const nlohmann::json& req_json); std::shared_ptr HandleStreaming(std::unique_ptr session, Request session_request, - const std::string& model_name, const std::string& response_id, - int64_t created_at, + const std::string& model_name, const std::string& model_id, + const std::string& response_id, int64_t created_at, const responses::ResponseCreateParams& params, const nlohmann::json& req_json, std::unique_ptr route_tracker); diff --git a/sdk_v2/cpp/src/service/web_service.cc b/sdk_v2/cpp/src/service/web_service.cc index 65f39d3f7..a88a9e07e 100644 --- a/sdk_v2/cpp/src/service/web_service.cc +++ b/sdk_v2/cpp/src/service/web_service.cc @@ -13,6 +13,7 @@ #include "service/handler_utils.h" #include "service/models_handlers.h" #include "service/responses_handler.h" +#include "telemetry/telemetry_action_tracker.h" #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include @@ -38,6 +40,10 @@ namespace fl { +namespace { + +} // namespace + // ======================================================================== // ForceCloseConnectionProvider // @@ -135,7 +141,9 @@ class StatusHandler : public HttpRequestHandler { public: explicit StatusHandler(ServiceContext& ctx) : ctx_(ctx) {} - std::shared_ptr handle(const std::shared_ptr&) override { + std::shared_ptr handle(const std::shared_ptr& request) override { + (void)request; + nlohmann::json body = { {"modelCachePath", ctx_.model_cache_dir}, {"endpoints", ctx_.bound_urls}, @@ -163,9 +171,13 @@ class ShutdownHandler : public HttpRequestHandler { : shutdown_fn_(std::move(shutdown_fn)) {} std::shared_ptr handle(const std::shared_ptr&) override { - shutdown_fn_(); - nlohmann::json body = {{"status", "shutting_down"}}; + std::thread([fn = shutdown_fn_] { + try { + fn(); + } catch (...) { + } + }).detach(); return JsonResponse(Status::CODE_200, body); } @@ -238,11 +250,13 @@ WebService::~WebService() { std::vector WebService::Start(const std::vector& endpoints) { auto& ctx = *impl_->context; - if (impl_->running.load()) { + if (impl_->running.exchange(true)) { ctx.logger.Log(LogLevel::Information, "Web service is already running."); FL_THROW(FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Web service is already running"); } + impl_->thread_tracker.Reset(); + // Create shared router and register all routes impl_->router = oatpp::web::server::HttpRouter::createShared(); @@ -275,89 +289,94 @@ std::vector WebService::Start(const std::vector& endpo std::vector bound_urls; - for (const auto& endpoint : endpoints) { - // Parse "http://host:port" — strip scheme, extract host:port - std::string host = "127.0.0.1"; - uint16_t port = 0; + try { + for (const auto& endpoint : endpoints) { + // Parse "http://host:port" — strip scheme, extract host:port + std::string host = "127.0.0.1"; + uint16_t port = 0; - std::string addr = endpoint; - auto scheme_end = addr.find("://"); - if (scheme_end != std::string::npos) { - addr = addr.substr(scheme_end + 3); - } + std::string addr = endpoint; + auto scheme_end = addr.find("://"); + if (scheme_end != std::string::npos) { + addr = addr.substr(scheme_end + 3); + } - if (!addr.empty() && addr.back() == '/') { - addr.pop_back(); - } + if (!addr.empty() && addr.back() == '/') { + addr.pop_back(); + } - auto colon = addr.rfind(':'); - if (colon != std::string::npos) { - host = addr.substr(0, colon); - port = static_cast(std::stoi(addr.substr(colon + 1))); - } else { - host = addr; - } + auto colon = addr.rfind(':'); + if (colon != std::string::npos) { + host = addr.substr(0, colon); + port = static_cast(std::stoi(addr.substr(colon + 1))); + } else { + host = addr; + } - auto tcp_provider = oatpp::network::tcp::server::ConnectionProvider::createShared({host.c_str(), port}); + auto tcp_provider = oatpp::network::tcp::server::ConnectionProvider::createShared({host.c_str(), port}); - // oatpp resolves ephemeral port 0 during construction via getsockname() - // and stores it in PROPERTY_PORT. getAddress().port is stale — use the property. - auto resolved_port = tcp_provider->getProperty(oatpp::network::ConnectionProvider::PROPERTY_PORT); + // oatpp resolves ephemeral port 0 during construction via getsockname() + // and stores it in PROPERTY_PORT. getAddress().port is stale — use the property. + auto resolved_port = tcp_provider->getProperty(oatpp::network::ConnectionProvider::PROPERTY_PORT); - if (resolved_port) { - port = static_cast(std::stoi(resolved_port.std_str())); - } + if (resolved_port) { + port = static_cast(std::stoi(resolved_port.std_str())); + } - // Wrap with our force-close provider so that connection invalidation also unblocks any pending recv() in worker - // threads (via CancelIoEx on Windows; shutdown() suffices on POSIX). Without this, HttpConnectionHandler::stop() - // can wait ~120s for a single keep-alive client to drop its connection. - auto provider = std::static_pointer_cast( - std::make_shared(tcp_provider)); + // Wrap with our force-close provider so that connection invalidation also unblocks any pending recv() in worker + // threads (via CancelIoEx on Windows; shutdown() suffices on POSIX). Without this, HttpConnectionHandler::stop() + // can wait ~120s for a single keep-alive client to drop its connection. + auto provider = std::static_pointer_cast( + std::make_shared(tcp_provider)); - auto server = std::make_shared(provider, impl_->connection_handler); + auto server = std::make_shared(provider, impl_->connection_handler); - impl_->providers.push_back(provider); - impl_->servers.push_back(server); + impl_->providers.push_back(provider); + impl_->servers.push_back(server); - // Start server on a background thread - impl_->listener_threads.emplace_back([server]() { - server->run(); - }); + // Start server on a background thread + impl_->listener_threads.emplace_back([server]() { + server->run(); + }); - // Wait until oatpp's server reports STATUS_RUNNING (i.e. the listener loop - // has started and is accepting connections). Bounded poll with a 5s timeout. - const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); - while (server->getStatus() != oatpp::network::Server::STATUS_RUNNING && - std::chrono::steady_clock::now() < deadline) { - std::this_thread::sleep_for(std::chrono::milliseconds(1)); - } + // Wait until oatpp's server reports STATUS_RUNNING (i.e. the listener loop + // has started and is accepting connections). Bounded poll with a 5s timeout. + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (server->getStatus() != oatpp::network::Server::STATUS_RUNNING && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } - if (server->getStatus() != oatpp::network::Server::STATUS_RUNNING) { - // Tear down anything we already started in this call so we don't leak - // listener threads. `running` is still false, so the destructor would skip Stop(). - impl_->running.store(true); - Stop(); - FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, - "Web service failed to reach RUNNING state for endpoint ", endpoint, " within 5s"); - } + if (server->getStatus() != oatpp::network::Server::STATUS_RUNNING) { + // Tear down anything we already started in this call so we don't leak listener threads. + Stop(); + FL_THROW(FOUNDRY_LOCAL_ERROR_INTERNAL, + "Web service failed to reach RUNNING state for endpoint ", endpoint, " within 5s"); + } - std::string bound_url = "http://" + host + ":" + std::to_string(port); - bound_urls.push_back(bound_url); + std::string bound_url = "http://" + host + ":" + std::to_string(port); + bound_urls.push_back(bound_url); - ctx.logger.Log(LogLevel::Information, fmt::format("Web service listening on {}", bound_url)); + ctx.logger.Log(LogLevel::Information, fmt::format("Web service listening on {}", bound_url)); + } + } catch (...) { + Stop(); + throw; } ctx.bound_urls = bound_urls; - impl_->running.store(true); return bound_urls; } void WebService::Stop() { - if (!impl_->running.load()) { + if (!impl_->running.load() && impl_->servers.empty() && impl_->providers.empty() && + impl_->listener_threads.empty() && !impl_->connection_handler) { return; } + impl_->thread_tracker.BeginStopping(); + // Stop accepting new connections first, then stop server loops. for (auto& provider : impl_->providers) { provider->stop(); @@ -367,6 +386,8 @@ void WebService::Stop() { server->stop(); } + impl_->thread_tracker.AbortAll(); + // Stop per-connection worker tasks before releasing router/handlers. // // HttpConnectionHandler::stop() invalidates every open connection (our ForceCloseConnectionProvider's invalidator @@ -378,16 +399,16 @@ void WebService::Stop() { impl_->connection_handler->stop(); } + // Join streaming producers after connection shutdown so no new streaming + // thread can be tracked while Stop() is taking its join snapshot. + impl_->thread_tracker.JoinAll(); + for (auto& thread : impl_->listener_threads) { if (thread.joinable()) { thread.join(); } } - // Streaming threads hold request/session state; Manager::Shutdown cancels active sessions - // and waits for drain before StopWebService reaches here. - impl_->thread_tracker.JoinAll(); - impl_->servers.clear(); impl_->listener_threads.clear(); impl_->providers.clear(); diff --git a/sdk_v2/cpp/src/service/web_service.h b/sdk_v2/cpp/src/service/web_service.h index b7d72bd17..fb44cfea2 100644 --- a/sdk_v2/cpp/src/service/web_service.h +++ b/sdk_v2/cpp/src/service/web_service.h @@ -4,6 +4,8 @@ #include "logger.h" +#include +#include #include #include #include @@ -21,47 +23,145 @@ class ResponseStore; /// Tracks streaming threads so they can be joined on shutdown. /// Handlers call Track() instead of std::thread::detach(). -/// Threads call Remove() when done to clean up immediately. class StreamingThreadTracker { + struct TrackedThread { + std::thread thread; + std::shared_ptr> done; + std::function abort; + }; + public: + StreamingThreadTracker() { + StartReaperLocked(); + } + + ~StreamingThreadTracker() { + StopReaper(); + JoinAll(); + } + /// Take ownership of a streaming thread. - void Track(std::thread t) { + void Track(std::thread t, std::shared_ptr> done, std::function abort) { std::lock_guard lock(mutex_); - threads_.push_back(std::move(t)); + StartReaperLocked(); + ReapCompletedLocked(); + if (stopping_) { + if (abort) { + abort(); + } + } + threads_.push_back(TrackedThread{std::move(t), std::move(done), std::move(abort)}); + if (threads_.back().done && threads_.back().done->load(std::memory_order_acquire)) { + completion_cv_.notify_one(); + } } - /// Called from within a thread to untrack itself after work is done. - /// Detaches the thread (can't join itself) and removes the entry. - void Remove(std::thread::id id) { + void NotifyCompleted() { + completion_cv_.notify_one(); + } + + void BeginStopping() { std::lock_guard lock(mutex_); - for (auto it = threads_.begin(); it != threads_.end(); ++it) { - if (it->get_id() == id) { - it->detach(); - threads_.erase(it); - return; + stopping_ = true; + AbortAllLocked(); + } + + void AbortAll() { + std::lock_guard lock(mutex_); + AbortAllLocked(); + } + + void Reset() { + std::lock_guard lock(mutex_); + stopping_ = false; + StartReaperLocked(); + } + + void AbortAllLocked() { + for (auto& tracked : threads_) { + if (tracked.abort) { + tracked.abort(); } } } /// Join all remaining threads. Called by WebService::Stop(). - /// Moves entries out before joining to avoid deadlock with Remove(). + /// Moves entries out before joining so completed streaming threads are cleaned up safely. void JoinAll() { - std::vector local; + StopReaper(); + + std::vector local; { std::lock_guard lock(mutex_); local = std::move(threads_); } - for (auto& t : local) { - if (t.joinable()) { - t.join(); + for (auto& tracked : local) { + if (tracked.thread.joinable()) { + tracked.thread.join(); } } } private: + bool HasCompletedLocked() const { + for (const auto& tracked : threads_) { + if (tracked.done && tracked.done->load(std::memory_order_acquire)) { + return true; + } + } + return false; + } + + void StartReaperLocked() { + if (reaper_thread_.joinable()) { + return; + } + reaper_stop_ = false; + reaper_thread_ = std::thread([this]() { ReaperLoop(); }); + } + + void StopReaper() { + { + std::lock_guard lock(mutex_); + reaper_stop_ = true; + } + completion_cv_.notify_all(); + if (reaper_thread_.joinable()) { + reaper_thread_.join(); + } + } + + void ReaperLoop() { + while (true) { + std::unique_lock lock(mutex_); + completion_cv_.wait(lock, [this]() { return reaper_stop_ || HasCompletedLocked(); }); + if (reaper_stop_) { + return; + } + ReapCompletedLocked(); + } + } + + void ReapCompletedLocked() { + for (auto it = threads_.begin(); it != threads_.end();) { + if (it->done && it->done->load(std::memory_order_acquire)) { + if (it->thread.joinable()) { + it->thread.join(); + } + it = threads_.erase(it); + } else { + ++it; + } + } + } + std::mutex mutex_; - std::vector threads_; + std::condition_variable completion_cv_; + std::vector threads_; + std::thread reaper_thread_; + bool stopping_ = false; + bool reaper_stop_ = false; }; /// Context shared with all HTTP controllers. diff --git a/sdk_v2/cpp/test/CMakeLists.txt b/sdk_v2/cpp/test/CMakeLists.txt index 3a5eb51b3..fd6024836 100644 --- a/sdk_v2/cpp/test/CMakeLists.txt +++ b/sdk_v2/cpp/test/CMakeLists.txt @@ -7,6 +7,7 @@ find_package(httplib CONFIG REQUIRED) # Links against the static library with access to internal src/ headers. # ========================================================================== add_executable(foundry_local_tests + test_main.cc internal_api/audio/audio_session_test.cc internal_api/audio/audio_transcription_contract_test.cc internal_api/audio/pcm_utils_test.cc @@ -67,7 +68,7 @@ target_compile_options(foundry_local_tests PRIVATE ${FOUNDRY_LOCAL_COMPILE_OPTIO target_link_libraries(foundry_local_tests PRIVATE foundry_local_static - GTest::gtest_main + GTest::gtest httplib::httplib ) @@ -127,6 +128,7 @@ gtest_discover_tests(foundry_local_tests # ========================================================================== add_executable(sdk_integration_tests + test_main.cc sdk_api/audio_transcriptions_test.cc sdk_api/cpp_api_test.cc sdk_api/catalog_test.cc @@ -163,7 +165,7 @@ target_include_directories(sdk_integration_tests target_link_libraries(sdk_integration_tests PRIVATE foundry_local_cpp - GTest::gtest_main + GTest::gtest nlohmann_json::nlohmann_json httplib::httplib ) @@ -211,6 +213,7 @@ set_tests_properties(sdk_integration_tests PROPERTIES TIMEOUT 1200) # their own Manager instances with different configurations. # ========================================================================== add_executable(cache_only_tests + test_main.cc sdk_api/cache_only_test.cc sdk_api/catalog_live_test.cc sdk_api/manager_web_service_test.cc @@ -228,7 +231,7 @@ target_include_directories(cache_only_tests target_link_libraries(cache_only_tests PRIVATE foundry_local_cpp - GTest::gtest_main + GTest::gtest nlohmann_json::nlohmann_json ) diff --git a/sdk_v2/cpp/test/internal_api/audio/audio_session_test.cc b/sdk_v2/cpp/test/internal_api/audio/audio_session_test.cc index 35efe2b2a..b966dd835 100644 --- a/sdk_v2/cpp/test/internal_api/audio/audio_session_test.cc +++ b/sdk_v2/cpp/test/internal_api/audio/audio_session_test.cc @@ -18,9 +18,9 @@ #include "logger.h" #include "model.h" #include "internal_api/null_session_manager.h" -#include "telemetry/telemetry_logger.h" #include "internal_api/test_helpers.h" #include "internal_api/test_model_cache.h" +#include "telemetry/telemetry_logger.h" #include diff --git a/sdk_v2/cpp/test/internal_api/azure_catalog_test.cc b/sdk_v2/cpp/test/internal_api/azure_catalog_test.cc index 49a05f096..0ebb48881 100644 --- a/sdk_v2/cpp/test/internal_api/azure_catalog_test.cc +++ b/sdk_v2/cpp/test/internal_api/azure_catalog_test.cc @@ -11,9 +11,9 @@ #include "catalog/catalog_client.h" #include "ep_detection/ep_detector.h" #include "exception.h" +#include "internal_api/test_helpers.h" #include "logger.h" #include "model_info.h" -#include "test_helpers.h" #include #include @@ -509,8 +509,7 @@ TEST(AzureCatalogClientTest, WithCachedModels_NoCachedIds_BehavesLikeRegularFetc return MakeOkResponse(MakeMockCatalogResponse({{"phi-4-mini", 3}})); }); - CatalogFetchInfo telemetry_info; - auto result = FetchAllModelInfosWithCachedModels(client, {}, logger, fl::test::TestTelemetrySink(), telemetry_info); + auto result = FetchAllModelInfosWithCachedModels(client, {}, logger, fl::test::TestTelemetrySink()); // Only the primary FetchAllModelInfos call — no extra fetch for cached models. EXPECT_EQ(http_call_count, 1); @@ -530,9 +529,8 @@ TEST(AzureCatalogClientTest, WithCachedModels_AlreadyInCatalog_NoExtraFetch) { }); // The cached ID matches what's already in the catalog — no extra fetch needed. - CatalogFetchInfo telemetry_info; auto result = FetchAllModelInfosWithCachedModels(client, {"phi-4-mini:3"}, logger, - fl::test::TestTelemetrySink(), telemetry_info); + fl::test::TestTelemetrySink()); EXPECT_EQ(http_call_count, 1); ASSERT_EQ(result.size(), 1u); @@ -570,9 +568,8 @@ TEST(AzureCatalogClientTest, WithCachedModels_UnresolvedId_TriggersSecondFetch) } }); - CatalogFetchInfo telemetry_info; auto result = FetchAllModelInfosWithCachedModels(client, {"old-model:1"}, logger, - fl::test::TestTelemetrySink(), telemetry_info); + fl::test::TestTelemetrySink()); EXPECT_EQ(http_call_count, 2); ASSERT_EQ(result.size(), 2u); @@ -612,9 +609,8 @@ TEST(AzureCatalogClientTest, WithCachedModels_FullyUnresolved_CreatesBYOEntry) { } }); - CatalogFetchInfo telemetry_info; auto result = FetchAllModelInfosWithCachedModels(client, {"custom-model:0"}, logger, - fl::test::TestTelemetrySink(), telemetry_info); + fl::test::TestTelemetrySink()); EXPECT_EQ(http_call_count, 2); diff --git a/sdk_v2/cpp/test/internal_api/base_model_catalog_test.cc b/sdk_v2/cpp/test/internal_api/base_model_catalog_test.cc index 4e331d038..204bf3d62 100644 --- a/sdk_v2/cpp/test/internal_api/base_model_catalog_test.cc +++ b/sdk_v2/cpp/test/internal_api/base_model_catalog_test.cc @@ -9,6 +9,7 @@ // - Variant grouping behavior // #include "catalog/base_model_catalog.h" +#include "exception.h" #include "internal_api/test_helpers.h" #include "logger.h" #include "model.h" @@ -108,6 +109,38 @@ class QueryingTestCatalog : public BaseModelCatalog { mutable std::vector id_fetch_results_; }; +class FailingRefreshCatalog : public BaseModelCatalog { + public: + explicit FailingRefreshCatalog(ILogger& logger) : BaseModelCatalog("failing-refresh-catalog", logger) {} + + void SetModels(std::vector models) { + models_ = std::move(models); + } + + void FailNextFetch() { + fail_next_fetch_ = true; + } + + int FetchCount() const { + return fetch_count_; + } + + protected: + std::vector FetchModels() const override { + ++fetch_count_; + if (fail_next_fetch_) { + fail_next_fetch_ = false; + FL_THROW(FOUNDRY_LOCAL_ERROR_NETWORK, "simulated refresh failure"); + } + return std::move(models_); + } + + private: + mutable std::vector models_; + mutable bool fail_next_fetch_ = false; + mutable int fetch_count_ = 0; +}; + // Helper: create a Model from basic fields. static Model MakeModel(const std::string& model_id, const std::string& name, int version, const std::string& alias, @@ -208,6 +241,76 @@ TEST_F(BaseModelCatalogTest, GetModel_VariantsAccessible) { EXPECT_EQ(container->Variants().size(), 2u); } +TEST_F(BaseModelCatalogTest, RefreshMergesNewVariantsForExistingAlias) { + TestCatalog catalog(logger_); + catalog.AddModel(MakeModel("phi-3-mini-cpu:1", "phi-3-mini-cpu", 1, "phi-3")); + + Model* container = catalog.GetModel("phi-3"); + ASSERT_NE(container, nullptr); + EXPECT_EQ(container->Variants().size(), 1u); + EXPECT_EQ(catalog.GetModelVariant("phi-3-mini-gpu:1"), nullptr); + + catalog.AddModel(MakeModel("phi-3-mini-gpu:1", "phi-3-mini-gpu", 1, "phi-3")); + catalog.InvalidateCache(); + + auto refreshed = catalog.ListModels(); + ASSERT_EQ(refreshed.size(), 1u); + EXPECT_EQ(refreshed[0]->Variants().size(), 2u); + EXPECT_NE(catalog.GetModelVariant("phi-3-mini-gpu:1"), nullptr); + EXPECT_EQ(catalog.GetModel("phi-3")->Id(), "phi-3-mini-gpu:1"); +} + +TEST_F(BaseModelCatalogTest, RefreshPreservesExplicitVariantSelection) { + TestCatalog catalog(logger_); + catalog.AddModel(MakeModel("phi-3-mini-cpu:1", "phi-3-mini-cpu", 1, "phi-3")); + + Model* container = catalog.GetModel("phi-3"); + ASSERT_NE(container, nullptr); + Model* cpu_variant = catalog.GetModelVariant("phi-3-mini-cpu:1"); + ASSERT_NE(cpu_variant, nullptr); + container->SelectVariant(*cpu_variant); + + catalog.AddModel(MakeModel("phi-3-mini-gpu:1", "phi-3-mini-gpu", 1, "phi-3")); + catalog.InvalidateCache(); + + auto refreshed = catalog.ListModels(); + ASSERT_EQ(refreshed.size(), 1u); + EXPECT_EQ(refreshed[0]->Variants().size(), 2u); + EXPECT_EQ(container->Id(), "phi-3-mini-cpu:1"); +} + +TEST_F(BaseModelCatalogTest, GetLatestVersionReturnsLatestSameModelName) { + TestCatalog catalog(logger_); + catalog.AddModel(MakeModel("phi-3-mini-cpu:1", "phi-3-mini-cpu", 1, "phi-3")); + catalog.AddModel(MakeModel("phi-3-mini-gpu:2", "phi-3-mini-gpu", 2, "phi-3")); + + Model* cpu_variant = catalog.GetModelVariant("phi-3-mini-cpu:1"); + ASSERT_NE(cpu_variant, nullptr); + + Model* latest = catalog.GetLatestVersion(cpu_variant); + ASSERT_NE(latest, nullptr); + EXPECT_EQ(latest->Info().name, "phi-3-mini-cpu"); +} + +TEST_F(BaseModelCatalogTest, ForcedRefreshFailureBacksOffWhenCatalogAlreadyPopulated) { + FailingRefreshCatalog catalog(logger_); + std::vector models; + models.push_back(MakeModel("phi-3-mini-cpu:1", "phi-3-mini-cpu", 1, "phi-3")); + catalog.SetModels(std::move(models)); + + ASSERT_NE(catalog.GetModel("phi-3"), nullptr); + EXPECT_EQ(catalog.FetchCount(), 1); + + catalog.FailNextFetch(); + catalog.InvalidateCache(); + + EXPECT_EQ(catalog.ListModels().size(), 1u); + EXPECT_EQ(catalog.FetchCount(), 2); + + EXPECT_EQ(catalog.ListModels().size(), 1u); + EXPECT_EQ(catalog.FetchCount(), 2); +} + TEST_F(BaseModelCatalogTest, GetModel_NotFound_ReturnsNullptr) { TestCatalog catalog(logger_); catalog.AddModel(MakeModel("phi-3-mini:1", "phi-3-mini", 1, "phi-3")); diff --git a/sdk_v2/cpp/test/internal_api/blob_download_state_test.cc b/sdk_v2/cpp/test/internal_api/blob_download_state_test.cc index 38dfdd7cf..e149f4ac4 100644 --- a/sdk_v2/cpp/test/internal_api/blob_download_state_test.cc +++ b/sdk_v2/cpp/test/internal_api/blob_download_state_test.cc @@ -252,6 +252,30 @@ TEST(BlobDownloadStateTest, LoadStateRejectsTotalChunksMismatch) { EXPECT_EQ(s, nullptr); } +TEST(BlobDownloadStateTest, LoadStateRejectsBlobIdentityMismatchAndStartsFresh) { + auto d = TempPath::CreateTempDir(); + auto local = d.path() / "blob.bin"; + { + auto s = BlobDownloadState::CreateNew("blob", local, kBlobSize, kChunkSize, kNumChunks, "etag-a"); + s->MarkChunkComplete(0); + s->MarkChunkComplete(3); + s->MarkChunkComplete(7); + ASSERT_TRUE(s->SaveState(fl::test::NullLog())); + } + + auto stale = BlobDownloadState::LoadState("blob", local, kBlobSize, kChunkSize, kNumChunks, "etag-b", + fl::test::NullLog()); + EXPECT_EQ(stale, nullptr); + + auto fresh = BlobDownloadState::CreateNew("blob", local, kBlobSize, kChunkSize, kNumChunks, "etag-b"); + std::vector expected_pending; + for (int32_t i = 0; i < kNumChunks; ++i) { + expected_pending.push_back(i); + } + EXPECT_EQ(fresh->GetPendingChunks(), expected_pending); + EXPECT_EQ(fresh->CalculateDownloadedSize(), 0); +} + TEST(BlobDownloadStateTest, DeleteStateRemovesSidecar) { auto d = TempPath::CreateTempDir(); auto local = d.path() / "blob.bin"; diff --git a/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc b/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc index 780cc186e..e50e6b13e 100644 --- a/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc +++ b/sdk_v2/cpp/test/internal_api/chat/chat_session_test.cc @@ -13,9 +13,9 @@ #include "logger.h" #include "model.h" #include "internal_api/null_session_manager.h" -#include "telemetry/telemetry_logger.h" #include "internal_api/test_helpers.h" #include "internal_api/test_model_cache.h" +#include "telemetry/telemetry_logger.h" #include "utils/string_utils.h" #include diff --git a/sdk_v2/cpp/test/internal_api/download_test.cc b/sdk_v2/cpp/test/internal_api/download_test.cc index 208826333..cd7be2b17 100644 --- a/sdk_v2/cpp/test/internal_api/download_test.cc +++ b/sdk_v2/cpp/test/internal_api/download_test.cc @@ -49,6 +49,10 @@ namespace { using fl::test::TempPath; +ITelemetry& TestTelemetry() { + return fl::test::TestTelemetrySink(); +} + /// Read entire file contents. std::string ReadFile(const fs::path& path) { std::ifstream f(path); @@ -467,6 +471,7 @@ TEST(BlobDownloadTest, FiltersOutInferenceModelJson) { }; BlobDownloadOptions opts; + opts.skip_completed_files = true; DownloadBlobsToDirectoryForTest(mock, "https://test.blob/c?sig=x", tmpdir.string(), opts); ASSERT_EQ(mock.downloaded_blobs.size(), 2u); @@ -530,6 +535,7 @@ TEST(BlobDownloadTest, SkipsExistingFilesWithCorrectSize) { }; BlobDownloadOptions opts; + opts.skip_completed_files = true; DownloadBlobsToDirectoryForTest(mock, "https://test.blob/c?sig=x", tmpdir.string(), opts); // Only the missing blob should be downloaded. @@ -537,6 +543,26 @@ TEST(BlobDownloadTest, SkipsExistingFilesWithCorrectSize) { EXPECT_EQ(mock.downloaded_blobs[0], "config.json"); } +TEST(BlobDownloadTest, RedownloadsSameSizeFilesWhenIdentityRevalidationIsRequired) { + auto tmpdir = TempPath::CreateTempDir(); + std::ofstream(tmpdir.path() / "weights.safetensors") << std::string(1000, 'X'); + + MockBlobDownloader mock; + mock.blobs_to_return = { + {"weights.safetensors", 1000, "etag-b"}, + {"config.json", 100, "etag-config"}, + }; + + BlobDownloadOptions opts; + opts.skip_completed_files = true; + opts.require_completed_file_identity = true; + DownloadBlobsToDirectoryForTest(mock, "https://test.blob/c?sig=x", tmpdir.string(), opts); + + ASSERT_EQ(mock.downloaded_blobs.size(), 2u); + EXPECT_EQ(mock.downloaded_blobs[0], "config.json"); + EXPECT_EQ(mock.downloaded_blobs[1], "weights.safetensors"); +} + TEST(BlobDownloadTest, RedownloadsFilesWithWrongSize) { auto tmpdir = TempPath::CreateTempDir(); // Existing file is truncated relative to the expected blob size. @@ -568,6 +594,7 @@ TEST(BlobDownloadTest, ReportsSkippedBytesInInitialProgress) { std::vector progress_values; BlobDownloadOptions opts; + opts.skip_completed_files = true; opts.progress = [&](float pct) { progress_values.push_back(pct); return 0; @@ -595,6 +622,7 @@ TEST(BlobDownloadTest, EmitsHundredPercentWhenEverythingIsCached) { std::vector progress_values; BlobDownloadOptions opts; + opts.skip_completed_files = true; opts.progress = [&](float pct) { progress_values.push_back(pct); return 0; @@ -855,13 +883,18 @@ TEST(VariantFixupTest, PreservesRootFileWhenNoSubdirs) { TEST(DownloadManagerTest, FullDownloadFlow) { auto tmpdir = TempPath::CreateTempDir(); + auto manager = std::make_unique(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); + + // Mock the registry client auto registry = std::make_unique( "eastus", fl::test::NullLog(), std::make_unique(fl::test::NullLog(), false), [](const std::string&) { return MakeRegistryResponse( R"({"blobSasUri": "https://storage.blob.core.windows.net/container?sig=test"})"); }); + manager->SetModelRegistryClient(std::move(registry)); + // Mock the blob downloader auto mock_downloader = std::make_unique(); mock_downloader->expected_sas_uri = "https://storage.blob.core.windows.net/container?sig=test"; @@ -869,10 +902,7 @@ TEST(DownloadManagerTest, FullDownloadFlow) { {"weights.safetensors", 1024}, {"config.json", 100}, }; - - auto manager = std::make_unique( - tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink(), - /*disable_region_fallback=*/false, std::move(registry), std::move(mock_downloader)); + manager->SetBlobDownloader(std::move(mock_downloader)); ModelInfo info; info.model_id = "test-model:1"; @@ -898,12 +928,164 @@ TEST(DownloadManagerTest, FullDownloadFlow) { EXPECT_FALSE(progress_values.empty()); } +TEST(DownloadManagerTest, RetrySkipsCompletedBlobWithoutSidecar) { + auto tmpdir = TempPath::CreateTempDir(); + + auto manager = std::make_unique(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); + + auto registry = std::make_unique( + "eastus", fl::test::NullLog(), std::make_unique(fl::test::NullLog(), false), + [](const std::string&) { + return MakeRegistryResponse( + R"({"blobSasUri": "https://storage.blob.core.windows.net/container?sig=test"})"); + }); + manager->SetModelRegistryClient(std::move(registry)); + + auto mock_downloader = std::make_unique(); + auto* mock_downloader_ptr = mock_downloader.get(); + mock_downloader->blobs_to_return = { + {"weights.safetensors", 1024}, + {"config.json", 100}, + }; + manager->SetBlobDownloader(std::move(mock_downloader)); + + ModelInfo info; + info.model_id = "test-model:1"; + info.name = "test-model"; + info.uri = "azureml://registries/test/models/test-model/versions/1"; + info.string_properties[FOUNDRY_LOCAL_MODEL_PROP_PUBLISHER_STR] = "TestPublisher"; + + auto model_dir = fs::path(tmpdir.string()) / "TestPublisher" / "test-model-1"; + fs::create_directories(model_dir); + { + std::ofstream signal(model_dir / "download.tmp"); + signal << "foundry-local-sidecar-safe-v1\n"; + } + { + std::ofstream f(model_dir / "weights.safetensors", std::ios::binary); + f.seekp(1023); + f.put('\0'); + } + ASSERT_EQ(fs::file_size(model_dir / "weights.safetensors"), static_cast(1024)); + + manager->DownloadModel(info); + + ASSERT_EQ(mock_downloader_ptr->downloaded_blobs.size(), 1u); + EXPECT_EQ(mock_downloader_ptr->downloaded_blobs[0], "config.json"); +} + +TEST(DownloadManagerTest, LegacyIncompleteRetryRedownloadsFullSizeBlobWithoutSidecar) { + auto tmpdir = TempPath::CreateTempDir(); + + auto manager = std::make_unique(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); + + auto registry = std::make_unique( + "eastus", fl::test::NullLog(), std::make_unique(fl::test::NullLog(), false), + [](const std::string&) { + return MakeRegistryResponse( + R"({"blobSasUri": "https://storage.blob.core.windows.net/container?sig=test"})"); + }); + manager->SetModelRegistryClient(std::move(registry)); + + auto mock_downloader = std::make_unique(); + auto* mock_downloader_ptr = mock_downloader.get(); + mock_downloader->blobs_to_return = { + {"weights.safetensors", 1024}, + {"config.json", 100}, + }; + manager->SetBlobDownloader(std::move(mock_downloader)); + + ModelInfo info; + info.model_id = "test-model:1"; + info.name = "test-model"; + info.uri = "azureml://registries/test/models/test-model/versions/1"; + info.string_properties[FOUNDRY_LOCAL_MODEL_PROP_PUBLISHER_STR] = "TestPublisher"; + + auto model_dir = fs::path(tmpdir.string()) / "TestPublisher" / "test-model-1"; + fs::create_directories(model_dir); + { + std::ofstream signal(model_dir / "download.tmp"); + signal << "legacy-incomplete-download"; + } + { + std::ofstream f(model_dir / "weights.safetensors", std::ios::binary); + f.seekp(1023); + f.put('\0'); + } + + manager->DownloadModel(info); + + ASSERT_EQ(mock_downloader_ptr->downloaded_blobs.size(), 2u); + EXPECT_NE(std::find(mock_downloader_ptr->downloaded_blobs.begin(), mock_downloader_ptr->downloaded_blobs.end(), + "weights.safetensors"), + mock_downloader_ptr->downloaded_blobs.end()); + EXPECT_NE(std::find(mock_downloader_ptr->downloaded_blobs.begin(), mock_downloader_ptr->downloaded_blobs.end(), + "config.json"), + mock_downloader_ptr->downloaded_blobs.end()); +} + +TEST(DownloadManagerTest, FailedLegacyRepairDoesNotPromoteToSidecarSafeRetry) { + auto tmpdir = TempPath::CreateTempDir(); + + int registry_calls = 0; + auto manager = std::make_unique(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); + auto registry = std::make_unique( + "eastus", fl::test::NullLog(), std::make_unique(fl::test::NullLog(), false), + [®istry_calls](const std::string&) { + ++registry_calls; + if (registry_calls == 1) { + return MakeRegistryResponse("temporary failure", 503); + } + return MakeRegistryResponse( + R"({"blobSasUri": "https://storage.blob.core.windows.net/container?sig=test"})"); + }); + manager->SetModelRegistryClient(std::move(registry)); + + auto mock_downloader = std::make_unique(); + auto* mock_downloader_ptr = mock_downloader.get(); + mock_downloader->blobs_to_return = { + {"weights.safetensors", 1024}, + {"config.json", 100}, + }; + manager->SetBlobDownloader(std::move(mock_downloader)); + + ModelInfo info; + info.model_id = "test-model:1"; + info.name = "test-model"; + info.uri = "azureml://registries/test/models/test-model/versions/1"; + info.string_properties[FOUNDRY_LOCAL_MODEL_PROP_PUBLISHER_STR] = "TestPublisher"; + + auto model_dir = fs::path(tmpdir.string()) / "TestPublisher" / "test-model-1"; + fs::create_directories(model_dir); + { + std::ofstream signal(model_dir / "download.tmp"); + signal << "legacy-incomplete-download"; + } + { + std::ofstream f(model_dir / "weights.safetensors", std::ios::binary); + f.seekp(1023); + f.put('\0'); + } + + EXPECT_THROW(manager->DownloadModel(info), fl::Exception); + EXPECT_TRUE(mock_downloader_ptr->downloaded_blobs.empty()); + + manager->DownloadModel(info); + + ASSERT_EQ(mock_downloader_ptr->downloaded_blobs.size(), 2u); + EXPECT_NE(std::find(mock_downloader_ptr->downloaded_blobs.begin(), mock_downloader_ptr->downloaded_blobs.end(), + "weights.safetensors"), + mock_downloader_ptr->downloaded_blobs.end()); +} + // --- Region resolution: detected region drives the download endpoint --- // Run one download and return the registry URL the manager hit. static std::string CaptureRegistryUrlForDownload(const std::string& config_region, const std::string& detected_region) { auto tmpdir = TempPath::CreateTempDir(); + auto manager = + std::make_unique(tmpdir.string(), config_region, 64, fl::test::NullLog(), TestTelemetry()); std::string captured_url; auto registry = std::make_unique( @@ -913,14 +1095,12 @@ static std::string CaptureRegistryUrlForDownload(const std::string& config_regio return MakeRegistryResponse( R"({"blobSasUri": "https://storage.blob.core.windows.net/container?sig=test"})"); }); + manager->SetModelRegistryClient(std::move(registry)); auto mock_downloader = std::make_unique(); mock_downloader->expected_sas_uri = "https://storage.blob.core.windows.net/container?sig=test"; mock_downloader->blobs_to_return = {{"config.json", 100}}; - - auto manager = std::make_unique( - tmpdir.string(), config_region, 64, fl::test::NullLog(), fl::test::TestTelemetrySink(), - /*disable_region_fallback=*/false, std::move(registry), std::move(mock_downloader)); + manager->SetBlobDownloader(std::move(mock_downloader)); ModelInfo info; info.model_id = "test-model:1"; @@ -958,8 +1138,7 @@ TEST(DownloadManagerTest, Region_FallsBackToDefaultRegistryRegionWhenNoConfigAnd TEST(DownloadManagerTest, SkipsAlreadyCachedModel) { auto tmpdir = TempPath::CreateTempDir(); - auto manager = std::make_unique(tmpdir.string(), "eastus", 64, fl::test::NullLog(), - fl::test::TestTelemetrySink()); + auto manager = std::make_unique(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "cached-model:1"; @@ -985,7 +1164,7 @@ TEST(DownloadManagerTest, SkipsAlreadyCachedModel) { TEST(DownloadManagerTest, IsModelCachedReturnsFalseForMissing) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "nonexistent:1"; @@ -996,7 +1175,7 @@ TEST(DownloadManagerTest, IsModelCachedReturnsFalseForMissing) { TEST(DownloadManagerTest, IsModelCachedReturnsFalseForIncomplete) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "incomplete:1"; @@ -1014,7 +1193,7 @@ TEST(DownloadManagerTest, IsModelCachedReturnsFalseForIncomplete) { TEST(DownloadManagerTest, IsModelCachedReturnsTrueForComplete) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "complete:2"; @@ -1033,7 +1212,7 @@ TEST(DownloadManagerTest, IsModelCachedReturnsTrueForComplete) { TEST(DownloadManagerTest, IsModelCachedReturnsFalseForEmptyDir) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "empty:1"; @@ -1049,7 +1228,7 @@ TEST(DownloadManagerTest, IsModelCachedReturnsFalseForEmptyDir) { TEST(DownloadManagerTest, VersionSuffixConversion) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "mymodel:42"; @@ -1069,7 +1248,7 @@ TEST(DownloadManagerTest, VersionSuffixConversion) { TEST(DownloadManagerTest, ThrowsOnEmptyUri) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "test:1"; @@ -1083,12 +1262,14 @@ TEST(DownloadManagerTest, ThrowsOnEmptyUri) { // proceed in parallel — covered by the unrelated-model test below. TEST(DownloadManagerTest, ConcurrentDownloadsOfSameModelSerialize) { auto tmpdir = TempPath::CreateTempDir(); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); auto registry = std::make_unique( "eastus", fl::test::NullLog(), std::make_unique(fl::test::NullLog(), false), [](const std::string&) { return MakeRegistryResponse(R"({"blobSasUri": "https://storage.blob.core.windows.net/c?sig=test"})"); }); + manager.SetModelRegistryClient(std::move(registry)); // Counting mock — increments an atomic on every DownloadBlob call. class CountingDownloader : public IBlobDownloader { @@ -1121,8 +1302,7 @@ TEST(DownloadManagerTest, ConcurrentDownloadsOfSameModelSerialize) { auto counting = std::make_unique(); auto* counting_raw = counting.get(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink(), - /*disable_region_fallback=*/false, std::move(registry), std::move(counting)); + manager.SetBlobDownloader(std::move(counting)); ModelInfo info; info.model_id = "concurrent-model:1"; @@ -1166,6 +1346,7 @@ TEST(DownloadManagerTest, ConcurrentDownloadsOfSameModelSerialize) { // second download can't enter until the first releases the mutex). TEST(DownloadManagerTest, ModelDownloadsSerializeUnderGlobalLock) { auto tmpdir = TempPath::CreateTempDir(); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); auto registry = std::make_unique( "eastus", fl::test::NullLog(), std::make_unique(fl::test::NullLog(), false), @@ -1173,6 +1354,7 @@ TEST(DownloadManagerTest, ModelDownloadsSerializeUnderGlobalLock) { return MakeRegistryResponse( R"({"blobSasUri": "https://storage.blob.core.windows.net/c?sig=test"})"); }); + manager.SetModelRegistryClient(std::move(registry)); // Tracks the peak number of downloads running at once. The global download // mutex must keep this at 1 even for different models. @@ -1211,8 +1393,7 @@ TEST(DownloadManagerTest, ModelDownloadsSerializeUnderGlobalLock) { auto probe = std::make_unique(); auto* probe_raw = probe.get(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink(), - /*disable_region_fallback=*/false, std::move(registry), std::move(probe)); + manager.SetBlobDownloader(std::move(probe)); auto make_info = [](const char* id, const char* publisher) { ModelInfo info; @@ -1257,6 +1438,7 @@ TEST(DownloadManagerTest, ModelDownloadsSerializeUnderGlobalLock) { // recheck WITHOUT re-downloading anything. TEST(DownloadManagerTest, WaitsForCrossProcessLockThenServesCachedResult) { auto tmpdir = TempPath::CreateTempDir(); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); // Registry + downloader that must stay untouched if the post-lock recheck works. auto registry = std::make_unique( @@ -1265,12 +1447,12 @@ TEST(DownloadManagerTest, WaitsForCrossProcessLockThenServesCachedResult) { return MakeRegistryResponse( R"({"blobSasUri": "https://storage.blob.core.windows.net/c?sig=test"})"); }); + manager.SetModelRegistryClient(std::move(registry)); auto mock = std::make_unique(); mock->blobs_to_return = {{"weights.bin", 100}}; // non-empty: a stray download would be visible auto* mock_raw = mock.get(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink(), - /*disable_region_fallback=*/false, std::move(registry), std::move(mock)); + manager.SetBlobDownloader(std::move(mock)); ModelInfo info; info.model_id = "wait-model:1"; @@ -1311,7 +1493,7 @@ TEST(DownloadManagerTest, WaitsForCrossProcessLockThenServesCachedResult) { // underlying directory_iterator would throw filesystem_error. TEST(DownloadManagerTest, IsModelCachedReturnsFalseWhenPathIsRegularFile) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "filemodel:1"; @@ -1379,7 +1561,7 @@ TEST(EndToEndTest, DISABLED_LiveCatalogAndDownload) { // 3. Download the model — use build output dir so reruns skip the download auto cache_path = fs::path(__FILE__).parent_path().parent_path() / "build" / "test_cache"; fs::create_directories(cache_path); - DownloadManager dm(cache_path.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager dm(cache_path.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); std::vector progress_values; std::string local_path = dm.DownloadModel(*smallest, [&](float pct) { @@ -1452,7 +1634,7 @@ TEST(EndToEndTest, DISABLED_LiveCatalogAndDownload) { TEST(DownloadManagerTest, RejectsParentEscapeInModelId) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "../evil:1"; @@ -1464,7 +1646,7 @@ TEST(DownloadManagerTest, RejectsParentEscapeInModelId) { TEST(DownloadManagerTest, RejectsBackslashInPublisher) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "test:1"; @@ -1475,7 +1657,7 @@ TEST(DownloadManagerTest, RejectsBackslashInPublisher) { TEST(DownloadManagerTest, RejectsForwardSlashInPublisher) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "test:1"; @@ -1488,7 +1670,7 @@ TEST(DownloadManagerTest, RejectsColonInBareModelId) { // model_id "drive:c:1" splits as bare="drive:c", version="1"; the bare half then // contains a stray ':' that would let a Windows drive letter slip through. auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "drive:c:1"; @@ -1497,9 +1679,20 @@ TEST(DownloadManagerTest, RejectsColonInBareModelId) { EXPECT_THROW(manager.GetModelCachePath(info), fl::Exception); } +TEST(DownloadManagerTest, RejectsNonNumericVersionSuffix) { + auto tmpdir = TempPath::CreateTempDir(); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); + + ModelInfo info; + info.model_id = "test:preview"; + info.string_properties[FOUNDRY_LOCAL_MODEL_PROP_PUBLISHER_STR] = "Publisher"; + + EXPECT_THROW(manager.GetModelCachePath(info), fl::Exception); +} + TEST(DownloadManagerTest, RejectsTrailingDotInPublisher) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "test:1"; @@ -1510,7 +1703,7 @@ TEST(DownloadManagerTest, RejectsTrailingDotInPublisher) { TEST(DownloadManagerTest, RejectsEmptyModelId) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = ""; @@ -1521,7 +1714,7 @@ TEST(DownloadManagerTest, RejectsEmptyModelId) { TEST(DownloadManagerTest, AcceptsNormalModelIdAndPublisher) { auto tmpdir = TempPath::CreateTempDir(); - DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), fl::test::TestTelemetrySink()); + DownloadManager manager(tmpdir.string(), "eastus", 64, fl::test::NullLog(), TestTelemetry()); ModelInfo info; info.model_id = "phi-3-mini:1"; @@ -1535,7 +1728,7 @@ TEST(DownloadManagerTest, AcceptsNormalModelIdAndPublisher) { // ======================================================================== // AzureBlobDownloader resume + cancel-cascade tests -// Use a subclass that overrides the protected GetBlobSize / DownloadChunkStreaming +// Use a subclass that overrides the protected GetBlobProperties / DownloadChunkStreaming // virtuals to bypass the real Azure SDK and simulate per-chunk behavior. // ======================================================================== @@ -1546,6 +1739,7 @@ namespace { class FakeChunkAzureDownloader : public AzureBlobDownloader { public: int64_t blob_size = 0; + std::string blob_identity; /// Per-call hook. Receives the chunk offset and size plus a `sink` callback /// that forwards bytes to the file writer. Allowed to: @@ -1570,7 +1764,9 @@ class FakeChunkAzureDownloader : public AzureBlobDownloader { FakeChunkAzureDownloader() : AzureBlobDownloader(fl::test::NullLog()) {} protected: - int64_t GetBlobSize(ChunkContext& /*ctx*/) override { return blob_size; } + BlobProperties GetBlobProperties(ChunkContext& /*ctx*/) override { + return BlobProperties{blob_size, blob_identity}; + } void DownloadChunkStreaming(ChunkContext& ctx, int64_t offset, int64_t size, std::vector& scratch, @@ -1801,6 +1997,42 @@ TEST(AzureBlobDownloaderResumeTest, CleansUpSidecarOnEmptyBlob) { EXPECT_EQ(d.chunk_call_count.load(), 0); } +TEST(AzureBlobDownloaderResumeTest, CompleteUnfinalizedSidecarDoesNotEmitRegressingProgress) { + auto tmpdir = TempPath::CreateTempDir(); + auto local = tmpdir.path() / "blob.bin"; + + constexpr int32_t kChunkSize = 2 * 1024 * 1024; + constexpr int32_t kNumChunks = 2; + constexpr int64_t kBlobSize = static_cast(kNumChunks) * kChunkSize; + + { + std::ofstream f(local, std::ios::binary); + f.seekp(kBlobSize - 1); + f.put('\0'); + } + { + auto state = BlobDownloadState::CreateNew("blob", local, kBlobSize, kChunkSize, kNumChunks); + for (int32_t i = 0; i < kNumChunks; ++i) { + state->MarkChunkComplete(i); + } + ASSERT_TRUE(state->SaveState(fl::test::NullLog())); + } + + FakeChunkAzureDownloader d; + d.blob_size = kBlobSize; + std::vector progress; + + d.DownloadBlob(/*sas_uri=*/"", "blob", local.string(), /*max_concurrency=*/2, + [&](int64_t bytes) { progress.push_back(bytes); }); + + ASSERT_FALSE(progress.empty()); + EXPECT_LT(progress.front(), kBlobSize); + for (size_t i = 1; i < progress.size(); ++i) { + EXPECT_GE(progress[i], progress[i - 1]); + } + EXPECT_EQ(progress.back(), kBlobSize); +} + TEST(AzureBlobDownloaderResumeTest, ChunkFailureCancelsInFlightPeersFast) { auto tmpdir = TempPath::CreateTempDir(); auto local = tmpdir.path() / "blob.bin"; diff --git a/sdk_v2/cpp/test/internal_api/ep_detector_test.cc b/sdk_v2/cpp/test/internal_api/ep_detector_test.cc index eb45e64e3..dc609b186 100644 --- a/sdk_v2/cpp/test/internal_api/ep_detector_test.cc +++ b/sdk_v2/cpp/test/internal_api/ep_detector_test.cc @@ -228,10 +228,11 @@ TEST_F(EpDetectorTest, DownloadFiltered_UnknownNamesSkipped) { std::vector names = {"CUDAExecutionProvider", "NonExistentProvider"}; auto result = detector->DownloadAndRegisterEps(&names, nullptr); - EXPECT_TRUE(result.success); + EXPECT_FALSE(result.success); ASSERT_EQ(result.registered_eps.size(), 1u); EXPECT_EQ(result.registered_eps[0], "CUDAExecutionProvider"); - EXPECT_TRUE(result.failed_eps.empty()); + ASSERT_EQ(result.failed_eps.size(), 1u); + EXPECT_EQ(result.failed_eps[0], "NonExistentProvider"); EXPECT_TRUE(mocks[0]->download_called_); } @@ -244,14 +245,12 @@ TEST_F(EpDetectorTest, DownloadFiltered_TelemetryCountsRequestedNamesIncludingUn std::vector names = {"CUDAExecutionProvider", "NonExistentProvider"}; auto result = detector->DownloadAndRegisterEps(&names, nullptr); - EXPECT_TRUE(result.success); + EXPECT_FALSE(result.success); ASSERT_EQ(telemetry.ep_attempt_calls.size(), 1u); EXPECT_EQ(telemetry.ep_attempt_calls[0].num_providers, 2); EXPECT_EQ(telemetry.ep_attempt_calls[0].attempts, 1); EXPECT_EQ(telemetry.ep_attempt_calls[0].succeeded, 1); - ASSERT_EQ(telemetry.ep_register_calls.size(), 1u); - EXPECT_EQ(telemetry.ep_register_calls[0].download_status, ActionStatus::kSuccess); - EXPECT_EQ(telemetry.ep_register_calls[0].register_status, ActionStatus::kSuccess); + EXPECT_EQ(telemetry.ep_attempt_calls[0].failed, 1); } TEST_F(EpDetectorTest, DownloadAll_CancelledProgressRecordsSkippedTelemetry) { @@ -269,26 +268,25 @@ TEST_F(EpDetectorTest, DownloadAll_CancelledProgressRecordsSkippedTelemetry) { EXPECT_TRUE(result.failed_eps.empty()); ASSERT_EQ(telemetry.ep_attempt_calls.size(), 1u); EXPECT_EQ(telemetry.ep_attempt_calls[0].status, ActionStatus::kCanceled); - EXPECT_FALSE(telemetry.ep_attempt_calls[0].user_agent.empty()); EXPECT_EQ(telemetry.ep_attempt_calls[0].attempts, 1); EXPECT_EQ(telemetry.ep_attempt_calls[0].succeeded, 0); EXPECT_EQ(telemetry.ep_attempt_calls[0].failed, 0); ASSERT_EQ(telemetry.ep_register_calls.size(), 1u); - EXPECT_FALSE(telemetry.ep_register_calls[0].user_agent.empty()); EXPECT_EQ(telemetry.ep_register_calls[0].download_status, ActionStatus::kSkipped); EXPECT_EQ(telemetry.ep_register_calls[0].register_status, ActionStatus::kCanceled); } -TEST_F(EpDetectorTest, DownloadFiltered_AllNamesUnknown_SucceedsWithNothing) { +TEST_F(EpDetectorTest, DownloadFiltered_AllNamesUnknownFailsWithRequestedName) { std::vector mocks; auto detector = MakeDetector(mocks, {{"CUDAExecutionProvider", true}}); std::vector names = {"FakeProvider"}; auto result = detector->DownloadAndRegisterEps(&names, nullptr); - EXPECT_TRUE(result.success); + EXPECT_FALSE(result.success); EXPECT_TRUE(result.registered_eps.empty()); - EXPECT_TRUE(result.failed_eps.empty()); + ASSERT_EQ(result.failed_eps.size(), 1u); + EXPECT_EQ(result.failed_eps[0], "FakeProvider"); EXPECT_FALSE(mocks[0]->download_called_); } \ No newline at end of file diff --git a/sdk_v2/cpp/test/internal_api/model_load_manager_test.cc b/sdk_v2/cpp/test/internal_api/model_load_manager_test.cc index 695f0e5d7..9d91a94f0 100644 --- a/sdk_v2/cpp/test/internal_api/model_load_manager_test.cc +++ b/sdk_v2/cpp/test/internal_api/model_load_manager_test.cc @@ -213,6 +213,19 @@ TEST_F(ModelLoadManagerUnloadTest, UnloadThrowsWhenSessionsLive) { instance_->ReleaseSession(); } +TEST_F(ModelLoadManagerUnloadTest, AcquiredLoadedModelLeaseBlocksUnloadUntilReleased) { + auto lease = mgr_->AcquireLoadedModel(fl::test::kTestChatModelAlias); + ASSERT_TRUE(lease); + EXPECT_EQ(lease.get(), instance_); + EXPECT_EQ(instance_->SessionRefCount(), 1); + + EXPECT_THROW(mgr_->UnloadModel(fl::test::kTestChatModelAlias), fl::Exception); + + lease.Reset(); + EXPECT_EQ(instance_->SessionRefCount(), 0); + EXPECT_TRUE(mgr_->UnloadModel(fl::test::kTestChatModelAlias)); +} + TEST_F(ModelLoadManagerUnloadTest, UnloadFailsWhenInUseAndSucceedsAfterSessionsReleased) { instance_->AcquireSession(); EXPECT_THROW(mgr_->UnloadModel(fl::test::kTestChatModelAlias), fl::Exception); diff --git a/sdk_v2/cpp/test/internal_api/session_manager_test.cc b/sdk_v2/cpp/test/internal_api/session_manager_test.cc index f0ae41251..d9119b6ea 100644 --- a/sdk_v2/cpp/test/internal_api/session_manager_test.cc +++ b/sdk_v2/cpp/test/internal_api/session_manager_test.cc @@ -10,9 +10,9 @@ #include "exception.h" #include "logger.h" #include "model.h" -#include "telemetry/telemetry_logger.h" #include "internal_api/test_helpers.h" #include "internal_api/test_model_cache.h" +#include "telemetry/telemetry_logger.h" #include @@ -145,6 +145,31 @@ TEST_F(SessionManagerTest, CheckInAndCheckOutRoundTrip) { EXPECT_EQ(mgr.CacheSize(), 0u); } +TEST(SessionManagerStandaloneTest, CheckOutModelMismatchLeavesCachedSession) { + StderrLogger logger; + SessionManager mgr(logger); + + mgr.CheckIn("resp-1", nullptr, "model-a"); + + EXPECT_EQ(mgr.CheckOut("resp-1", "model-b"), nullptr); + EXPECT_EQ(mgr.CacheSize(), 1u); + + auto checked_out = mgr.CheckOut("resp-1", "model-a"); + EXPECT_EQ(checked_out, nullptr); + EXPECT_EQ(mgr.CacheSize(), 0u); +} + +TEST(SessionManagerStandaloneTest, CheckInAfterCancelAllDropsSession) { + StderrLogger logger; + SessionManager mgr(logger); + mgr.CancelAll(); + + mgr.CheckIn("resp-1", nullptr); + + EXPECT_EQ(mgr.CacheSize(), 0u); + EXPECT_EQ(mgr.CheckOut("resp-1"), nullptr); +} + TEST_F(SessionManagerTest, CheckOutRemovesFromCache) { SessionManager mgr(GetLogger()); auto session = MakeSession(); diff --git a/sdk_v2/cpp/test/internal_api/web_service_test.cc b/sdk_v2/cpp/test/internal_api/web_service_test.cc index 9cceca717..5a5293d05 100644 --- a/sdk_v2/cpp/test/internal_api/web_service_test.cc +++ b/sdk_v2/cpp/test/internal_api/web_service_test.cc @@ -85,7 +85,21 @@ class CapturingTelemetry : public ITelemetry { void RecordEpDownloadAttempt(const EpDownloadAttemptInfo&) override {} void RecordEpDownloadAndRegister(const EpDownloadAndRegisterInfo&) override {} void RecordDownload(const DownloadInfo&) override {} - void RecordCatalogFetch(const CatalogFetchInfo&) override {} + void RecordCatalogFetch(const CatalogFetchInfo& info) override { + std::lock_guard lock(mutex_); + catalog_fetches.push_back(info); + } + + int Count(Action action) { + std::lock_guard lock(mutex_); + int n = 0; + for (const auto& c : actions) { + if (c.action == action) { + ++n; + } + } + return n; + } std::optional Find(Action action) { std::lock_guard lock(mutex_); @@ -96,8 +110,10 @@ class CapturingTelemetry : public ITelemetry { } return std::nullopt; } + std::vector actions; std::vector model_usages; + std::vector catalog_fetches; std::mutex mutex_; }; diff --git a/sdk_v2/cpp/test/test_main.cc b/sdk_v2/cpp/test/test_main.cc new file mode 100644 index 000000000..548082fb2 --- /dev/null +++ b/sdk_v2/cpp/test/test_main.cc @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include + +#ifdef _WIN32 +#include +#else +#include +#endif + +namespace { + +void MarkTestProcessAsRunningUnitTestsForTelemetry() { +#ifdef _WIN32 + ::SetEnvironmentVariableA("ORT_RUNNING_UNIT_TESTS", "1"); +#else + setenv("ORT_RUNNING_UNIT_TESTS", "1", 1); +#endif +} + +} // namespace + +int main(int argc, char** argv) { + MarkTestProcessAsRunningUnitTestsForTelemetry(); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/sdk_v2/js/native/src/catalog.cc b/sdk_v2/js/native/src/catalog.cc index 8372187a2..d8a210b90 100644 --- a/sdk_v2/js/native/src/catalog.cc +++ b/sdk_v2/js/native/src/catalog.cc @@ -5,6 +5,7 @@ #include "addon_data.h" #include "errors.h" #include "model.h" +#include "manager.h" #include @@ -16,10 +17,24 @@ namespace foundry_local_node { namespace { +std::shared_ptr LockManagerOrThrow( + const std::weak_ptr& manager_keepalive, + const std::shared_ptr& lifecycle) { + if (!lifecycle || lifecycle->disposed.load(std::memory_order_acquire)) { + throw foundry_local::Error("Manager has been disposed", FOUNDRY_LOCAL_ERROR_INVALID_USAGE); + } + auto manager = manager_keepalive.lock(); + if (!manager) { + throw foundry_local::Error("Manager has been disposed", FOUNDRY_LOCAL_ERROR_INVALID_USAGE); + } + return manager; +} + // Wrap a ModelList (rvalue) into a JS array of Model handles, each pinning the // passed-in manager reference. -Napi::Value WrapModelList(Napi::Env env, foundry_local::ModelList ml, - Napi::ObjectReference manager) { +Napi::Value WrapModelList(Napi::Env env, foundry_local::ModelList ml, Napi::ObjectReference manager, + std::weak_ptr manager_keepalive, + std::shared_ptr lifecycle) { auto list = std::make_shared(std::move(ml)); auto models = list->Models(); Napi::Array arr = Napi::Array::New(env, models.size()); @@ -28,6 +43,8 @@ Napi::Value WrapModelList(Napi::Env env, foundry_local::ModelList ml, token.impl = models[i].get(); token.keepalive = list; token.manager = Napi::Reference::New(manager.Value(), 1); + token.manager_keepalive = manager_keepalive; + token.lifecycle = lifecycle; arr.Set(static_cast(i), Model::NewInstance(env, std::move(token))); } return arr; @@ -35,7 +52,9 @@ Napi::Value WrapModelList(Napi::Env env, foundry_local::ModelList ml, // Wrap an owning unique_ptr into a JS Model (or undefined when null). Napi::Value WrapOwnedModelOrUndefined(Napi::Env env, std::unique_ptr owned, - Napi::ObjectReference manager) { + Napi::ObjectReference manager, + std::weak_ptr manager_keepalive, + std::shared_ptr lifecycle) { if (!owned) { return env.Undefined(); } @@ -46,11 +65,13 @@ Napi::Value WrapOwnedModelOrUndefined(Napi::Env env, std::unique_ptr>(std::move(owned)); token.keepalive = holder; token.manager = std::move(manager); + token.manager_keepalive = std::move(manager_keepalive); + token.lifecycle = std::move(lifecycle); return Model::NewInstance(env, std::move(token)); } -// Extract IModel* from a JS Model arg, or return nullptr if not a Model. -foundry_local::IModel* ExtractIModel(const Napi::Value& v) { +// Extract Model* from a JS Model arg, or return nullptr if not a Model. +Model* ExtractModel(const Napi::Value& v) { if (!v.IsObject()) { return nullptr; } @@ -63,8 +84,7 @@ foundry_local::IModel* ExtractIModel(const Napi::Value& v) { if (!obj.InstanceOf(ctor)) { return nullptr; } - Model* m = Napi::ObjectWrap::Unwrap(obj); - return m != nullptr ? m->native_impl() : nullptr; + return Napi::ObjectWrap::Unwrap(obj); } Napi::ObjectReference CloneManager(const Napi::ObjectReference& mgr) { @@ -109,11 +129,15 @@ Catalog::Catalog(const Napi::CallbackInfo& info) : Napi::ObjectWrap(inf } impl_ = token->impl; manager_ = std::move(token->manager); + manager_keepalive_ = std::move(token->manager_keepalive); + lifecycle_ = std::move(token->lifecycle); } Napi::Value Catalog::GetName(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; std::string_view name = impl_->GetName(); return Napi::String::New(env, std::string(name)); }); @@ -125,14 +149,20 @@ Napi::Value Catalog::GetModels(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); Napi::ObjectReference mgr = CloneManager(manager_); return CallChecked( - env, [&]() -> Napi::Value { return WrapModelList(env, impl_->GetModels(), std::move(mgr)); }); + env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; + return WrapModelList(env, impl_->GetModels(), std::move(mgr), manager_keepalive_, lifecycle_); + }); } Napi::Value Catalog::GetCachedModels(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); Napi::ObjectReference mgr = CloneManager(manager_); return CallChecked(env, [&]() -> Napi::Value { - return WrapModelList(env, impl_->GetCachedModels(), std::move(mgr)); + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; + return WrapModelList(env, impl_->GetCachedModels(), std::move(mgr), manager_keepalive_, lifecycle_); }); } @@ -140,7 +170,9 @@ Napi::Value Catalog::GetLoadedModels(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); Napi::ObjectReference mgr = CloneManager(manager_); return CallChecked(env, [&]() -> Napi::Value { - return WrapModelList(env, impl_->GetLoadedModels(), std::move(mgr)); + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; + return WrapModelList(env, impl_->GetLoadedModels(), std::move(mgr), manager_keepalive_, lifecycle_); }); } @@ -155,8 +187,10 @@ Napi::Value Catalog::GetModel(const Napi::CallbackInfo& info) { std::string alias = info[0].As(); Napi::ObjectReference mgr = CloneManager(manager_); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; auto owned = impl_->GetModel(alias); - return WrapOwnedModelOrUndefined(env, std::move(owned), std::move(mgr)); + return WrapOwnedModelOrUndefined(env, std::move(owned), std::move(mgr), manager_keepalive_, lifecycle_); }); } @@ -169,8 +203,10 @@ Napi::Value Catalog::GetModelVariant(const Napi::CallbackInfo& info) { std::string model_id = info[0].As(); Napi::ObjectReference mgr = CloneManager(manager_); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; auto owned = impl_->GetModelVariant(model_id); - return WrapOwnedModelOrUndefined(env, std::move(owned), std::move(mgr)); + return WrapOwnedModelOrUndefined(env, std::move(owned), std::move(mgr), manager_keepalive_, lifecycle_); }); } @@ -180,15 +216,20 @@ Napi::Value Catalog::GetLatestVersion(const Napi::CallbackInfo& info) { Napi::TypeError::New(env, "getLatestVersion(model: Model)").ThrowAsJavaScriptException(); return env.Undefined(); } - foundry_local::IModel* arg = ExtractIModel(info[0]); - if (arg == nullptr) { + Model* arg = ExtractModel(info[0]); + if (arg == nullptr || arg->native_impl() == nullptr) { Napi::TypeError::New(env, "getLatestVersion: argument must be a Model").ThrowAsJavaScriptException(); return env.Undefined(); } Napi::ObjectReference mgr = CloneManager(manager_); return CallChecked(env, [&]() -> Napi::Value { - auto owned = impl_->GetLatestVersion(*arg); - return WrapOwnedModelOrUndefined(env, std::move(owned), std::move(mgr)); + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + if (arg->manager_disposed() || !arg->manager_keepalive()) { + throw foundry_local::Error("Manager has been disposed", FOUNDRY_LOCAL_ERROR_INVALID_USAGE); + } + (void)manager_alive; + auto owned = impl_->GetLatestVersion(*arg->native_impl()); + return WrapOwnedModelOrUndefined(env, std::move(owned), std::move(mgr), manager_keepalive_, lifecycle_); }); } diff --git a/sdk_v2/js/native/src/catalog.h b/sdk_v2/js/native/src/catalog.h index 29ff89cd9..6d8ce8fa3 100644 --- a/sdk_v2/js/native/src/catalog.h +++ b/sdk_v2/js/native/src/catalog.h @@ -16,13 +16,18 @@ #include +#include #include namespace foundry_local_node { +struct ManagerLifecycle; + struct CatalogCtorToken { foundry_local::ICatalog* impl = nullptr; Napi::ObjectReference manager; // pins the owning Manager + std::weak_ptr manager_keepalive; + std::shared_ptr lifecycle; }; class Catalog : public Napi::ObjectWrap { @@ -43,6 +48,8 @@ class Catalog : public Napi::ObjectWrap { foundry_local::ICatalog* impl_ = nullptr; Napi::ObjectReference manager_; + std::weak_ptr manager_keepalive_; + std::shared_ptr lifecycle_; }; } // namespace foundry_local_node diff --git a/sdk_v2/js/native/src/manager.cc b/sdk_v2/js/native/src/manager.cc index 319974e80..28eaeda6c 100644 --- a/sdk_v2/js/native/src/manager.cc +++ b/sdk_v2/js/native/src/manager.cc @@ -20,6 +20,24 @@ namespace foundry_local_node { namespace { +struct WorkerLease { + explicit WorkerLease(std::shared_ptr lifecycle) : lifecycle_(std::move(lifecycle)) { + if (lifecycle_) { + lifecycle_->active_workers.fetch_add(1, std::memory_order_acq_rel); + } + } + ~WorkerLease() { + if (lifecycle_) { + lifecycle_->active_workers.fetch_sub(1, std::memory_order_acq_rel); + } + } + std::shared_ptr lifecycle_; +}; + +std::shared_ptr MakeWorkerLease(std::shared_ptr lifecycle) { + return std::make_shared(std::move(lifecycle)); +} + Napi::Value ConvertEndpoints(Napi::Env env, std::vector& endpoints) { Napi::Array out = Napi::Array::New(env, endpoints.size()); for (size_t i = 0; i < endpoints.size(); ++i) { @@ -213,16 +231,14 @@ Manager::Manager(const Napi::CallbackInfo& info) : Napi::ObjectWrap(inf for (const auto& entry : additional_settings) { kvp.Set(entry.first.c_str(), entry.second.c_str()); } - if (has_disable_nonessential_telemetry) { - kvp.Set("DisableNonessentialTelemetry", disable_nonessential_telemetry ? "true" : "false"); - } config.SetAdditionalOptions(kvp); - } else if (has_disable_nonessential_telemetry) { + } + if (has_disable_nonessential_telemetry) { foundry_local::KeyValuePairs kvp; kvp.Set("DisableNonessentialTelemetry", disable_nonessential_telemetry ? "true" : "false"); config.SetAdditionalOptions(kvp); } - impl_ = std::make_unique(std::move(config)); + impl_ = std::make_shared(std::move(config)); }); } @@ -256,6 +272,8 @@ Napi::Value Manager::GetCatalog(const Napi::CallbackInfo& info) { CatalogCtorToken token; token.impl = &cat; token.manager = std::move(owner); + token.manager_keepalive = impl_; + token.lifecycle = lifecycle_; return Catalog::NewInstance(env, std::move(token)); }); } @@ -263,6 +281,17 @@ Napi::Value Manager::GetCatalog(const Napi::CallbackInfo& info) { Napi::Value Manager::Dispose(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); // Idempotent — releasing an already-null unique_ptr is a no-op. + if (lifecycle_->active_sessions.load(std::memory_order_acquire) > 0) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, + "Manager has active sessions; dispose sessions before disposing the manager"); + return env.Undefined(); + } + if (lifecycle_->active_workers.load(std::memory_order_acquire) > 0) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, + "Manager has active native workers; await them before disposing the manager"); + return env.Undefined(); + } + lifecycle_->disposed.store(true, std::memory_order_release); impl_.reset(); return env.Undefined(); } @@ -317,11 +346,13 @@ namespace { // progress callback. Mirrors the pattern in model.cc's DownloadWorker. class EpDownloadWorker : public Napi::AsyncWorker { public: - EpDownloadWorker(Napi::Env env, foundry_local::Manager* impl, std::vector ep_names, - Napi::ObjectReference owner, Napi::ThreadSafeFunction tsfn) + EpDownloadWorker(Napi::Env env, std::shared_ptr impl, std::shared_ptr worker_lease, + std::vector ep_names, Napi::ObjectReference owner, + Napi::ThreadSafeFunction tsfn) : Napi::AsyncWorker(env), deferred_(Napi::Promise::Deferred::New(env)), impl_(impl), + worker_lease_(std::move(worker_lease)), ep_names_(std::move(ep_names)), owner_(std::move(owner)), tsfn_(std::move(tsfn)) {} @@ -385,7 +416,8 @@ class EpDownloadWorker : public Napi::AsyncWorker { } Napi::Promise::Deferred deferred_; - foundry_local::Manager* impl_; + std::shared_ptr impl_; + std::shared_ptr worker_lease_; std::vector ep_names_; Napi::ObjectReference owner_; Napi::ThreadSafeFunction tsfn_; @@ -442,7 +474,8 @@ Napi::Value Manager::DownloadAndRegisterEps(const Napi::CallbackInfo& info) { } Napi::ObjectReference owner = Napi::Reference::New(info.This().As(), 1); - auto* w = new EpDownloadWorker(env, impl_.get(), std::move(ep_names), std::move(owner), std::move(tsfn)); + auto* w = new EpDownloadWorker(env, impl_, MakeWorkerLease(lifecycle_), std::move(ep_names), std::move(owner), + std::move(tsfn)); Napi::Promise p = w->Promise(); w->Queue(); return p; diff --git a/sdk_v2/js/native/src/manager.h b/sdk_v2/js/native/src/manager.h index 0cdd35207..c321a1c12 100644 --- a/sdk_v2/js/native/src/manager.h +++ b/sdk_v2/js/native/src/manager.h @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. // -// Napi::ObjectWrap over std::unique_ptr. +// Napi::ObjectWrap over std::shared_ptr. // // Surface: // - ctor accepts { appName, modelCacheDir?, serviceEndpoint? } @@ -15,10 +15,17 @@ #include +#include #include namespace foundry_local_node { +struct ManagerLifecycle { + std::atomic disposed{false}; + std::atomic active_sessions{0}; + std::atomic active_workers{0}; +}; + class Manager : public Napi::ObjectWrap { public: static Napi::Function Init(Napi::Env env); @@ -55,7 +62,8 @@ class Manager : public Napi::ObjectWrap { // on env and returns true. Callers should return env.Undefined() when true. bool ThrowIfDisposed(Napi::Env env); - std::unique_ptr impl_; + std::shared_ptr impl_; + std::shared_ptr lifecycle_ = std::make_shared(); }; } // namespace foundry_local_node diff --git a/sdk_v2/js/native/src/model.cc b/sdk_v2/js/native/src/model.cc index c7164778f..a6605acdc 100644 --- a/sdk_v2/js/native/src/model.cc +++ b/sdk_v2/js/native/src/model.cc @@ -4,6 +4,7 @@ #include "addon_data.h" #include "errors.h" +#include "manager.h" #include "promise_worker.h" #include @@ -18,6 +19,24 @@ namespace foundry_local_node { namespace { +struct WorkerLease { + explicit WorkerLease(std::shared_ptr lifecycle) : lifecycle_(std::move(lifecycle)) { + if (lifecycle_) { + lifecycle_->active_workers.fetch_add(1, std::memory_order_acq_rel); + } + } + ~WorkerLease() { + if (lifecycle_) { + lifecycle_->active_workers.fetch_sub(1, std::memory_order_acq_rel); + } + } + std::shared_ptr lifecycle_; +}; + +std::shared_ptr MakeWorkerLease(std::shared_ptr lifecycle) { + return std::make_shared(std::move(lifecycle)); +} + const char* DeviceTypeToString(flDeviceType dt) { switch (dt) { case FOUNDRY_LOCAL_DEVICE_CPU: @@ -73,6 +92,46 @@ void SetPromptTemplate(Napi::Env env, Napi::Object obj, const foundry_local::Mod obj.Set("promptTemplate", template_obj); } +std::shared_ptr LockManager( + const std::weak_ptr& manager_keepalive) { + return manager_keepalive.lock(); +} + +void ThrowFoundryLocalError(Napi::Env env, int code, const std::string& msg) { + Napi::Error err = Napi::Error::New(env, msg); + Napi::Object value = err.Value(); + value.Set("name", Napi::String::New(env, "FoundryLocalError")); + value.Set("code", Napi::Number::New(env, code)); + err.ThrowAsJavaScriptException(); +} + +std::shared_ptr LockManagerOrThrow( + const std::weak_ptr& manager_keepalive, + const std::shared_ptr& lifecycle) { + if (!lifecycle || lifecycle->disposed.load(std::memory_order_acquire)) { + throw foundry_local::Error("Manager has been disposed", FOUNDRY_LOCAL_ERROR_INVALID_USAGE); + } + auto manager = LockManager(manager_keepalive); + if (!manager) { + throw foundry_local::Error("Manager has been disposed", FOUNDRY_LOCAL_ERROR_INVALID_USAGE); + } + return manager; +} + +std::shared_ptr LockManagerOrThrowJs( + Napi::Env env, const std::weak_ptr& manager_keepalive, + const std::shared_ptr& lifecycle) { + if (!lifecycle || lifecycle->disposed.load(std::memory_order_acquire)) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Manager has been disposed"); + return nullptr; + } + auto manager = LockManager(manager_keepalive); + if (!manager) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Manager has been disposed"); + } + return manager; +} + void SetModelSettings(Napi::Env env, Napi::Object obj, const foundry_local::ModelInfo& info) { auto settings = info.GetModelSettings(); if (!settings.has_value()) { @@ -144,7 +203,9 @@ Napi::Object SnapshotModelInfo(Napi::Env env, const foundry_local::ModelInfo& in // Drain a ModelList into a JS array, with each entry wrapped as a JS Model // whose keepalive holds the shared ModelList. Napi::Array WrapModelList(Napi::Env env, std::shared_ptr list, - Napi::ObjectReference manager) { + Napi::ObjectReference manager, + std::weak_ptr manager_keepalive, + std::shared_ptr lifecycle) { auto models = list->Models(); Napi::Array arr = Napi::Array::New(env, models.size()); for (size_t i = 0; i < models.size(); ++i) { @@ -153,6 +214,8 @@ Napi::Array WrapModelList(Napi::Env env, std::shared_ptr::New(manager.Value(), 1); + token.manager_keepalive = manager_keepalive; + token.lifecycle = lifecycle; arr.Set(static_cast(i), Model::NewInstance(env, std::move(token))); } return arr; @@ -198,12 +261,20 @@ Model::Model(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info) { } impl_ = token->impl; keepalive_ = std::move(token->keepalive); + manager_keepalive_ = std::move(token->manager_keepalive); + lifecycle_ = std::move(token->lifecycle); manager_ = std::move(token->manager); } +bool Model::manager_disposed() const noexcept { + return !lifecycle_ || lifecycle_->disposed.load(std::memory_order_acquire); +} + Napi::Value Model::GetInfo(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; foundry_local::ModelInfo mi = impl_->GetInfo(); Napi::Object snapshot = SnapshotModelInfo(env, mi); snapshot.Set("cached", Napi::Boolean::New(env, impl_->IsCached())); @@ -214,6 +285,8 @@ Napi::Value Model::GetInfo(const Napi::CallbackInfo& info) { Napi::Value Model::IsCached(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; return Napi::Boolean::New(env, impl_->IsCached()); }); } @@ -221,6 +294,8 @@ Napi::Value Model::IsCached(const Napi::CallbackInfo& info) { Napi::Value Model::IsLoaded(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; return Napi::Boolean::New(env, impl_->IsLoaded()); }); } @@ -228,6 +303,8 @@ Napi::Value Model::IsLoaded(const Napi::CallbackInfo& info) { Napi::Value Model::GetPath(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; std::string_view p = impl_->GetPath(); return Napi::String::New(env, std::string(p)); }); @@ -237,18 +314,19 @@ Napi::Value Model::GetVariants(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); Napi::ObjectReference owner_clone = Napi::Reference::New(manager_.Value(), 1); return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; auto list = std::make_shared(impl_->GetVariants()); - return WrapModelList(env, std::move(list), std::move(owner_clone)); + return WrapModelList(env, std::move(list), std::move(owner_clone), manager_keepalive_, lifecycle_); }); } // ── Async lifecycle ───────────────────────────────────────────────────────── // // Load/Unload/Download dispatch the underlying virtual call onto a libuv -// worker so the event loop stays responsive. The Model itself is pinned -// against GC for the duration of the worker via an ObjectReference to the -// parent Manager (the Manager owns the catalog whose ModelList views the -// IModel*). +// worker so the event loop stays responsive. Each worker captures the model's +// native keepalives so explicit Manager.dispose() or JS GC cannot release the +// underlying Manager/ModelList/owned IModel before the worker finishes. Napi::Value Model::Load(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); @@ -258,8 +336,16 @@ Napi::Value Model::Load(const Napi::CallbackInfo& info) { } Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); foundry_local::IModel* m = impl_; + auto keepalive = keepalive_; + auto manager_keepalive = LockManagerOrThrowJs(env, manager_keepalive_, lifecycle_); + if (!manager_keepalive) { + return env.Undefined(); + } + auto worker_lease = MakeWorkerLease(lifecycle_); return PromiseWorkerVoid::Run( - env, [m]() { m->Load(); }, std::move(owner)); + env, [m, keepalive = std::move(keepalive), manager_keepalive = std::move(manager_keepalive), + worker_lease = std::move(worker_lease)]() { m->Load(); }, + std::move(owner)); } Napi::Value Model::Unload(const Napi::CallbackInfo& info) { @@ -270,8 +356,16 @@ Napi::Value Model::Unload(const Napi::CallbackInfo& info) { } Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); foundry_local::IModel* m = impl_; + auto keepalive = keepalive_; + auto manager_keepalive = LockManagerOrThrowJs(env, manager_keepalive_, lifecycle_); + if (!manager_keepalive) { + return env.Undefined(); + } + auto worker_lease = MakeWorkerLease(lifecycle_); return PromiseWorkerVoid::Run( - env, [m]() { m->Unload(); }, std::move(owner)); + env, [m, keepalive = std::move(keepalive), manager_keepalive = std::move(manager_keepalive), + worker_lease = std::move(worker_lease)]() { m->Unload(); }, + std::move(owner)); } namespace { @@ -282,11 +376,15 @@ namespace { // worker queues and released in OnOK/OnError. class DownloadWorker : public Napi::AsyncWorker { public: - DownloadWorker(Napi::Env env, foundry_local::IModel* impl, Napi::ObjectReference owner, - Napi::ThreadSafeFunction tsfn) + DownloadWorker(Napi::Env env, foundry_local::IModel* impl, std::shared_ptr keepalive, + std::shared_ptr manager_keepalive, std::shared_ptr worker_lease, + Napi::ObjectReference owner, Napi::ThreadSafeFunction tsfn) : Napi::AsyncWorker(env), deferred_(Napi::Promise::Deferred::New(env)), impl_(impl), + keepalive_(std::move(keepalive)), + manager_keepalive_(std::move(manager_keepalive)), + worker_lease_(std::move(worker_lease)), owner_(std::move(owner)), tsfn_(std::move(tsfn)) {} @@ -350,6 +448,9 @@ class DownloadWorker : public Napi::AsyncWorker { Napi::Promise::Deferred deferred_; foundry_local::IModel* impl_; + std::shared_ptr keepalive_; + std::shared_ptr manager_keepalive_; + std::shared_ptr worker_lease_; Napi::ObjectReference owner_; Napi::ThreadSafeFunction tsfn_; std::string err_msg_; @@ -366,6 +467,11 @@ Napi::Value Model::Download(const Napi::CallbackInfo& info) { return env.Undefined(); } + auto manager_keepalive = LockManagerOrThrowJs(env, manager_keepalive_, lifecycle_); + if (!manager_keepalive) { + return env.Undefined(); + } + Napi::ThreadSafeFunction tsfn; if (info.Length() >= 1 && info[0].IsFunction()) { tsfn = Napi::ThreadSafeFunction::New(env, info[0].As(), @@ -379,7 +485,9 @@ Napi::Value Model::Download(const Napi::CallbackInfo& info) { } Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); - auto* w = new DownloadWorker(env, impl_, std::move(owner), std::move(tsfn)); + auto* w = new DownloadWorker(env, impl_, keepalive_, std::move(manager_keepalive), MakeWorkerLease(lifecycle_), + std::move(owner), + std::move(tsfn)); Napi::Promise p = w->Promise(); w->Queue(); return p; @@ -395,6 +503,8 @@ Napi::Value Model::RemoveFromCache(const Napi::CallbackInfo& info) { // V1's contract is `removeFromCache(): void` so we do not bounce to a // worker. return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + (void)manager_alive; impl_->RemoveFromCache(); return env.Undefined(); }); @@ -422,6 +532,10 @@ Napi::Value Model::SelectVariant(const Napi::CallbackInfo& info) { } return CallChecked(env, [&]() -> Napi::Value { + auto manager_alive = LockManagerOrThrow(manager_keepalive_, lifecycle_); + auto variant_manager_alive = LockManagerOrThrow(variant->manager_keepalive_, variant->lifecycle_); + (void)manager_alive; + (void)variant_manager_alive; impl_->SelectVariant(*variant->impl_); return env.Undefined(); }); diff --git a/sdk_v2/js/native/src/model.h b/sdk_v2/js/native/src/model.h index 3ae4b6e63..e90dc3977 100644 --- a/sdk_v2/js/native/src/model.h +++ b/sdk_v2/js/native/src/model.h @@ -32,6 +32,8 @@ namespace foundry_local_node { +struct ManagerLifecycle; + struct ModelCtorToken { // The IModel accessor. Never null when the token is constructed. foundry_local::IModel* impl = nullptr; @@ -39,6 +41,8 @@ struct ModelCtorToken { // or a std::shared_ptr) alive for the JS Model's // lifetime. std::shared_ptr keepalive; + std::weak_ptr manager_keepalive; + std::shared_ptr lifecycle; // Pins the parent Manager so its native handle (and the Catalog's flCatalog* // which the IModel views into) cannot be released first. Napi::ObjectReference manager; @@ -62,6 +66,9 @@ class Model : public Napi::ObjectWrap { // Internal accessor used by Session / ChatSession ctors so they can clone // the parent Manager ObjectReference and pin it for the session lifetime. const Napi::ObjectReference& manager() const noexcept { return manager_; } + std::shared_ptr manager_keepalive() const noexcept { return manager_keepalive_.lock(); } + std::shared_ptr manager_lifecycle() const noexcept { return lifecycle_; } + bool manager_disposed() const noexcept; private: Napi::Value GetInfo(const Napi::CallbackInfo& info); @@ -78,6 +85,8 @@ class Model : public Napi::ObjectWrap { foundry_local::IModel* impl_ = nullptr; std::shared_ptr keepalive_; + std::weak_ptr manager_keepalive_; + std::shared_ptr lifecycle_; Napi::ObjectReference manager_; }; diff --git a/sdk_v2/js/native/src/session.cc b/sdk_v2/js/native/src/session.cc index 0d1247463..d71f2896e 100644 --- a/sdk_v2/js/native/src/session.cc +++ b/sdk_v2/js/native/src/session.cc @@ -5,6 +5,7 @@ #include "addon_data.h" #include "errors.h" #include "items.h" +#include "manager.h" #include "model.h" #include "promise_worker.h" #include "request.h" @@ -22,6 +23,36 @@ namespace foundry_local_node { namespace { +struct SessionLease { + explicit SessionLease(std::shared_ptr lifecycle) + : lifecycle_(std::move(lifecycle)) { + if (lifecycle_) { + lifecycle_->active_sessions.fetch_add(1, std::memory_order_acq_rel); + } + } + + ~SessionLease() { + if (lifecycle_) { + lifecycle_->active_sessions.fetch_sub(1, std::memory_order_acq_rel); + } + } + + std::shared_ptr lifecycle_; +}; + +std::shared_ptr MakeSessionLease(std::shared_ptr lifecycle) { + return std::make_shared(std::move(lifecycle)); +} + +template +void ReleaseSessionState(std::shared_ptr& impl, + std::shared_ptr& manager_keepalive, + std::shared_ptr& session_lease) { + impl.reset(); + manager_keepalive.reset(); + session_lease.reset(); +} + const char* FinishReasonToString(flFinishReason r) { switch (r) { case FOUNDRY_LOCAL_FINISH_STOP: @@ -97,8 +128,10 @@ foundry_local::Request* UnwrapRequest(Napi::Env env, const Napi::Value& v) { // Pins both the Manager (so the Model handle the Session holds stays alive) // and the Request (so the C++ Request the worker reads stays alive). template -Napi::Value ProcessRequestOn(Napi::Env env, SessT* sess, const Napi::Value& request_arg, - Napi::ObjectReference manager_ref) { +Napi::Value ProcessRequestOn(Napi::Env env, std::shared_ptr sess, const Napi::Value& request_arg, + Napi::ObjectReference manager_ref, + std::shared_ptr manager_keepalive, + std::shared_ptr session_lease) { foundry_local::Request* req = UnwrapRequest(env, request_arg); if (req == nullptr) return env.Undefined(); // pending exception Napi::ObjectReference req_pin = Napi::Reference::New(request_arg.As(), 1); @@ -107,12 +140,15 @@ Napi::Value ProcessRequestOn(Napi::Env env, SessT* sess, const Napi::Value& requ struct Pins { Napi::ObjectReference manager; Napi::ObjectReference request; + std::shared_ptr native_manager; + std::shared_ptr session_lease; }; - auto pins = std::make_shared(Pins{std::move(manager_ref), std::move(req_pin)}); + auto pins = std::make_shared( + Pins{std::move(manager_ref), std::move(req_pin), std::move(manager_keepalive), std::move(session_lease)}); return PromiseWorker::Run( env, - [sess, req, pins]() -> Result { + [sess = std::move(sess), req, pins]() -> Result { (void)pins; // keepalive captured by reference count return std::make_shared(sess->ProcessRequest(*req)); }, @@ -148,6 +184,8 @@ struct StreamCtx { Napi::Promise::Deferred deferred; Napi::ObjectReference manager; Napi::ObjectReference request; + std::shared_ptr native_manager; + std::shared_ptr session_lease; std::shared_ptr response; std::string err_msg; int err_code = 0; @@ -180,9 +218,9 @@ void FinalizeStream(Napi::Env env, void* /*data*/, StreamCtx* ctx) { template class StreamWorker : public Napi::AsyncWorker { public: - static Napi::Promise Run(Napi::Env env, SessT* sess, foundry_local::Request* req, + static Napi::Promise Run(Napi::Env env, std::shared_ptr sess, foundry_local::Request* req, Napi::Function jsCallback, StreamCtx* ctx) { - auto* w = new StreamWorker(env, sess, req, jsCallback, ctx); + auto* w = new StreamWorker(env, std::move(sess), req, jsCallback, ctx); Napi::Promise p = ctx->deferred.Promise(); w->Queue(); return p; @@ -215,9 +253,6 @@ class StreamWorker : public Napi::AsyncWorker { return 0; }); ctx_->response = std::make_shared(sess_->ProcessRequest(*req_)); - // Drop the callback so any stale shared state in the lambda is released - // before the Session is re-used for a follow-up request. - sess_->SetStreamingCallback(nullptr); } catch (const foundry_local::Error& e) { ctx_->errored = true; ctx_->err_code = static_cast(e.Code()); @@ -230,6 +265,9 @@ class StreamWorker : public Napi::AsyncWorker { ctx_->errored = true; ctx_->err_msg = "Unknown native exception"; } + // Drop the callback so any stale shared state in the lambda is released + // before the Session is re-used for a follow-up request. + sess_->SetStreamingCallback(nullptr); } // Promise resolution happens in FinalizeStream — overriding OnOK/OnError @@ -239,10 +277,10 @@ class StreamWorker : public Napi::AsyncWorker { void OnError(const Napi::Error& /*unused*/) override { tsfn_.Release(); } private: - StreamWorker(Napi::Env env, SessT* sess, foundry_local::Request* req, + StreamWorker(Napi::Env env, std::shared_ptr sess, foundry_local::Request* req, Napi::Function jsCallback, StreamCtx* ctx) : Napi::AsyncWorker(env), - sess_(sess), + sess_(std::move(sess)), req_(req), ctx_(ctx), tsfn_(Napi::ThreadSafeFunction::New(env, jsCallback, "foundry_local_stream", @@ -250,15 +288,17 @@ class StreamWorker : public Napi::AsyncWorker { FinalizeStream, static_cast(nullptr))) {} - SessT* sess_; + std::shared_ptr sess_; foundry_local::Request* req_; StreamCtx* ctx_; Napi::ThreadSafeFunction tsfn_; }; template -Napi::Value ProcessStreamingRequestOn(Napi::Env env, SessT* sess, const Napi::CallbackInfo& info, - Napi::ObjectReference manager_ref) { +Napi::Value ProcessStreamingRequestOn(Napi::Env env, std::shared_ptr sess, const Napi::CallbackInfo& info, + Napi::ObjectReference manager_ref, + std::shared_ptr manager_keepalive, + std::shared_ptr session_lease) { if (info.Length() < 2 || !info[1].IsFunction()) { Napi::TypeError::New(env, "processStreamingRequest(request: Request, onItem: (item) => void)") .ThrowAsJavaScriptException(); @@ -272,12 +312,14 @@ Napi::Value ProcessStreamingRequestOn(Napi::Env env, SessT* sess, const Napi::Ca auto* ctx = new StreamCtx{Napi::Promise::Deferred::New(env), std::move(manager_ref), std::move(req_pin), + std::move(manager_keepalive), + std::move(session_lease), nullptr, "", 0, false, false}; - return StreamWorker::Run(env, sess, req, info[1].As(), ctx); + return StreamWorker::Run(env, std::move(sess), req, info[1].As(), ctx); } } // namespace @@ -322,7 +364,15 @@ ChatSession::ChatSession(const Napi::CallbackInfo& info) : Napi::ObjectWrap(*native); + auto manager_keepalive = model->manager_keepalive(); + if (model->manager_disposed() || !manager_keepalive) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Manager has been disposed"); + return; + } + impl_ = std::make_shared(*native); + lifecycle_ = model->manager_lifecycle(); + session_lease_ = MakeSessionLease(lifecycle_); + manager_keepalive_ = std::move(manager_keepalive); } catch (const foundry_local::Error& e) { ThrowFoundryLocalError(env, static_cast(e.Code()), e.what()); return; @@ -333,6 +383,10 @@ ChatSession::ChatSession(const Napi::CallbackInfo& info) : Napi::ObjectWrap::New(model->manager().Value(), 1); } +ChatSession::~ChatSession() { + ReleaseSessionState(impl_, manager_keepalive_, session_lease_); +} + bool ChatSession::ThrowIfDisposed(Napi::Env env) { if (impl_ == nullptr) { ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, @@ -350,14 +404,14 @@ Napi::Value ChatSession::ProcessRequest(const Napi::CallbackInfo& info) { return env.Undefined(); } Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); - return ProcessRequestOn(env, impl_.get(), info[0], std::move(owner)); + return ProcessRequestOn(env, impl_, info[0], std::move(owner), manager_keepalive_, session_lease_); } Napi::Value ChatSession::ProcessStreamingRequest(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (ThrowIfDisposed(env)) return env.Undefined(); Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); - return ProcessStreamingRequestOn(env, impl_.get(), info, std::move(owner)); + return ProcessStreamingRequestOn(env, impl_, info, std::move(owner), manager_keepalive_, session_lease_); } Napi::Value ChatSession::SetOptions(const Napi::CallbackInfo& info) { @@ -437,7 +491,7 @@ Napi::Value ChatSession::UndoTurns(const Napi::CallbackInfo& info) { } Napi::Value ChatSession::Dispose(const Napi::CallbackInfo& info) { - impl_.reset(); + ReleaseSessionState(impl_, manager_keepalive_, session_lease_); return info.Env().Undefined(); } @@ -479,7 +533,15 @@ EmbeddingsSession::EmbeddingsSession(const Napi::CallbackInfo& info) return; } try { - impl_ = std::make_unique(*native); + auto manager_keepalive = model->manager_keepalive(); + if (model->manager_disposed() || !manager_keepalive) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Manager has been disposed"); + return; + } + impl_ = std::make_shared(*native); + lifecycle_ = model->manager_lifecycle(); + session_lease_ = MakeSessionLease(lifecycle_); + manager_keepalive_ = std::move(manager_keepalive); } catch (const foundry_local::Error& e) { ThrowFoundryLocalError(env, static_cast(e.Code()), e.what()); return; @@ -490,6 +552,10 @@ EmbeddingsSession::EmbeddingsSession(const Napi::CallbackInfo& info) manager_ = Napi::Reference::New(model->manager().Value(), 1); } +EmbeddingsSession::~EmbeddingsSession() { + ReleaseSessionState(impl_, manager_keepalive_, session_lease_); +} + bool EmbeddingsSession::ThrowIfDisposed(Napi::Env env) { if (impl_ == nullptr) { ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, @@ -507,7 +573,7 @@ Napi::Value EmbeddingsSession::ProcessRequest(const Napi::CallbackInfo& info) { return env.Undefined(); } Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); - return ProcessRequestOn(env, impl_.get(), info[0], std::move(owner)); + return ProcessRequestOn(env, impl_, info[0], std::move(owner), manager_keepalive_, session_lease_); } Napi::Value EmbeddingsSession::SetOptions(const Napi::CallbackInfo& info) { @@ -526,7 +592,7 @@ Napi::Value EmbeddingsSession::SetOptions(const Napi::CallbackInfo& info) { } Napi::Value EmbeddingsSession::Dispose(const Napi::CallbackInfo& info) { - impl_.reset(); + ReleaseSessionState(impl_, manager_keepalive_, session_lease_); return info.Env().Undefined(); } @@ -573,7 +639,15 @@ AudioSession::AudioSession(const Napi::CallbackInfo& info) return; } try { - impl_ = std::make_unique(*native); + auto manager_keepalive = model->manager_keepalive(); + if (model->manager_disposed() || !manager_keepalive) { + ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, "Manager has been disposed"); + return; + } + impl_ = std::make_shared(*native); + lifecycle_ = model->manager_lifecycle(); + session_lease_ = MakeSessionLease(lifecycle_); + manager_keepalive_ = std::move(manager_keepalive); } catch (const foundry_local::Error& e) { ThrowFoundryLocalError(env, static_cast(e.Code()), e.what()); return; @@ -584,6 +658,10 @@ AudioSession::AudioSession(const Napi::CallbackInfo& info) manager_ = Napi::Reference::New(model->manager().Value(), 1); } +AudioSession::~AudioSession() { + ReleaseSessionState(impl_, manager_keepalive_, session_lease_); +} + bool AudioSession::ThrowIfDisposed(Napi::Env env) { if (impl_ == nullptr) { ThrowFoundryLocalError(env, FOUNDRY_LOCAL_ERROR_INVALID_USAGE, @@ -601,14 +679,14 @@ Napi::Value AudioSession::ProcessRequest(const Napi::CallbackInfo& info) { return env.Undefined(); } Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); - return ProcessRequestOn(env, impl_.get(), info[0], std::move(owner)); + return ProcessRequestOn(env, impl_, info[0], std::move(owner), manager_keepalive_, session_lease_); } Napi::Value AudioSession::ProcessStreamingRequest(const Napi::CallbackInfo& info) { Napi::Env env = info.Env(); if (ThrowIfDisposed(env)) return env.Undefined(); Napi::ObjectReference owner = Napi::Reference::New(manager_.Value(), 1); - return ProcessStreamingRequestOn(env, impl_.get(), info, std::move(owner)); + return ProcessStreamingRequestOn(env, impl_, info, std::move(owner), manager_keepalive_, session_lease_); } Napi::Value AudioSession::SetOptions(const Napi::CallbackInfo& info) { @@ -627,7 +705,7 @@ Napi::Value AudioSession::SetOptions(const Napi::CallbackInfo& info) { } Napi::Value AudioSession::Dispose(const Napi::CallbackInfo& info) { - impl_.reset(); + ReleaseSessionState(impl_, manager_keepalive_, session_lease_); return info.Env().Undefined(); } diff --git a/sdk_v2/js/native/src/session.h b/sdk_v2/js/native/src/session.h index 2b1db7b86..796439e54 100644 --- a/sdk_v2/js/native/src/session.h +++ b/sdk_v2/js/native/src/session.h @@ -32,11 +32,14 @@ namespace foundry_local_node { +struct ManagerLifecycle; + class ChatSession : public Napi::ObjectWrap { public: static Napi::Function Init(Napi::Env env); explicit ChatSession(const Napi::CallbackInfo& info); + ~ChatSession(); private: Napi::Value ProcessRequest(const Napi::CallbackInfo& info); @@ -51,8 +54,11 @@ class ChatSession : public Napi::ObjectWrap { bool ThrowIfDisposed(Napi::Env env); - std::unique_ptr impl_; + std::shared_ptr impl_; Napi::ObjectReference manager_; + std::shared_ptr manager_keepalive_; + std::shared_ptr lifecycle_; + std::shared_ptr session_lease_; }; // Napi::ObjectWrap over foundry_local::EmbeddingsSession. @@ -73,6 +79,7 @@ class EmbeddingsSession : public Napi::ObjectWrap { static Napi::Function Init(Napi::Env env); explicit EmbeddingsSession(const Napi::CallbackInfo& info); + ~EmbeddingsSession(); private: Napi::Value ProcessRequest(const Napi::CallbackInfo& info); @@ -82,8 +89,11 @@ class EmbeddingsSession : public Napi::ObjectWrap { bool ThrowIfDisposed(Napi::Env env); - std::unique_ptr impl_; + std::shared_ptr impl_; Napi::ObjectReference manager_; + std::shared_ptr manager_keepalive_; + std::shared_ptr lifecycle_; + std::shared_ptr session_lease_; }; // Napi::ObjectWrap over foundry_local::AudioSession. @@ -101,6 +111,7 @@ class AudioSession : public Napi::ObjectWrap { static Napi::Function Init(Napi::Env env); explicit AudioSession(const Napi::CallbackInfo& info); + ~AudioSession(); private: Napi::Value ProcessRequest(const Napi::CallbackInfo& info); @@ -111,8 +122,11 @@ class AudioSession : public Napi::ObjectWrap { bool ThrowIfDisposed(Napi::Env env); - std::unique_ptr impl_; + std::shared_ptr impl_; Napi::ObjectReference manager_; + std::shared_ptr manager_keepalive_; + std::shared_ptr lifecycle_; + std::shared_ptr session_lease_; }; } // namespace foundry_local_node diff --git a/sdk_v2/js/script/copy-native.mjs b/sdk_v2/js/script/copy-native.mjs index 494d26671..34682f7dc 100644 --- a/sdk_v2/js/script/copy-native.mjs +++ b/sdk_v2/js/script/copy-native.mjs @@ -55,6 +55,7 @@ const wanted = (() => { "onnxruntime-genai.dll", "onnxruntime_providers_shared.dll", "Microsoft.Windows.AI.MachineLearning.dll", + "DirectML.dll", ]; } if (process.platform === "darwin") { diff --git a/sdk_v2/js/script/install-native.cjs b/sdk_v2/js/script/install-native.cjs index bbb372091..ae32ce0b4 100644 --- a/sdk_v2/js/script/install-native.cjs +++ b/sdk_v2/js/script/install-native.cjs @@ -30,7 +30,6 @@ const PLATFORM_MAP = { 'win32-x64': 'win-x64', 'win32-arm64': 'win-arm64', 'linux-x64': 'linux-x64', - 'linux-arm64': 'linux-arm64', 'darwin-arm64': 'osx-arm64', }; const platformKey = `${os.platform()}-${os.arch()}`; diff --git a/sdk_v2/js/script/pack-prebuilds.mjs b/sdk_v2/js/script/pack-prebuilds.mjs index 36e1d0913..5dbf558aa 100644 --- a/sdk_v2/js/script/pack-prebuilds.mjs +++ b/sdk_v2/js/script/pack-prebuilds.mjs @@ -43,18 +43,14 @@ const destDir = resolve(pkgRoot, "prebuilds", `${process.platform}-${process.arc mkdirSync(destDir, { recursive: true }); // foundry_local is required; ORT/GenAI siblings are excluded (fetched at install -// time). The WinML EP catalog DLL is an optional sibling bundled on Windows. +// time). The WinML EP catalog and DirectML DLLs are required siblings on Windows. const wanted = (() => { - if (process.platform === "win32") return ["foundry_local.dll"]; + if (process.platform === "win32") return ["foundry_local.dll", "Microsoft.Windows.AI.MachineLearning.dll", "DirectML.dll"]; if (process.platform === "darwin") return ["libfoundry_local.dylib"]; return ["libfoundry_local.so"]; })(); -// Optional native siblings — copied when present, skipped (with a warning) when -// not. The reg-free WinML 2.x runtime ships next to foundry_local.dll on Windows -// so WinML hardware EPs work out of the box without an install-time download. -const optional = - process.platform === "win32" ? ["Microsoft.Windows.AI.MachineLearning.dll"] : []; +const optional = []; let copied = 0; const available = new Set(readdirSync(sourceDir)); diff --git a/sdk_v2/js/src/catalog.ts b/sdk_v2/js/src/catalog.ts index 978a0f359..c08c8d579 100644 --- a/sdk_v2/js/src/catalog.ts +++ b/sdk_v2/js/src/catalog.ts @@ -4,6 +4,7 @@ // // The underlying native catalog operations are synchronous; the async surface here is for parity with the C# / // Python SDKs. `getModel`, `getModelVariant`, and `getLatestVersion` throw when the alias / id is not found. +// After the parent manager is disposed, native catalog methods throw a FoundryLocalError. import type { NativeCatalog, NativeModel } from "./detail/native.js"; import type { IModel } from "./imodel.js"; diff --git a/sdk_v2/js/test/manager-dispose.test.ts b/sdk_v2/js/test/manager-dispose.test.ts index fe7c7ba19..26bbe1c4e 100644 --- a/sdk_v2/js/test/manager-dispose.test.ts +++ b/sdk_v2/js/test/manager-dispose.test.ts @@ -62,6 +62,21 @@ describeIfBuilt("FoundryLocalManager.dispose", () => { } }); + it("a cached catalog handle throws after manager disposal and does not block a new manager", () => { + const mgr = freshManager("cached-catalog"); + const catalog = mgr.catalog; + mgr.dispose(); + + expect(() => catalog.name).toThrow(/disposed/i); + + const next = freshManager("after-cached-catalog"); + try { + expect(next.disposed).toBe(false); + } finally { + next.dispose(); + } + }); + it("Symbol.dispose is wired and idempotent", () => { const mgr = freshManager("symbol-dispose"); mgr[Symbol.dispose](); @@ -70,6 +85,15 @@ describeIfBuilt("FoundryLocalManager.dispose", () => { expect(mgr.disposed).toBe(true); }); + it("dispose() rejects while a native EP worker is in flight", async () => { + const mgr = freshManager("async-worker"); + const pending = mgr.downloadAndRegisterEps(["__not-a-provider__"]); + expect(() => mgr.dispose()).toThrow(/active native workers/i); + await expect(pending).rejects.toThrow(/not available/i); + mgr.dispose(); + expect(mgr.disposed).toBe(true); + }); + it("`using` declaration disposes at scope exit", () => { let captured: FoundryLocalManager | undefined; { diff --git a/sdk_v2/python/src/foundry_local_sdk/_native/lib_loader.py b/sdk_v2/python/src/foundry_local_sdk/_native/lib_loader.py index 1369759c8..0658c08d7 100644 --- a/sdk_v2/python/src/foundry_local_sdk/_native/lib_loader.py +++ b/sdk_v2/python/src/foundry_local_sdk/_native/lib_loader.py @@ -15,6 +15,7 @@ import pathlib import platform import sys +import sysconfig logger = logging.getLogger(__name__) @@ -29,7 +30,7 @@ def _lib_name() -> str: def _platform_subdir() -> str: if sys.platform == "win32": - return "win-x64" + return "win-arm64" if sysconfig.get_platform().lower() == "win-arm64" else "win-x64" if sys.platform == "darwin": return "osx-arm64" if platform.machine() == "arm64" else "osx-x64" return "linux-x64" @@ -248,7 +249,7 @@ def prepare_native_dependencies(foundry_local_dir: pathlib.Path) -> list: if add_dll_directory is not None: for d in {ort_path.parent, genai_path.parent, foundry_local_dir}: try: - add_dll_directory(str(d)) + handles.append(add_dll_directory(str(d))) except OSError as exc: logger.warning("os.add_dll_directory(%s) failed: %s", d, exc) @@ -278,22 +279,10 @@ def prepare_native_dependencies(foundry_local_dir: pathlib.Path) -> list: logger.warning("Failed to preload ORT (%s): %s", ort_path, exc) return handles - # macOS only: GenAI's static initializer does its own dlopen("libonnxruntime.dylib"), - # which on Darwin only matches by leafname against dyld search paths or images - # whose install_name leaf is exactly "libonnxruntime.dylib". ORT's dylib has a - # versioned install_name instead (the soversion, e.g. "@rpath/libonnxruntime.1.dylib"), - # so neither match path succeeds and GenAI aborts in dyld init before any of our code - # runs. The second name GenAI tries is "/libonnxruntime.dylib", so a symlink - # there fixes it. Linux dlopen consults the loaded-soname table by leafname and finds - # our already-RTLD_GLOBAL'd image without help. - if sys.platform == "darwin": - symlink_path = genai_path.parent / "libonnxruntime.dylib" - try: - if not symlink_path.exists(): - symlink_path.symlink_to(ort_path) - logger.info("Created macOS GenAI->ORT symlink: %s -> %s", symlink_path, ort_path) - except OSError as exc: - logger.warning("Failed to create macOS ORT symlink at %s: %s", symlink_path, exc) + # GenAI's lazy InitApi() first checks ORT_LIB_PATH before falling back to a + # platform leafname. Point it at the exact ORT dylib/so we resolved instead + # of mutating the installed GenAI package to add an unversioned alias. + os.environ["ORT_LIB_PATH"] = str(ort_path) try: handles.append(ctypes.CDLL(str(genai_path), mode=ctypes.RTLD_GLOBAL)) diff --git a/sdk_v2/python/src/foundry_local_sdk/items.py b/sdk_v2/python/src/foundry_local_sdk/items.py index 0f1f3c14f..809fe2b62 100644 --- a/sdk_v2/python/src/foundry_local_sdk/items.py +++ b/sdk_v2/python/src/foundry_local_sdk/items.py @@ -11,7 +11,7 @@ from enum import IntEnum -_API_VERSION = 1 # FOUNDRY_LOCAL_API_VERSION +_API_VERSION = 2 # FOUNDRY_LOCAL_API_VERSION # Sentinel for absent time fields in flSpeech* structs. INT64_MIN, matches # C ABI FOUNDRY_LOCAL_DURATION_UNSET. cffi cdef does not process #define. _DURATION_UNSET: int = -(2**63) diff --git a/sdk_v2/python/src/foundry_local_sdk/request.py b/sdk_v2/python/src/foundry_local_sdk/request.py index fe6879e81..5c3a90ca5 100644 --- a/sdk_v2/python/src/foundry_local_sdk/request.py +++ b/sdk_v2/python/src/foundry_local_sdk/request.py @@ -12,7 +12,7 @@ from foundry_local_sdk.items import Item from foundry_local_sdk.session_types import RequestOptions -_API_VERSION = 1 # FOUNDRY_LOCAL_API_VERSION +_API_VERSION = 2 # FOUNDRY_LOCAL_API_VERSION class Request: diff --git a/sdk_v2/python/src/foundry_local_sdk/response.py b/sdk_v2/python/src/foundry_local_sdk/response.py index df6748bc3..8df2e2f24 100644 --- a/sdk_v2/python/src/foundry_local_sdk/response.py +++ b/sdk_v2/python/src/foundry_local_sdk/response.py @@ -12,7 +12,7 @@ from foundry_local_sdk.items import Item from foundry_local_sdk.session_types import FinishReason, TokenUsage -_API_VERSION = 1 # FOUNDRY_LOCAL_API_VERSION +_API_VERSION = 2 # FOUNDRY_LOCAL_API_VERSION class Response: diff --git a/sdk_v2/python/src/foundry_local_sdk/session.py b/sdk_v2/python/src/foundry_local_sdk/session.py index 136609085..e913c0138 100644 --- a/sdk_v2/python/src/foundry_local_sdk/session.py +++ b/sdk_v2/python/src/foundry_local_sdk/session.py @@ -17,7 +17,7 @@ from foundry_local_sdk.response import Response from foundry_local_sdk.session_types import RequestOptions -_API_VERSION = 1 # FOUNDRY_LOCAL_API_VERSION +_API_VERSION = 2 # FOUNDRY_LOCAL_API_VERSION # Sentinel placed on the stream queue by the background thread when inference finishes. _DONE = object()