infiniflow/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-10-12 02:16:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(deepdoc): share ONNX model weights across pooled native sessions (#20035)

Browse Source
This commit is contained in:
xugangqiang
2026-09-22 20:12:38 +08:00
committed by GitHub
parent cf1b20ca26
commit 3a23ba00a5
11 changed files with 755 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
build.sh
+152 -75
View File
@@ -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:-<none>}" >&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:-<none>}" >&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
go.mod
+1 -1
View File
@@ -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
go.sum
+2 -2
View File
@@ -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=
internal/deepdoc/native/ocr_rec.go
+9 -5
View File
@@ -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)
internal/deepdoc/native/session.go
+126 -8
View File
@@ -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
}
internal/deepdoc/native/session_pool.go
+10 -1
View File
@@ -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)
})
}
internal/deepdoc/native/shared_weights_memory_test.go
+141
View File
@@ -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()
}
}
internal/deepdoc/native/shared_weights_test.go
+169
View File
@@ -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))
}
}
internal/development.md
+25
View File
@@ -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<ver>-linux-x86_64.zip`) against the published
> `onnxruntime-v<ver>-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 <zip>` + `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<ver>-<sha256:0:16>/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
ragflow_deps/download_deps.py
+59 -2
View File
@@ -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.
ragflow_deps/download_go_deps.py
+61 -2
View File
@@ -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.