From 3a23ba00a53beac5ca1f96fcaaf4c53371903dc9 Mon Sep 17 00:00:00 2001 From: xugangqiang Date: Tue, 22 Sep 2026 20:12:38 +0800 Subject: [PATCH] feat(deepdoc): share ONNX model weights across pooled native sessions (#20035) --- build.sh | 227 ++++++++++++------ go.mod | 2 +- go.sum | 4 +- internal/deepdoc/native/ocr_rec.go | 14 +- internal/deepdoc/native/session.go | 134 ++++++++++- internal/deepdoc/native/session_pool.go | 11 +- .../native/shared_weights_memory_test.go | 141 +++++++++++ .../deepdoc/native/shared_weights_test.go | 169 +++++++++++++ internal/development.md | 25 ++ ragflow_deps/download_deps.py | 61 ++++- ragflow_deps/download_go_deps.py | 63 ++++- 11 files changed, 755 insertions(+), 96 deletions(-) create mode 100644 internal/deepdoc/native/shared_weights_memory_test.go create mode 100644 internal/deepdoc/native/shared_weights_test.go diff --git a/build.sh b/build.sh index 083c811bcb..478e4af76a 100755 --- a/build.sh +++ b/build.sh @@ -246,6 +246,108 @@ check_ort_version_consistency() { echo -e "${GREEN}✓ ONNX Runtime native version consistent: ${env_go}${NC}" } +# Check the ONNX Runtime static archive the Go DeepDoc backend links. +# +# This mirrors check_office_oxide_deps: it verifies the on-disk archive +# actually matches the released artifact, so a stale/incompatible .a — one +# re-issued under the SAME release tag and asset name, or an older download — +# fails fast HERE with an actionable message, instead of producing a binary +# that links fine but crashes at runtime when the shared-initializer feature +# (SessionGetInitializer*, added by infiniflow/ragflow-build) runs. +# +# The onnxruntime_go binding reaches ORT purely through the OrtApi +# function-pointer table (onnxruntime_go/shared_initializer.go -> +# onnxruntime_wrapper.c -> ort_api->SessionGetInitializer*), so a stale .a +# that lacks those custom slots links successfully and only fails at runtime. +# The check below catches that by comparing the local release ZIP against the +# published {asset}.sha256 sidecar (same source of truth the download scripts +# use). Devs who ran a download script have the zip under ragflow_deps/; CI +# seeds from /opt and has no zip, where the bake is authoritative. +check_onnxruntime_deps() { + local ort_version + ort_version="$(grep -m1 -E 'DeepDocORTVersion[[:space:]]*=[[:space:]]*"' \ + "${PROJECT_ROOT}/internal/common/environments.go" \ + | sed -E 's/.*"([^"]+)".*/\1/')" + if [ -z "$ort_version" ]; then + echo -e "${RED}Error: cannot parse DeepDocORTVersion from internal/common/environments.go${NC}" >&2 + return 1 + fi + + # Locate the ORT static archives under the version dir. Populate the + # globals ORT_A_FILES (space-separated .a paths) and ORT_A_SHA256 (sha256 + # of the primary libonnxruntime.a) for the caller's version-stamped + # cache path. + ORT_A_FILES="" + local primary="" + local f + while IFS= read -r f; do + case "$(basename "$f")" in + *cuda*|*tensorrt*|*coreml*|*dml*|*migraphx*) continue ;; + esac + case "$f" in + */onnxruntime-linux-x64-static_lib-"${ort_version}"*/lib/*.a) + ORT_A_FILES="$ORT_A_FILES $f" + [ -z "$primary" ] && primary="$f" + case "$(basename "$f")" in + libonnxruntime.a) primary="$f" ;; + esac + ;; + esac + done < <(find "$ONNXRUNTIME_STATIC_PREFIX" -type f -name '*.a' 2>/dev/null) + + if [ -z "$ORT_A_FILES" ]; then + local avail + avail="$(find "$ONNXRUNTIME_STATIC_PREFIX" -maxdepth 1 -type d \ + -name 'onnxruntime-linux-x64-static_lib-*' -exec basename {} \; 2>/dev/null | tr '\n' ' ')" + echo -e "${RED}Error: no ONNX Runtime ${ort_version} static lib under $ONNXRUNTIME_STATIC_PREFIX${NC}" >&2 + echo " available: ${avail:-}" >&2 + echo " DeepDocORTVersion=${ort_version}; fetch it with:" >&2 + echo " uv run python3 ragflow_deps/download_go_deps.py" >&2 + echo " or pre-seed /opt/ragflow-native-libs/onnxruntime (CI image)." >&2 + return 1 + fi + + # sha256 comparison against the published sidecar (best-effort, dev path). + local asset="onnxruntime-v${ort_version}-linux-x86_64.zip" + local zip_path="${PROJECT_ROOT}/ragflow_deps/${asset}" + if [ -f "$zip_path" ]; then + local sidecar_url="https://github.com/infiniflow/ragflow-build/releases/download/onnxruntime-v${ort_version}/${asset}.sha256" + local expected="" + if command -v curl >/dev/null 2>&1; then + expected="$(curl -fsSL "$sidecar_url" 2>/dev/null | awk '{print $1}')" + elif command -v python3 >/dev/null 2>&1; then + expected="$(python3 - "$sidecar_url" <<'PY' 2>/dev/null +import sys, urllib.request +try: + print(urllib.request.urlopen(sys.argv[1]).read().decode().split()[0]) +except Exception: + pass +PY +)" + fi + if [ -n "$expected" ]; then + local actual + actual="$(sha256sum "$zip_path" | awk '{print $1}')" + if [ "$actual" != "$expected" ]; then + echo -e "${RED}Error: ONNX Runtime archive ${asset} is stale${NC}" >&2 + echo " expected sha256: $expected" >&2 + echo " local sha256: $actual" >&2 + echo " A re-issued archive under the same tag/asset name was not picked up." >&2 + echo " Refresh it (this re-downloads + re-extracts the .a):" >&2 + echo " rm -f ${zip_path}" >&2 + echo " uv run python3 ragflow_deps/download_go_deps.py # or ragflow_deps/download_deps.py" >&2 + return 1 + fi + fi + else + echo " (no local ${asset}; skipping sha256 check — CI seed assumed authoritative)" + fi + + # Content hash of the primary archive for the version-stamped cache path. + ORT_A_SHA256="$(sha256sum "$primary" | awk '{print $1}')" + echo "✓ onnxruntime v${ort_version} static lib verified under $ONNXRUNTIME_STATIC_PREFIX" +} + # Check office_oxide native library check_office_oxide_deps() { print_section "Checking office_oxide native library" @@ -690,82 +792,57 @@ setup_cgo_env() { # below. Mirrors the existing _seed_from_system calls for the other # native libs so CI never downloads ORT at build/test time. _seed_from_system "onnxruntime" || true - if [ -d "$ONNXRUNTIME_STATIC_PREFIX" ]; then - # Collect every .a, but skip GPU-only providers we never build - # against (would pull in CUDA/cuDNN/TensorRT which we don't ship). - # - # Select the ORT static lib that matches the Go deepdoc backend's - # required version (DeepDocORTVersion in internal/common/environments.go). - # Go and Python build/link against independent ORT versions, so the - # static_lib prefix legitimately holds more than one - # onnxruntime-linux-x64-static_lib-* dir at once (e.g. the Python-side - # 1.23.x next to the Go-side 1.29.0). We pick the dir that matches - # DeepDocORTVersion rather than failing when a second version dir is - # present. This avoids silently linking the wrong version while still - # keeping the bake self-documenting. - local ort_version - ort_version="$(grep -m1 -E 'DeepDocORTVersion[[:space:]]*=[[:space:]]*"' \ - "${PROJECT_ROOT}/internal/common/environments.go" \ - | sed -E 's/.*"([^"]+)".*/\1/')" - if [ -z "$ort_version" ]; then - echo " Error: cannot parse DeepDocORTVersion from internal/common/environments.go" >&2 - return 1 - fi - local ort_a="" - while IFS= read -r f; do - case "$(basename "$f")" in - *cuda*|*tensorrt*|*coreml*|*dml*|*migraphx*) continue ;; - esac - case "$f" in - # Only collect .a from the dir matching the required version. - */onnxruntime-linux-x64-static_lib-"${ort_version}"*/lib/*.a) - ort_a="$ort_a $f" ;; - esac - done < <(find "$ONNXRUNTIME_STATIC_PREFIX" -type f -name '*.a' 2>/dev/null) + check_onnxruntime_deps || return 1 - if [ -n "$ort_a" ]; then - # Export exactly one symbol (OrtGetApiBase) for the binding's - # dlopen(NULL)+dlsym lookup via --dynamic-list (NOT --version-script - # with "local: *", which hides Go's runtime type symbols and breaks - # the PIE absolute relocations). No --whole-archive, so GNU ld's - # archive-level GC drops any ORT kernel/EP object nothing references. - # - # The dynamic list is written to a STABLE, project-scoped path (not - # mktemp) so CGO_LDFLAGS is reproducible across builds and survives - # across invocations; its content never changes, so overwriting is - # safe and nothing leaks in /tmp. .cache/ is gitignored. - local ort_dynamic_list="${PROJECT_ROOT}/.cache/ort_dynamic_list.txt" - mkdir -p "$(dirname "$ort_dynamic_list")" - printf '{\n OrtGetApiBase;\n};\n' > "$ort_dynamic_list" - # --undefined=OrtGetApiBase force-pulls the archive member that - # defines OrtGetApiBase (the Go binding reaches ORT only via - # dlsym("OrtGetApiBase"), so nothing references it at link time and - # it would otherwise be GC'd). From there the minimal build's CPU-EP - # registration call chain pulls in the operators the models use. - # The explicit space between $ort_dynamic_list and $ort_a keeps the - # two as separate linker arguments regardless of $ort_a's leading - # space. - export CGO_LDFLAGS="$CGO_LDFLAGS -Wl,--undefined=OrtGetApiBase -Wl,--dynamic-list=$ort_dynamic_list $ort_a -lstdc++" - echo " onnxruntime (static) → $ONNXRUNTIME_STATIC_PREFIX" - # The re2 regex-library collision between onnxruntime.a and - # librag_tokenizer_c_api.a is fixed at the .a level in build_cpp(): - # the tokenizer's bundled re2 symbols are renamed into a private - # namespace (ragtokre2_) so the two re2 copies never share a symbol - # name. --dynamic-list (above) exports OrtGetApiBase into the process - # dynamic symbol table, which is what the binding's dlopen(NULL)+dlsym - # lookup needs at runtime (no --export-dynamic required). - else - local avail - avail="$(find "$ONNXRUNTIME_STATIC_PREFIX" -maxdepth 1 -type d \ - -name 'onnxruntime-linux-x64-static_lib-*' -exec basename {} \; 2>/dev/null | tr '\n' ' ')" - echo " Error: no ONNX Runtime ${ort_version} static lib under $ONNXRUNTIME_STATIC_PREFIX" >&2 - echo " available: ${avail:-}" >&2 - echo " DeepDocORTVersion=${ort_version}; bake/download the matching ORT (or update DeepDocORTVersion)." >&2 - return 1 - fi - else - echo " onnxruntime static_lib not found ($ONNXRUNTIME_STATIC_PREFIX); the in-process DeepDoc backend cannot link ORT" >&2 - fi + # Version-stamp the archive paths so an in-place .a upgrade invalidates + # Go's build cache. Go's cache keys CGO_LDFLAGS as a string and does NOT + # hash the referenced .a contents, so swapping the .a in place (same path, + # e.g. a re-issue under the same tag/asset name) silently reuses a stale + # linked binary. Stamp with the primary archive's sha256 so the flag string + # changes when the content changes → automatic relink. + local ort_version + ort_version="$(grep -m1 -E 'DeepDocORTVersion[[:space:]]*=[[:space:]]*"' \ + "${PROJECT_ROOT}/internal/common/environments.go" \ + | sed -E 's/.*"([^"]+)".*/\1/')" + local ort_versioned_dir="${ONNXRUNTIME_STATIC_PREFIX}/v${ort_version}-${ORT_A_SHA256:0:16}" + mkdir -p "$ort_versioned_dir" + local ort_a_versioned="" + local f + for f in $ORT_A_FILES; do + ln -sf "$f" "$ort_versioned_dir/$(basename "$f")" + ort_a_versioned="$ort_a_versioned $ort_versioned_dir/$(basename "$f")" + done + + # Export exactly one symbol (OrtGetApiBase) for the binding's + # dlopen(NULL)+dlsym lookup via --dynamic-list (NOT --version-script + # with "local: *", which hides Go's runtime type symbols and breaks + # the PIE absolute relocations). No --whole-archive, so GNU ld's + # archive-level GC drops any ORT kernel/EP object nothing references. + # + # The dynamic list is written to a STABLE, project-scoped path (not + # mktemp) so CGO_LDFLAGS is reproducible across builds and survives + # across invocations; its content never changes, so overwriting is + # safe and nothing leaks in /tmp. .cache/ is gitignored. + local ort_dynamic_list="${PROJECT_ROOT}/.cache/ort_dynamic_list.txt" + mkdir -p "$(dirname "$ort_dynamic_list")" + printf '{\n OrtGetApiBase;\n};\n' > "$ort_dynamic_list" + # --undefined=OrtGetApiBase force-pulls the archive member that + # defines OrtGetApiBase (the Go binding reaches ORT only via + # dlsym("OrtGetApiBase"), so nothing references it at link time and + # it would otherwise be GC'd). From there the minimal build's CPU-EP + # registration call chain pulls in the operators the models use. + # The explicit space between $ort_dynamic_list and the versioned .a + # list keeps the two as separate linker arguments regardless of the + # list's leading space. + export CGO_LDFLAGS="$CGO_LDFLAGS -Wl,--undefined=OrtGetApiBase -Wl,--dynamic-list=$ort_dynamic_list $ort_a_versioned -lstdc++" + echo " onnxruntime (static) → $ONNXRUNTIME_STATIC_PREFIX" + # The re2 regex-library collision between onnxruntime.a and + # librag_tokenizer_c_api.a is fixed at the .a level in build_cpp(): + # the tokenizer's bundled re2 symbols are renamed into a private + # namespace (ragtokre2_) so the two re2 copies never share a symbol + # name. --dynamic-list (above) exports OrtGetApiBase into the process + # dynamic symbol table, which is what the binding's dlopen(NULL)+dlsym + # lookup needs at runtime (no --export-dynamic required). # ── platform-specific system libraries ──────────────────────────── case "$(uname -s)" in diff --git a/go.mod b/go.mod index 56f5f7db1d..3b56c8cfb9 100644 --- a/go.mod +++ b/go.mod @@ -45,7 +45,7 @@ require ( github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.3 github.com/infiniflow/infinity-go-sdk v0.0.0-00010101000000-000000000000 - github.com/infiniflow/onnxruntime_go v1.23.1 + github.com/infiniflow/onnxruntime_go v1.29.0 github.com/jackc/pgx/v5 v5.10.0 github.com/kaptinlin/jsonrepair v0.4.8 github.com/larksuite/oapi-sdk-go/v3 v3.9.9 diff --git a/go.sum b/go.sum index 0c21699a5c..de2dd663fc 100644 --- a/go.sum +++ b/go.sum @@ -351,8 +351,8 @@ github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/infiniflow/infinity/go v0.0.0-20260806040857-d755c5ad25d9 h1:mqPwCOTktwl0V1f3tkfCRFonGQ3XrQ43gvH+TOWec6o= github.com/infiniflow/infinity/go v0.0.0-20260806040857-d755c5ad25d9/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0= -github.com/infiniflow/onnxruntime_go v1.23.1 h1:PvhtCkWAmnMwy16QrkohkBKUOH/EOrEZm0LO/QykIo0= -github.com/infiniflow/onnxruntime_go v1.23.1/go.mod h1:dFFL0KdgeeEDXoE9cjrg/ipQg0wEvD0tkGIg2UH0jdE= +github.com/infiniflow/onnxruntime_go v1.29.0 h1:pI6v4Bcc7BJdKRF+ZZZkaf4HV6Of6Kl6XryaW2aGbGQ= +github.com/infiniflow/onnxruntime_go v1.29.0/go.mod h1:dFFL0KdgeeEDXoE9cjrg/ipQg0wEvD0tkGIg2UH0jdE= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= diff --git a/internal/deepdoc/native/ocr_rec.go b/internal/deepdoc/native/ocr_rec.go index ab87a4fa3d..44a1cf8765 100644 --- a/internal/deepdoc/native/ocr_rec.go +++ b/internal/deepdoc/native/ocr_rec.go @@ -302,14 +302,18 @@ type recSession struct { poisoned bool } -func newRecSession(modelPath, inName string, inShape []int64, outName string) (*recSession, error) { - opts, err := newSessionOptions() +func newRecSession(modelPath, inName string, inShape []int64, outName string, weights *weightSet) (*recSession, error) { + // Build the same options as NewSession: one intra-op thread, BFC arena + // disabled, and the shared initializers injected when weights != nil. See + // newSessionOptions in session.go for the rationale. Reusing it keeps the + // rec pool's session configuration in lockstep with the det/DLA/TSR pool. + opts, err := newSessionOptions(weights) if err != nil { return nil, err } - // See NewSession: the C session does not take ownership of opts, so release - // it once the session is built. newSessionOptions centralizes the - // intra-op-thread and arena settings shared with NewSession. + // The C session copies these options at creation time, so the options handle + // can be released once the session is built. The shared weight buffers are + // owned by the process-wide weightCache and outlive every session. defer opts.Destroy() sess, err := ort.NewDynamicAdvancedSession(modelPath, []string{inName}, []string{outName}, opts) diff --git a/internal/deepdoc/native/session.go b/internal/deepdoc/native/session.go index 8ddd3fd687..fed1df7d27 100644 --- a/internal/deepdoc/native/session.go +++ b/internal/deepdoc/native/session.go @@ -18,6 +18,7 @@ package native import ( "context" "fmt" + "log" "sync" ort "github.com/infiniflow/onnxruntime_go" @@ -93,6 +94,87 @@ type session struct { poisoned bool } +// weightSet is a process-wide cached copy of one model's constant initializers +// (its weights). Every pooled session of the same modelPath injects these +// shared buffers into its SessionOptions so the weight buffers live in memory +// exactly once, instead of each of the ~220 live rec/det/DLA/TSR sessions +// deserializing its own independent copy. The SharedInitializers (and the +// malloc'd buffers they wrap) are owned here and must outlive every session +// that references them; they are never Destroy'd until process exit, which is +// safe because the DeepDoc models are fixed for the life of the process. +type weightSet struct { + names []string + vals []*ort.SharedInitializer +} + +var ( + // weightMu guards weightCache. Extraction is cheap but must not race. + weightMu sync.Mutex + weightCache = map[string]*weightSet{} +) + +// sharedWeights extracts the constant initializers from modelPath once and +// caches them keyed by model path. Each subsequent call for the same model +// returns the cached buffers, so every pooled session of that model shares a +// single copy of its weights. Extraction uses a throwaway AdvancedSession (the +// only onnxruntime_go type exposing the GetInitializer* API); the session is +// Destroy'd right after extraction, leaving the SharedInitializers (user-owned +// malloc'd copies) alive in the cache. Returns (nil, nil) when the model has no +// initializers worth sharing. +func sharedWeights(modelPath, inName string, inShape []int64, outName string) (*weightSet, error) { + weightMu.Lock() + defer weightMu.Unlock() + if ws, ok := weightCache[modelPath]; ok { + return ws, nil + } + // Throwaway extraction session: the weights live in the graph, not in any + // input/output tensor, so we only need valid in/out tensors to load the + // model. We never Run it, so the output tensor's shape is irrelevant. + inT, err := ort.NewTensor(ort.NewShape(inShape...), make([]float32, prod(inShape))) + if err != nil { + return nil, fmt.Errorf("allocate extraction input for %s: %w", modelPath, err) + } + defer inT.Destroy() + outT, err := ort.NewTensor(ort.NewShape(1), []float32{0}) + if err != nil { + return nil, fmt.Errorf("allocate extraction output for %s: %w", modelPath, err) + } + defer outT.Destroy() + ext, err := ort.NewAdvancedSession(modelPath, + []string{inName}, []string{outName}, + []ort.Value{inT}, []ort.Value{outT}, nil) + if err != nil { + return nil, fmt.Errorf("open extraction session for %s: %w", modelPath, err) + } + defer ext.Destroy() + + count, err := ext.GetInitializerCount() + if err != nil { + return nil, fmt.Errorf("initializer count for %s: %w", modelPath, err) + } + if count == 0 { + return nil, nil + } + ws := &weightSet{ + names: make([]string, 0, count), + vals: make([]*ort.SharedInitializer, 0, count), + } + for i := 0; i < count; i++ { + name, err := ext.GetInitializerName(i) + if err != nil { + return nil, fmt.Errorf("initializer name %d for %s: %w", i, modelPath, err) + } + val, err := ext.GetInitializer(name) + if err != nil { + return nil, fmt.Errorf("get initializer %q for %s: %w", name, modelPath, err) + } + ws.names = append(ws.names, name) + ws.vals = append(ws.vals, val) + } + weightCache[modelPath] = ws + return ws, nil +} + // NewSession opens modelPath. inShape describes the fixed input tensor // dimensions; output tensors are allocated per Run (their shape is // model-determined, so no outShape argument is needed). The session runs @@ -100,16 +182,41 @@ type session struct { // the process-wide ceiling owned by the inference budget the process owner // registers (see inference_limit.go). Input/output tensors are allocated and // freed on every Run (see session.Run), so a pooled session holds only its -// weights in steady state. InitORT must have been called first. +// weights in steady state. Weight sharing is applied transparently: the model's +// constant initializers are extracted once per modelPath and injected into the +// session options, so every session of the same model shares a single copy of +// the weight buffers. InitORT must have been called first. func NewSession(modelPath, inName string, inShape []int64, outName string) (*session, error) { - opts, err := newSessionOptions() + weights, werr := sharedWeights(modelPath, inName, inShape, outName) + if werr != nil { + // Degrade gracefully: a model still loads and runs correctly without + // sharing; it just deserializes its own weight copy. + log.Printf("deepdoc/native: weight sharing unavailable for %s: %v", + modelPath, werr) + weights = nil + } + return newRawSession(modelPath, inName, inShape, outName, weights) +} + +// newRawSession opens modelPath with no weight sharing unless weights != nil, +// in which case each shared initializer is injected into the session options +// before the model is loaded. The session is a DynamicAdvancedSession: it owns +// no input/output tensors. Every Run allocates fresh input/output tensors and +// frees them afterwards (see session.Run), so a pooled session's steady-state +// native memory is just its weights plus ORT's plan cache — never the pinned +// in/out buffers. It is the single point where the ORT running session is +// created; sharedWeights funnels its one-shot extraction source through a +// separate AdvancedSession (the only type exposing GetInitializer*). +func newRawSession(modelPath, inName string, inShape []int64, outName string, weights *weightSet) (*session, error) { + opts, err := newSessionOptions(weights) if err != nil { return nil, err } - // newSessionOptions builds opts; the C session does not take ownership of it - // (see newSessionOptions / onnxruntime_go), so release it once the session is - // built. Releasing after a successful build is the documented pattern and - // avoids leaking the native SessionOptions handle on every session. + // The C session copies these options (including the shared-initializer + // references) at creation time, so the options handle can be released once + // the session is built. The shared weight buffers themselves are owned by + // the process-wide weightCache and outlive every session, so releasing opts + // here does not free them. defer opts.Destroy() sess, err := ort.NewDynamicAdvancedSession(modelPath, []string{inName}, []string{outName}, opts) @@ -127,8 +234,10 @@ func NewSession(modelPath, inName string, inShape []int64, outName string) (*ses // one intra-op thread (so each Run costs exactly one CPU thread — see the // intraOpThreads constant and inference_limit.go) and the BFC arena disabled // (idle sessions then keep only their weights; activation tensors are allocated -// per Run and freed after). -func newSessionOptions() (*ort.SessionOptions, error) { +// per Run and freed after). When weights != nil, each shared initializer is +// injected so the model's constant buffers live in memory a single time across +// every pooled session of the same model. +func newSessionOptions(weights *weightSet) (*ort.SessionOptions, error) { opts, err := ort.NewSessionOptions() if err != nil { return nil, err @@ -151,6 +260,15 @@ func newSessionOptions() (*ort.SessionOptions, error) { opts.Destroy() return nil, err } + if weights != nil { + for i, v := range weights.vals { + if err := opts.AddInitializer(weights.names[i], v); err != nil { + opts.Destroy() + return nil, fmt.Errorf("inject shared initializer %q: %w", + weights.names[i], err) + } + } + } return opts, nil } diff --git a/internal/deepdoc/native/session_pool.go b/internal/deepdoc/native/session_pool.go index 37233af748..4992f670be 100644 --- a/internal/deepdoc/native/session_pool.go +++ b/internal/deepdoc/native/session_pool.go @@ -33,6 +33,7 @@ package native import ( "context" + "log" "reflect" "strconv" "strings" @@ -278,6 +279,14 @@ var recSessions = newSessionPool[recKey, *recSession](recMaxShapePools, recShape func getRecSession(ctx context.Context, modelPath, inName string, inShape []int64, outName string) (*recSession, func(), error) { key := recKeyOf(modelPath, inName, inShape, outName) return recSessions.Get(ctx, key, func() (*recSession, error) { - return newRecSession(modelPath, inName, inShape, outName) + // Weight sharing is best-effort: a failure here degrades to a normal + // (non-shared) session rather than breaking OCR-rec entirely. + weights, werr := sharedWeights(modelPath, inName, inShape, outName) + if werr != nil { + log.Printf("deepdoc/native: rec weight sharing unavailable for %s: %v", + modelPath, werr) + weights = nil + } + return newRecSession(modelPath, inName, inShape, outName, weights) }) } diff --git a/internal/deepdoc/native/shared_weights_memory_test.go b/internal/deepdoc/native/shared_weights_memory_test.go new file mode 100644 index 0000000000..2d4026cab2 --- /dev/null +++ b/internal/deepdoc/native/shared_weights_memory_test.go @@ -0,0 +1,141 @@ +//go:build cgo && manual + +package native + +// Manual (local-only) verification that the patched ONNX Runtime static +// library actually shares one copy of the model weights across many sessions, +// instead of each session deserializing its own private copy. +// +// Empirical method (run twice, one mode per process to avoid glibc arena +// cross-contamination): +// MODE=nonshared -> open N sessions that each deserialize their own weights +// MODE=shared -> open N sessions that all inject the same cached weight set +// In each mode we create the sessions one at a time (never freeing between) and +// record the per-session RSS delta. The non-shared slope includes the +// per-session weight copy (~model weights); the shared slope must not — the +// weights live in the single cached set allocated once at extraction. If +// sharing works, the shared per-session delta is ~model-weights smaller than +// the non-shared one. +// +// This is the runtime counterpart to the ORT source proof: +// - SessionOptions::AddInitializer stores the OrtValue pointer (no data copy) +// (session_options.cc) +// - the session consumes it via `ort_value = *(initializers_to_share_map[name])` +// — a shallow OrtValue copy that keeps the same user-owned data buffer +// (session_state_utils.cc). +// +// rec.ort is used because its input tensor is tiny (~184KB), so per-session +// input-tensor allocation does not mask the weight-sharing signal. +// +// Run with: +// MODE=shared MODEL_DIR=/path/to/models \ +// go test -tags cgo,manual -run TestWeightSharingMemory -v . +// MODE=nonshared ... (same, different MODE) + +import ( + "context" + "os" + "path/filepath" + "runtime" + "strconv" + "strings" + "testing" +) + +// The native package's TestMain skips the whole binary unless the external +// golden testdata was fetched. This manual verification needs only MODEL_DIR +// (the .ort weights), not the golden fixtures, so opt in here. +func init() { + testdataFetchAttempted = true +} + +func rssMB() int64 { + data, err := os.ReadFile("/proc/self/statm") + if err != nil { + return 0 + } + fields := strings.Fields(string(data)) + if len(fields) < 2 { + return 0 + } + pages, err := strconv.ParseInt(fields[1], 10, 64) + if err != nil { + return 0 + } + return (pages * int64(os.Getpagesize())) >> 20 +} + +func fileSizeMB(path string) int64 { + fi, err := os.Stat(path) + if err != nil { + return 0 + } + return fi.Size() >> 20 +} + +// TestWeightSharingMemory opens N rec sessions in shared or non-shared mode and +// reports the total RSS delta. Compare the two modes to confirm sharing removes +// the per-session weight copy from the slope. +func TestWeightSharingMemory(t *testing.T) { + if os.Getenv("MODEL_DIR") == "" { + t.Skip("set MODEL_DIR to run the weight-sharing memory verification") + } + mode := os.Getenv("MODE") + if mode != "shared" && mode != "nonshared" { + t.Skip("set MODE=shared or MODE=nonshared") + } + if err := InitORT(); err != nil { + t.Fatalf("InitORT: %v", err) + } + + modelPath := filepath.Join(os.Getenv("MODEL_DIR"), "rec.ort") + inName := "x" + inShape := []int64{1, 3, 48, 320} + outName := "softmax_11.tmp_0" + + const N = 50 + + var ws *weightSet + if mode == "shared" { + var e error + ws, e = sharedWeights(modelPath, inName, inShape, outName) + if e != nil { + t.Fatalf("sharedWeights: %v", e) + } + t.Logf("extracted %d initializers (one shared copy, allocated once)", len(ws.vals)) + } + + sessions := make([]*recSession, 0, N) + rss0 := rssMB() + prev := rss0 + var totalDelta int64 + for i := 0; i < N; i++ { + var h *recSession + var e error + if mode == "shared" { + h, e = newRecSession(modelPath, inName, inShape, outName, ws) + } else { + h, e = newRecSession(modelPath, inName, inShape, outName, nil) + } + if e != nil { + t.Fatalf("newRecSession #%d (%s): %v", i+1, mode, e) + } + sessions = append(sessions, h) + runtime.GC() + cur := rssMB() + totalDelta += cur - prev + prev = cur + } + // Exercise every session so ORT actually uses the weights. + ctx := context.Background() + input := make([]float32, prod(inShape)) + for _, h := range sessions { + if _, err := h.Run(ctx, input); err != nil { + t.Fatalf("Run: %v", err) + } + } + t.Logf("MODE=%s: N=%d sessions, total RSS delta=%d MB (from %d MB)", mode, N, totalDelta, rss0) + for _, h := range sessions { + h.Destroy() + } +} diff --git a/internal/deepdoc/native/shared_weights_test.go b/internal/deepdoc/native/shared_weights_test.go new file mode 100644 index 0000000000..6900b5ea41 --- /dev/null +++ b/internal/deepdoc/native/shared_weights_test.go @@ -0,0 +1,169 @@ +//go:build cgo && integration + +package native + +// Integration tests for the per-model weight-sharing integration point. They +// require MODEL_DIR (gated by skipIfNoModels) and a statically linked ONNX +// Runtime resolved via dlopen(NULL), so they self-skip wherever the DeepDoc +// model snapshot or the static ORT archives are absent. When CI bakes both and +// sets DEEPDOC_NATIVE_REQUIRED=1 they run as a hard gate instead of a skip. +// +// Run with: +// MODEL_DIR=... go test -tags cgo,integration -run TestWeightSharing ./internal/deepdoc/native/... + +import ( + "context" + "os" + "path/filepath" + "reflect" + "testing" +) + +// These integration tests need only MODEL_DIR (the .ort weights), not the +// golden fixtures, so opt the package run in without the fetch_testdata tag +// (which would try to download fixtures we don't use here). +func init() { + testdataFetchAttempted = true +} + +// deepdocModels lists every .ort the in-process backend loads, with the +// fixed-shape I/O signature each recognizer uses. The rec.ort width here is +// arbitrary — extraction does not run the graph, so only the rank matters. +var deepdocModels = []struct { + name string + inName string + inShape []int64 + outName string +}{ + {"layout.ort", "images", []int64{1, 3, dlaInputSize, dlaInputSize}, "output0"}, + {"tsr.ort", "images", []int64{1, 3, tsrInputSize, tsrInputSize}, "output0"}, + {"det.ort", "x", []int64{1, 3, 736, 736}, "sigmoid_0.tmp_0"}, + {"rec.ort", "x", []int64{1, 3, 48, 320}, "softmax_11.tmp_0"}, +} + +// TestWeightSharingExtractsAndCaches proves the integration point for every +// model: each model's constant initializers are extracted exactly once per +// modelPath and cached, and a repeat lookup — or any NewSession/recSession for +// that model — reuses the same shared weight set instead of deserializing a +// fresh copy. +func TestWeightSharingExtractsAndCaches(t *testing.T) { + skipIfNoModels(t) + + // Force a clean cache so the extraction count is deterministic. + weightMu.Lock() + weightCache = map[string]*weightSet{} + weightMu.Unlock() + + for _, m := range deepdocModels { + modelPath := filepath.Join(os.Getenv("MODEL_DIR"), m.name) + + ws1, err := sharedWeights(modelPath, m.inName, m.inShape, m.outName) + if err != nil { + t.Fatalf("sharedWeights(%s) first call: %v", m.name, err) + } + if len(ws1.vals) == 0 { + t.Fatalf("sharedWeights(%s) extracted 0 initializers; expected > 0", m.name) + } + + ws2, err := sharedWeights(modelPath, m.inName, m.inShape, m.outName) + if err != nil { + t.Fatalf("sharedWeights(%s) second call: %v", m.name, err) + } + if ws2 != ws1 { + t.Fatalf("sharedWeights(%s) not cached: second call returned %p, first %p", + m.name, ws2, ws1) + } + + weightMu.Lock() + cached := weightCache[modelPath] + weightMu.Unlock() + if cached != ws1 { + t.Fatalf("weightCache[%s] = %p, expected the cached set %p", m.name, cached, ws1) + } + } +} + +// TestWeightSharingProducesIdenticalInference proves that injecting the shared +// weight buffers into a fixed-shape (layout.ort) session yields numerically +// identical inference to a session that deserializes its own private copy of +// the weights. Sharing changes the memory footprint, not the arithmetic; if +// extraction copied the weights wrong, the two outputs diverge. +func TestWeightSharingProducesIdenticalInference(t *testing.T) { + skipIfNoModels(t) + + m := deepdocModels[0] // layout.ort + modelPath := filepath.Join(os.Getenv("MODEL_DIR"), m.name) + inShape := m.inShape + + // Non-shared: ORT deserializes its own private copy of the constants. + alone, err := newRawSession(modelPath, m.inName, inShape, m.outName, nil) + if err != nil { + t.Fatalf("newRawSession (non-shared): %v", err) + } + defer alone.Destroy() + + // Shared: weights extracted once and injected via AddInitializer. + shared, err := NewSession(modelPath, m.inName, inShape, m.outName) + if err != nil { + t.Fatalf("NewSession (shared): %v", err) + } + defer shared.Destroy() + + ctx := context.Background() + input := make([]float32, prod(inShape)) + outAlone, err := alone.Run(ctx, input) + if err != nil { + t.Fatalf("Run (non-shared): %v", err) + } + outShared, err := shared.Run(ctx, input) + if err != nil { + t.Fatalf("Run (shared): %v", err) + } + if !reflect.DeepEqual(outAlone, outShared) { + t.Fatalf("shared weights changed inference output (%d vs %d elems diverge)", + len(outAlone), len(outShared)) + } +} + +// TestWeightSharingRecProducesIdenticalInference covers the dynamic-width +// (OCR-rec) injection point: newRecSession shares the same weight set as the +// one used by the pooled recSession. +func TestWeightSharingRecProducesIdenticalInference(t *testing.T) { + skipIfNoModels(t) + + m := deepdocModels[3] // rec.ort + modelPath := filepath.Join(os.Getenv("MODEL_DIR"), m.name) + inShape := m.inShape + + weights, err := sharedWeights(modelPath, m.inName, inShape, m.outName) + if err != nil { + t.Fatalf("sharedWeights(rec.ort): %v", err) + } + + alone, err := newRecSession(modelPath, m.inName, inShape, m.outName, nil) + if err != nil { + t.Fatalf("newRecSession (non-shared): %v", err) + } + defer alone.Destroy() + + shared, err := newRecSession(modelPath, m.inName, inShape, m.outName, weights) + if err != nil { + t.Fatalf("newRecSession (shared): %v", err) + } + defer shared.Destroy() + + ctx := context.Background() + input := make([]float32, prod(inShape)) + outAlone, err := alone.Run(ctx, input) + if err != nil { + t.Fatalf("Run (non-shared rec): %v", err) + } + outShared, err := shared.Run(ctx, input) + if err != nil { + t.Fatalf("Run (shared rec): %v", err) + } + if !reflect.DeepEqual(outAlone, outShared) { + t.Fatalf("shared rec weights changed inference output (%d vs %d elems diverge)", + len(outAlone), len(outShared)) + } +} diff --git a/internal/development.md b/internal/development.md index 8d4ee7ae06..3bcebaab80 100644 --- a/internal/development.md +++ b/internal/development.md @@ -117,6 +117,31 @@ python3 ragflow_deps/download_go_deps.py > `pyproject.toml` is versioned independently and is intentionally not part of > this check. +> **Note**: `build.sh` also guards ONNX Runtime **archive integrity** and the +> **link cache**, because a silently-stale `.a` is easy to miss: +> - `check_onnxruntime_deps` compares the local release ZIP under `ragflow_deps/` +> (`onnxruntime-v-linux-x86_64.zip`) against the published +> `onnxruntime-v-linux-x86_64.zip.sha256` sidecar — the same source of +> truth the download scripts use. If they differ (the archive was re-issued +> under the same tag/asset name, or an older download is present) the build +> fails fast with `Error: ONNX Runtime archive ... is stale`, printing the +> expected/actual sha256 and the exact refresh command +> (`rm -f ` + `uv run python3 ragflow_deps/download_go_deps.py`). CI seeds +> ORT from `/opt` and has no local zip, so the check is skipped there (the bake +> is authoritative). This matters because the `onnxruntime_go` binding reaches +> ORT only through the OrtApi function-pointer table +> (`ort_api->SessionGetInitializer*`), so a `.a` lacking a custom slot links +> successfully and only crashes at runtime — not at link time. +> - The link passes `CGO_LDFLAGS` a version-stamped path +> `onnxruntime/static_lib/v-/libonnxruntime.a`. Go's build +> cache keys `CGO_LDFLAGS` as a string and does NOT hash the referenced `.a`, +> so swapping the `.a` in place (same path, re-issued under the same name) +> would otherwise silently reuse a stale linked binary. Stamping the path with +> the archive's sha256 changes the flag string whenever the content changes → +> automatic relink. If you ever see a stale-`.a` crash after an ORT re-issue, +> re-run `uv run python3 ragflow_deps/download_go_deps.py` (or `download_deps.py`) +> so the stamp moves; a plain `go clean -cache` also forces it. + ### 1.5 Build RAGFlow - Build binary diff --git a/ragflow_deps/download_deps.py b/ragflow_deps/download_deps.py index f7c7bb6052..2a25cf5405 100755 --- a/ragflow_deps/download_deps.py +++ b/ragflow_deps/download_deps.py @@ -31,6 +31,7 @@ # is unaffected by where these files live locally. import argparse +import hashlib import os import shutil import sys @@ -56,7 +57,10 @@ from huggingface_hub import snapshot_download # Source of the native static archives: infiniflow/ragflow-build (our own # ORT-only minimal build), NOT the third-party csukuangfj/onnxruntime-libs # account. The release tag is `onnxruntime-v{ORT_VERSION}` and the asset is -# `onnxruntime-v{ORT_VERSION}-linux-x86_64.zip`. +# `onnxruntime-v{ORT_VERSION}-linux-x86_64.zip`. The archive is occasionally +# re-issued under this SAME tag/asset name with patched content; download_deps.py +# detects that via a `{asset}.sha256` sidecar and re-downloads/re-extracts, so a +# stale local copy never silently lingers. ORT_VERSION = "1.29.0" @@ -75,6 +79,15 @@ def _ort_normalized_dir(version): return f"onnxruntime-linux-x64-static_lib-{version}-glibc2_28" +def _sha256_of(path): + """sha256 of a file, streamed in chunks so large archives don't blow memory.""" + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + def get_urls(use_china_mirrors=False) -> list[str | list[str]]: if use_china_mirrors: return [ @@ -213,8 +226,52 @@ if __name__ == "__main__": download_url = url[0] if isinstance(url, list) else url filename = url[1] if isinstance(url, list) else url.split("/")[-1] print(f"Downloading {filename} from {download_url}...") - if not os.path.exists(filename): + + # The ONNX Runtime archive is re-issued under the SAME release tag and + # asset name whenever its content changes (e.g. the patched build that + # exports SessionGetInitializer*). A pure existence check would then + # keep a colleague's stale local copy and fail to link onnxruntime_go. + # Verify against the published .sha256 sidecar so a re-issued archive + # is always re-downloaded and re-extracted. Every other archive keeps + # the legacy existence-based skip. + is_ort = filename == _ort_asset_name(ORT_VERSION) + expected_sha = None + if is_ort: + sidecar_url = download_url + ".sha256" + try: + with urllib.request.urlopen(sidecar_url) as resp: + expected_sha = resp.read().decode().split()[0] + except Exception as exc: # noqa: BLE001 - best-effort; fall back to legacy + print(f" WARNING: could not fetch {sidecar_url} ({exc}); skipping checksum for {filename}") + + needs_download = True + if os.path.exists(filename): + if expected_sha is not None: + actual = _sha256_of(filename) + if actual == expected_sha: + print(f" ✓ {filename} checksum matches released {expected_sha}; skipping download") + needs_download = False + else: + print(f" {filename} checksum mismatch (local {actual} != released {expected_sha}); re-downloading") + else: + needs_download = False + + if needs_download: urllib.request.urlretrieve(download_url, filename) + if expected_sha is not None: + actual = _sha256_of(filename) + if actual != expected_sha: + print(f" ERROR: {filename} checksum mismatch after download (got {actual}, expected {expected_sha})", file=sys.stderr) + sys.exit(1) + print(f" ✓ {filename} checksum verified ({actual})") + # Force re-extract below: drop any previously extracted version dir + # so the same-named re-issued archive actually refreshes the .a files. + if is_ort: + native_libs = os.path.expanduser("~/ragflow-native-libs") + version_dir = os.path.join(native_libs, "onnxruntime", "static_lib", _ort_normalized_dir(ORT_VERSION)) + if os.path.isdir(version_dir): + print(f" Removing stale extracted ONNX Runtime dir: {version_dir}") + shutil.rmtree(version_dir) # Extract native static libraries to ~/ragflow-native-libs for Go build. # Ensures build.sh can find them without network access. diff --git a/ragflow_deps/download_go_deps.py b/ragflow_deps/download_go_deps.py index d0dbd46d67..a333c4bc39 100755 --- a/ragflow_deps/download_go_deps.py +++ b/ragflow_deps/download_go_deps.py @@ -52,6 +52,7 @@ # cl100k estimate instead (see internal/tokenizer/embedding_token_limits.md). import argparse +import hashlib import os import shutil import sys @@ -71,7 +72,11 @@ import requests # Source of the native static archives: infiniflow/ragflow-build (our own # ORT-only minimal build), NOT the third-party csukuangfj/onnxruntime-libs # account. The release tag is `onnxruntime-v{ORT_VERSION}` and the asset is -# `onnxruntime-v{ORT_VERSION}-linux-x86_64.zip`. +# `onnxruntime-v{ORT_VERSION}-linux-x86_64.zip`. The archive is occasionally +# re-issued under this SAME tag/asset name with patched content; both download +# scripts (this one and download_deps.py) detect that via a `{asset}.sha256` +# sidecar and re-download/re-extract, so a stale local copy never silently +# lingers. ORT_VERSION = "1.29.0" @@ -90,6 +95,15 @@ def _ort_normalized_dir(version): return f"onnxruntime-linux-x64-static_lib-{version}-glibc2_28" +def _sha256_of(path): + """sha256 of a file, streamed in chunks so large archives don't blow memory.""" + h = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return h.hexdigest() + + # Mirrors internal/common.DeepDocModelFiles (Go in-process DeepDoc backend). # These are the ONLY weights the Go backend loads; the full InfiniFlow/deepdoc # repo also ships .onnx (Python-only), which this Go-only script deliberately @@ -409,8 +423,53 @@ if __name__ == "__main__": download_url = url[0] if isinstance(url, list) else url filename = url[1] if isinstance(url, list) else url.split("/")[-1] print(f"Downloading {filename} from {download_url}...") - if not os.path.exists(filename): + + # The ONNX Runtime archive is re-issued under the SAME release tag and + # asset name whenever its content changes (e.g. the patched build that + # exports SessionGetInitializer*). A pure existence check would then + # keep a colleague's stale local copy and fail to link onnxruntime_go. + # Verify against the published .sha256 sidecar so a re-issued archive + # is always re-downloaded and re-extracted. Every other archive keeps + # the legacy existence-based skip. + is_ort = filename == _ort_asset_name(ORT_VERSION) + expected_sha = None + if is_ort: + sidecar_url = download_url + ".sha256" + try: + resp = requests.get(sidecar_url, timeout=30) + resp.raise_for_status() + expected_sha = resp.text.split()[0] + except Exception as exc: # noqa: BLE001 - best-effort; fall back to legacy + print(f" WARNING: could not fetch {sidecar_url} ({exc}); skipping checksum for {filename}") + + needs_download = True + if os.path.exists(filename): + if expected_sha is not None: + actual = _sha256_of(filename) + if actual == expected_sha: + print(f" ✓ {filename} checksum matches released {expected_sha}; skipping download") + needs_download = False + else: + print(f" {filename} checksum mismatch (local {actual} != released {expected_sha}); re-downloading") + else: + needs_download = False + + if needs_download: download_with_progress(download_url, filename) + if expected_sha is not None: + actual = _sha256_of(filename) + if actual != expected_sha: + print(f" ERROR: {filename} checksum mismatch after download (got {actual}, expected {expected_sha})", file=sys.stderr) + sys.exit(1) + print(f" ✓ {filename} checksum verified ({actual})") + # Force re-extract below: drop any previously extracted version dir + # so the same-named re-issued archive actually refreshes the .a files. + if is_ort: + native_libs = os.path.expanduser("~/ragflow-native-libs") + version_dir = os.path.join(native_libs, "onnxruntime", "static_lib", _ort_normalized_dir(ORT_VERSION)) + if os.path.isdir(version_dir): + print(f" Removing stale extracted ONNX Runtime dir: {version_dir}") + shutil.rmtree(version_dir) # Extract native static libraries to ~/ragflow-native-libs for Go build. # Ensures build.sh can find them without network access.