diff --git a/docs/src/bundles.md b/docs/src/bundles.md index 99bd1c5..07aecda 100644 --- a/docs/src/bundles.md +++ b/docs/src/bundles.md @@ -51,8 +51,12 @@ a top-level `input_shapes` block (see `write_bundle` below); the [Object Detection](object_detection.md) example compiles one weight set for several image shapes this way. -The per-input batch axis is derived from the position of `n`/`b`; at inference the request's size -along that axis must equal one of `batching.compiled_batch_sizes`. Each tensor parses into a +The per-input batch axis is derived from the position of `n`/`b`. When every executable input and +output has a batch axis, the scheduler coalesces queued requests along it and pads the dispatch +with zero rows up to the smallest compiled size that fits, so a request may carry any number of +rows up to the largest of `batching.compiled_batch_sizes`; otherwise the request's size along that +axis must equal one of them. `runtime.batch_sizes: largest` loads only the largest size (see +[Node configuration](node_config.md)). Each tensor parses into a [`TensorSpec`](@ref) with a [`Dim`](@ref) per axis, and the compiled sizes form the [`BatchingSpec`](@ref). The writer stamps `format_version: "2.0"`; the loader accepts `2.0` or `2`. diff --git a/docs/src/node_config.md b/docs/src/node_config.md index dea41a2..ee041b5 100644 --- a/docs/src/node_config.md +++ b/docs/src/node_config.md @@ -73,6 +73,9 @@ global: allow_cpu_fallback: false numerics: auto # f32 | auto | tf32; see below executable_cache: true # cache compiled programs under each bundle's .cache/; see below + profile: default # default | regulated; see below + batch_sizes: all # all | largest compiled batch size per model (profile default) + xla_flags: {} # XLA DebugOptions overrides, e.g. xla_gpu_deterministic_ops: true weight_cache_fraction: 1.0 # arena fraction for all weights (pinned + on-demand); 0 disables weight_cache_wiggle_fraction: 0.1 # arena fraction kept free; drives startup auto-sizing autotune: true # XLA GPU compile autotuner; false = default kernels, no trials @@ -111,7 +114,9 @@ change model outputs; the cost is tensor-core throughput for f32 matmuls on TF32 `tf32` compiles exactly like `auto` (TF32 is permitted and XLA/cuBLAS pick the kernels; StableHLO has no way to force TF32 for convolutions) but turns the hardware requirement into a guarantee: startup fails on hardware that cannot run TF32, so a mixed fleet cannot silently serve divergent -numerics. On CUDA workers a startup probe logs whether TF32 arithmetic is actually in use and, +numerics. Startup also fails if `NVIDIA_TF32_OVERRIDE=0` is set, or if the startup probe matmul does +not actually run in TF32, so a capable worker that is not using TF32 cannot serve a deployment +validated on it. On CUDA workers a startup probe logs whether TF32 arithmetic is actually in use and, under `f32`, proves the pin bit-exactly; the per-model outcome (ops pinned, algorithms rewritten or stripped) is recorded in each "model loaded" log line. @@ -125,6 +130,31 @@ them. The directory watcher never reacts to anything under `.cache/`, so cache w the model that produced them. The bundle directory must be writable by the worker; when it is not, the cache logs a warning and every program is compiled as before. +`runtime.batch_sizes` selects which of a bundle's compiled batch sizes are loaded. `all` (the +default) loads every `model.b{N}.mlir` the manifest declares, and the scheduler picks the largest +size a dispatch can fill. `largest` loads only the largest declared size, per input-shape variant, +and never reads the others (they need not even be present). Every dispatch then runs that one +program, padded with zero rows when fewer requests are queued and sliced back per request, so a +row's result no longer depends on how many other requests were coalesced with it; the cost is the +full batch's compute on every dispatch, including a lone one-row request. Unbatched bundles are +unaffected. Workers with different settings can share a bundle directory: a `largest` worker keeps +the cached programs of the sizes it skipped. + +`runtime.xla_flags` is a mapping of XLA `DebugOptions` field names (no leading `--`) to values, +applied to every compile. At startup each name is checked against the XLA linked into the worker +and each value against the field's type (booleans, integers, numbers, strings, and enum fields by +enumerator name), so a typo fails startup instead of being ignored. The flags are part of the +executable cache key. An explicit `xla_gpu_autotune_level` here wins over `autotune: false`. + +`runtime.profile` applies a named set of defaults. `regulated` defaults `batch_sizes` to +`largest` and adds `xla_gpu_exclude_nondeterministic_ops: true` to `xla_flags`; anything set +explicitly wins, so `batch_sizes: all` or `xla_gpu_exclude_nondeterministic_ops: false` turns a +piece back off. The profile does not change `numerics`; set `f32` or `tf32` to match what the +deployment was validated on. The worker states the precision prominently at startup: a bannered +info line under `f32` or `tf32`, both of which are attested by the startup probe, and a bannered +warning under `auto`, naming the precision actually in effect on its device (TF32 on Ampere and +newer, full f32 on older GPUs), because `auto` guarantees neither. + `model_control_mode` sets how the loaded model set evolves: `dynamic` (the default) watches the repository and loads, unloads, reloads, and renames bundles online as files change (a renamed directory with unchanged contents renames the model in place, no recompile); `static` fixes the @@ -305,6 +335,8 @@ overrides were applied, is logged at startup. | `INFERENCE_SERVER_RUNTIME_AUTOTUNE_CACHE_DIR` | `runtime.autotune_cache_dir` | path | | `INFERENCE_SERVER_RUNTIME_NUMERICS` | `runtime.numerics` | `f32` \| `auto` \| `tf32` | | `INFERENCE_SERVER_RUNTIME_EXECUTABLE_CACHE` | `runtime.executable_cache` | bool | +| `INFERENCE_SERVER_RUNTIME_PROFILE` | `runtime.profile` | `default` \| `regulated` | +| `INFERENCE_SERVER_RUNTIME_BATCH_SIZES` | `runtime.batch_sizes` | `all` \| `largest` | | `INFERENCE_SERVER_RUNTIME_SHARED_HOST_WEIGHTS` | `runtime.shared_host_weights` | bool | | `INFERENCE_SERVER_RUNTIME_SHARED_HOST_WEIGHTS_MODE` | `runtime.shared_host_weights_mode` | octal string | | `INFERENCE_SERVER_SCHEDULER_DISCIPLINE` | `scheduler.discipline` | `fair` \| `fifo` \| `edf` | diff --git a/packages/ReactantServer/src/bundle.jl b/packages/ReactantServer/src/bundle.jl index 865b229..4c9e270 100644 --- a/packages/ReactantServer/src/bundle.jl +++ b/packages/ReactantServer/src/bundle.jl @@ -41,27 +41,37 @@ end # Discover the per-batch-size StableHLO modules for one variant prefix. The prefix is `model` for # a single-shape bundle and `model.v{i}` for variant `i` of a multi-shape bundle. A variant has # either per-batch-size files `.b{N}.mlir` (keyed by N) or a single `.mlir` -# (keyed by 0, used for any batch size). -function _discover_batch_modules(dir::AbstractString, m::Manifest, prefix::AbstractString) - modules = Dict{Int, Vector{UInt8}}() +# (keyed by 0, used for any batch size). Under `BATCH_SIZES_LARGEST` only the largest size's file is +# read; the others are neither read nor required to exist. +function _discover_batch_modules( + dir::AbstractString, m::Manifest, prefix::AbstractString; + batch_sizes::BatchSizeMode = BATCH_SIZES_ALL + ) + files = Dict{Int, String}() rx = Regex("^" * replace(prefix, "." => "\\.") * "\\.b(\\d+)\\.mlir\$") for f in readdir(dir) mt = match(rx, f) mt === nothing && continue - modules[parse(Int, mt.captures[1])] = read(joinpath(dir, f)) + files[parse(Int, mt.captures[1])] = joinpath(dir, f) end - if isempty(modules) + if isempty(files) single = joinpath(dir, prefix * ".mlir") isfile(single) || throw(BundleError("bundle '$(m.name)' has no $(prefix).mlir or $(prefix).b{N}.mlir")) - modules[0] = read(single) - return modules + return Dict{Int, Vector{UInt8}}(0 => read(single)) + end + declared = m.batching.compiled_batch_sizes + wanted = if batch_sizes == BATCH_SIZES_LARGEST + [isempty(declared) ? maximum(keys(files)) : maximum(declared)] + else + declared end - for sz in m.batching.compiled_batch_sizes - haskey(modules, sz) || + for sz in wanted + haskey(files, sz) || throw(BundleError("bundle '$(m.name)' declares batch size $sz but has no $(prefix).b$sz.mlir")) end - return modules + keep = batch_sizes == BATCH_SIZES_LARGEST ? wanted : collect(keys(files)) + return Dict{Int, Vector{UInt8}}(sz => read(files[sz]) for sz in keep) end # Discover every variant's StableHLO module(s), keyed by variant. A single-shape bundle (no @@ -69,28 +79,34 @@ end # multi-shape bundle yields one entry per declared `input_shapes` variant `i`, read from # `model.v{i}.*.mlir`; the variant key is the same variable-axis size vector the manifest resolved # and the runtime derives from a request, so dispatch lines up with what was compiled. -function _discover_modules(dir::AbstractString, m::Manifest) +function _discover_modules(dir::AbstractString, m::Manifest; batch_sizes::BatchSizeMode = BATCH_SIZES_ALL) if isempty(m.input_shapes) - return Dict{VariantKey, Dict{Int, Vector{UInt8}}}(VariantKey() => _discover_batch_modules(dir, m, "model")) + return Dict{VariantKey, Dict{Int, Vector{UInt8}}}( + VariantKey() => _discover_batch_modules(dir, m, "model"; batch_sizes) + ) end out = Dict{VariantKey, Dict{Int, Vector{UInt8}}}() for (i, vkey) in enumerate(m.input_shapes) - out[vkey] = _discover_batch_modules(dir, m, "model.v$(i - 1)") + out[vkey] = _discover_batch_modules(dir, m, "model.v$(i - 1)"; batch_sizes) end return out end """ - load_bundle_entry(dir; validator=NullSignatureValidator()) -> ModelEntry + load_bundle_entry(dir; validator=NullSignatureValidator(), batch_sizes=BATCH_SIZES_ALL) -> ModelEntry Parse and validate the bundle directory `dir` into an uncompiled `ModelEntry` (its `executable` and `sched` slots are `nothing`). The model's name is the directory's basename: renaming the directory renames the model. A `name` declared in the manifest is informational and ignored (the directory name is injected before parsing so every downstream consumer, including error messages and the metadata RPC, agrees on the served name). Used by both `load_bundles` and the directory -watcher (see watcher.jl) to load a single bundle. +watcher (see watcher.jl) to load a single bundle. `batch_sizes` selects which compiled batch +sizes are read (see [`BatchSizeMode`](@ref)). """ -function load_bundle_entry(dir::AbstractString; validator::SignatureValidator = NullSignatureValidator()) +function load_bundle_entry( + dir::AbstractString; validator::SignatureValidator = NullSignatureValidator(), + batch_sizes::BatchSizeMode = BATCH_SIZES_ALL + ) manifest_path = joinpath(dir, "manifest.yaml") raw = YAML.load_file(manifest_path; dicttype = Dict{String, Any}) raw isa AbstractDict || throw(BundleError("manifest in $dir is not a mapping")) @@ -108,7 +124,7 @@ function load_bundle_entry(dir::AbstractString; validator::SignatureValidator = return MetaEntry(m.name, m, m.meta_calls, mreg.run) end - mlir_bytes = _discover_modules(dir, m) + mlir_bytes = _discover_modules(dir, m; batch_sizes) weights_path = joinpath(dir, "weights.safetensors") isfile(weights_path) || throw(BundleError("bundle '$(m.name)' missing weights.safetensors")) @@ -126,8 +142,11 @@ function load_bundle_entry(dir::AbstractString; validator::SignatureValidator = return ModelEntry(m.name, m, mlir_bytes, weights_path, weights, nothing, nothing, pre, post) end -function _load_one_bundle!(reg::ModelRegistry, dir::AbstractString, validator::SignatureValidator) - entry = load_bundle_entry(dir; validator = validator) +function _load_one_bundle!( + reg::ModelRegistry, dir::AbstractString, validator::SignatureValidator, + batch_sizes::BatchSizeMode + ) + entry = load_bundle_entry(dir; validator, batch_sizes) (haskey(reg.by_name, entry.name) || haskey(reg.meta, entry.name)) && throw(BundleError("duplicate model name '$(entry.name)'")) if entry isa MetaEntry @@ -139,7 +158,7 @@ function _load_one_bundle!(reg::ModelRegistry, dir::AbstractString, validator::S end """ - load_bundles(model_dirs; validator=NullSignatureValidator(), include=nothing) -> ModelRegistry + load_bundles(model_dirs; validator=NullSignatureValidator(), include=nothing, batch_sizes=BATCH_SIZES_ALL) -> ModelRegistry Discover every subdirectory containing a manifest.yaml under each model dir, load and validate it, and register it. The runtime fills each entry's executable slot afterwards. @@ -147,12 +166,13 @@ validate it, and register it. The runtime fills each entry's executable slot aft When `include` is a non-empty collection of model names, only bundles whose directory name is in the set are loaded. The directory name IS the model name (see `load_bundle_entry`), so filtering by directory avoids parsing skipped manifests. Names in `include` that are not found -in any model dir produce a warning. +in any model dir produce a warning. `batch_sizes` is passed to every [`load_bundle_entry`](@ref). """ function load_bundles( model_dirs::AbstractVector{<:AbstractString}; validator::SignatureValidator = NullSignatureValidator(), - include = nothing + include = nothing, + batch_sizes::BatchSizeMode = BATCH_SIZES_ALL ) want = include === nothing ? nothing : Set{String}(String(x) for x in include) reg = ModelRegistry() @@ -166,7 +186,7 @@ function load_bundles( if want !== nothing && !(name in want) continue end - _load_one_bundle!(reg, child, validator) + _load_one_bundle!(reg, child, validator, batch_sizes) push!(found, name) end end diff --git a/packages/ReactantServer/src/runtime/backend.jl b/packages/ReactantServer/src/runtime/backend.jl index e0b598b..191d557 100644 --- a/packages/ReactantServer/src/runtime/backend.jl +++ b/packages/ReactantServer/src/runtime/backend.jl @@ -21,6 +21,13 @@ function device_ordinal end # (backend, device) -> Int # and load the compiled program from the bundle's `.cache/` when present instead of compiling. function compile_artifact end # (backend, pool, mlir_bytes, num_parameters, num_outputs; numerics_stats, cache) -> executable +# Check `runtime.xla_flags` (XLA DebugOptions field name => value) and convert it to the form +# compile_artifact applies: name/value pairs sorted by name. A backend that compiles through XLA +# rejects unknown names and values of the wrong type; this default only converts, since a backend +# with no XLA (MockBackend) has nothing to check against. +resolve_xla_flags(::AbstractBackend, flags::AbstractDict) = + sort!(Pair{Symbol, Any}[Symbol(k) => v for (k, v) in flags]; by = first) + # Whether compile_artifact can serialize compiled programs into the per-bundle executable cache. # Default false (MockBackend has nothing to cache); the Reactant backend always can (see # xla_serialization.jl for where the bindings come from). diff --git a/packages/ReactantServer/src/runtime/executable_cache.jl b/packages/ReactantServer/src/runtime/executable_cache.jl index 3b4c94f..07a6e78 100644 --- a/packages/ReactantServer/src/runtime/executable_cache.jl +++ b/packages/ReactantServer/src/runtime/executable_cache.jl @@ -119,17 +119,29 @@ Record the bundle's current MLIR source hashes (file name to sha256) in the cach delete every cached executable whose source changed hash or is no longer part of the bundle. Returns the names of the invalidated sources. Weights are not part of the record. Fails open: an I/O error logs a warning and returns an empty list. + +`retain` names sources that are still part of the bundle but were not read by this worker +(`runtime.batch_sizes: largest` skips the smaller batch sizes). Their recorded hashes and cached +programs are kept as they are, so a worker loading only the largest size does not delete programs +that a worker loading every size, on the same bundle directory, still uses. """ -function sync_mlir_hashes!(cache_dir::AbstractString, sources::Dict{String, String}) +function sync_mlir_hashes!( + cache_dir::AbstractString, sources::Dict{String, String}; + retain = () + ) try prev = _read_hashes(cache_dir) - stale = String[get(sources, src, nothing) == sha ? "" : src for (src, sha) in prev] + record = copy(sources) + for src in retain + haskey(prev, src) && !haskey(record, src) && (record[src] = prev[src]) + end + stale = String[get(record, src, nothing) == sha ? "" : src for (src, sha) in prev] filter!(!isempty, stale) for src in stale n = _sweep_entries!(cache_dir, src) @info "executable cache: MLIR source changed; dropped cached programs" cache_dir source = src entries = n end - prev == sources || _write_hashes(cache_dir, sources) + prev == record || _write_hashes(cache_dir, record) return stale catch err @warn "executable cache: could not sync MLIR hashes (cache disabled for this bundle)" cache_dir exception = err @@ -166,6 +178,14 @@ end "Remove an entry that failed to load, so the next start recompiles instead of retrying it." drop_entry(path::AbstractString) = (rm(path; force = true); nothing) +# Every per-batch-size module file the manifest declares, across all variants. A declared file this +# worker did not read is still part of the bundle, so its cache record is retained (see +# `sync_mlir_hashes!`). +function _declared_module_files(m::Manifest) + vkeys = isempty(m.input_shapes) ? [VariantKey()] : m.input_shapes + return Set{String}(module_filename(m, vkey, sz) for vkey in vkeys for sz in m.batching.compiled_batch_sizes) +end + # Build the per-module cache slots for a bundle entry: sync the hash record, then one slot per # (variant, batch size). Returns `nothing` when the bundle cannot host a cache. function executable_cache_slots(entry::ModelEntry) @@ -180,7 +200,7 @@ function executable_cache_slots(entry::ModelEntry) sources[f] = sha shas[(vkey, sz)] = (f, sha) end - sync_mlir_hashes!(cache_dir, sources) + sync_mlir_hashes!(cache_dir, sources; retain = _declared_module_files(entry.manifest)) return Dict{Tuple{VariantKey, Int}, ExecutableCacheSlot}( k => ExecutableCacheSlot(cache_dir, f, sha) for (k, (f, sha)) in shas ) diff --git a/packages/ReactantServer/src/runtime/memory_pool.jl b/packages/ReactantServer/src/runtime/memory_pool.jl index 7ed43df..03972c0 100644 --- a/packages/ReactantServer/src/runtime/memory_pool.jl +++ b/packages/ReactantServer/src/runtime/memory_pool.jl @@ -12,16 +12,22 @@ mutable struct MemoryPool ctx::Any # backend compilation context (Reactant) or nothing autotune::Bool # runtime.autotune; false => compile with the GPU autotuner disabled numerics::NumericsMode # runtime.numerics; f32 matmul/conv precision policy (see NumericsMode) + xla_flags::Vector{Pair{Symbol, Any}} # runtime.xla_flags, checked by resolve_xla_flags; sorted by name end -# Back-compat constructors: default autotune on and hardware-adaptive numerics, so existing call -# sites and test mocks that pass the original five or six fields keep the current behavior. +# Back-compat constructors: default autotune on, hardware-adaptive numerics, and no XLA flags, so +# existing call sites and test mocks that pass the original five to seven fields keep the current +# behavior. MemoryPool(backend::AbstractBackend, client, device, platform::String, ctx) = MemoryPool(backend, client, device, platform, ctx, true, NUMERICS_AUTO) MemoryPool(backend::AbstractBackend, client, device, platform::String, ctx, autotune::Bool) = MemoryPool(backend, client, device, platform, ctx, autotune, NUMERICS_AUTO) +MemoryPool(backend::AbstractBackend, client, device, platform::String, ctx, autotune::Bool, numerics::NumericsMode) = + MemoryPool(backend, client, device, platform, ctx, autotune, numerics, Pair{Symbol, Any}[]) function resolve_client(backend::AbstractBackend, cfg::RuntimeConfig) + # Checked before the client exists, so a misspelled or mistyped flag fails startup immediately. + xla_flags = resolve_xla_flags(backend, cfg.xla_flags) platform = cfg.backend == CUDA_BACKEND ? "cuda" : "cpu" try client = make_client( @@ -29,7 +35,7 @@ function resolve_client(backend::AbstractBackend, cfg::RuntimeConfig) autotune_cache = cfg.autotune_cache, autotune_cache_dir = cfg.autotune_cache_dir ) device = select_device(backend, client, cfg.device_ordinal) - return MemoryPool(backend, client, device, platform, make_context(backend), cfg.autotune, cfg.numerics) + return MemoryPool(backend, client, device, platform, make_context(backend), cfg.autotune, cfg.numerics, xla_flags) catch err if cfg.backend == CUDA_BACKEND && cfg.allow_cpu_fallback @warn "CUDA backend unavailable; falling back to CPU" exception = (err, catch_backtrace()) @@ -38,7 +44,7 @@ function resolve_client(backend::AbstractBackend, cfg::RuntimeConfig) autotune_cache = cfg.autotune_cache, autotune_cache_dir = cfg.autotune_cache_dir ) device = select_device(backend, client, 0) - return MemoryPool(backend, client, device, "cpu", make_context(backend), cfg.autotune, cfg.numerics) + return MemoryPool(backend, client, device, "cpu", make_context(backend), cfg.autotune, cfg.numerics, xla_flags) end rethrow() end diff --git a/packages/ReactantServer/src/runtime/reactant_backend.jl b/packages/ReactantServer/src/runtime/reactant_backend.jl index 23d705e..3e40ced 100644 --- a/packages/ReactantServer/src/runtime/reactant_backend.jl +++ b/packages/ReactantServer/src/runtime/reactant_backend.jl @@ -212,14 +212,57 @@ end # When autotuning is disabled, force xla_gpu_autotune_level=0: XLA uses default gemm/conv algorithm # selection with no device timing trials. This removes the autotuner's run-to-run non-determinism and # the compile-time scratch that otherwise inflates the startup memory probe on the first (un-cached) -# start. When enabled, pass no override so the compile is byte-identical to the previous behavior. +# start. `runtime.xla_flags` is layered on top, so an explicit xla_gpu_autotune_level there wins. +# With neither, pass no override so the compile is byte-identical to the previous behavior. function _compile_options(pool::MemoryPool, device_id::Int) - pool.autotune && return _RXLA.make_compile_options(; device_id = Int64(device_id)) + pool.autotune && isempty(pool.xla_flags) && + return _RXLA.make_compile_options(; device_id = Int64(device_id)) + base = pool.autotune ? (;) : (; xla_gpu_autotune_level = Int32(0)) return _RXLA.make_compile_options(; - device_id = Int64(device_id), xla_debug_options = (; xla_gpu_autotune_level = Int32(0)) + device_id = Int64(device_id), xla_debug_options = merge(base, NamedTuple(pool.xla_flags)) ) end +# runtime.xla_flags against the DebugOptions proto of the linked XLA: an unknown name or a value +# that does not fit the field's type is a ConfigError at startup, rather than a Setfield error on the +# first compile or, worse, a flag that is silently not what the operator meant. +function resolve_xla_flags(::ReactantBackend, flags::AbstractDict) + T = Reactant.Proto.xla.DebugOptions + out = Pair{Symbol, Any}[] + for (k, v) in flags + name = Symbol(k) + hasfield(T, name) || + throw(ConfigError("runtime.xla_flags.$k is not an XLA DebugOptions field in Reactant $(pkgversion(Reactant))")) + push!(out, name => _xla_flag_value(fieldtype(T, name), String(k), v)) + end + return sort!(out; by = first) +end + +_xla_flag_error(k, v, what) = throw(ConfigError("runtime.xla_flags.$k must be $what, got $(repr(v))")) +_xla_flag_value(::Type{Bool}, k, v) = v isa Bool ? v : _xla_flag_error(k, v, "a boolean") +function _xla_flag_value(::Type{T}, k, v) where {T <: Integer} + (v isa Integer && !(v isa Bool)) || _xla_flag_error(k, v, "an integer") + typemin(T) <= v <= typemax(T) || _xla_flag_error(k, v, "an integer that fits $T") + return T(v) +end +_xla_flag_value(::Type{T}, k, v) where {T <: AbstractFloat} = + (v isa Real && !(v isa Bool)) ? T(v) : _xla_flag_error(k, v, "a number") +_xla_flag_value(::Type{String}, k, v) = v isa AbstractString ? String(v) : _xla_flag_error(k, v, "a string") +# Enum fields take the enumerator's name; repeated and message fields are not supported. +function _xla_flag_value(::Type{T}, k, v) where {T} + names = try + [string(Symbol(x)) => x for x in instances(T)] + catch + throw(ConfigError("runtime.xla_flags.$k has type $T, which runtime.xla_flags cannot set")) + end + if v isa AbstractString + for (n, x) in names + n == v && return x + end + end + return _xla_flag_error(k, v, "one of " * join(first.(names), ", ")) +end + # ── The executable cache (executable_cache.jl) on the Reactant backend ─────────────────────────── # Serialize a (rewritten) module back to a StableHLO portable artifact at the current version. These @@ -259,7 +302,9 @@ end function _cache_policy(pool::MemoryPool, tf32_capable::Bool) return string( "reactant=", pkgversion(Reactant), ";autotune=", pool.autotune, ";numerics=", pool.numerics, - ";tf32=", tf32_capable, ";format=", EXEC_CACHE_FORMAT + ";tf32=", tf32_capable, ";format=", EXEC_CACHE_FORMAT, + # Only when set, so programs cached before runtime.xla_flags existed keep their keys. + isempty(pool.xla_flags) ? "" : ";xla=" * join([string(k, "=", v) for (k, v) in pool.xla_flags], ",") ) end diff --git a/packages/ReactantServer/src/runtime/tf32.jl b/packages/ReactantServer/src/runtime/tf32.jl index c0fea38..66a9a5c 100644 --- a/packages/ReactantServer/src/runtime/tf32.jl +++ b/packages/ReactantServer/src/runtime/tf32.jl @@ -330,7 +330,7 @@ end function _run_probe_leg(backend, pool::MemoryPool, numerics::NumericsMode) legpool = MemoryPool( pool.backend, pool.client, pool.device, pool.platform, pool.ctx, - pool.autotune, numerics + pool.autotune, numerics, pool.xla_flags ) exec = compile_artifact(backend, legpool, _probe_artifact(pool.ctx), 2, 1) a = b = nothing @@ -355,24 +355,43 @@ end """ tf32_probe(backend, pool) -> (; tf32_active, pinned_exact) -Startup numerics attestation. Leg 1 (informational) compiles the probe matmul at DEFAULT -precision (`auto` semantics) and reports whether TF32 arithmetic was actually used by this -worker's hardware+stack. Leg 2 runs only under `numerics = f32`: it compiles through the -as-configured pool (the real `pin_f32!` production path) and **throws** unless the result is -bitwise-exact f32; a pin that does not hold is a bug, not a tolerance. Results are logged; the -returned fields are `true`/`false`, or `nothing` for a leg that did not run or was indeterminate. +Startup numerics attestation. Leg 1 compiles the probe matmul at DEFAULT precision and reports +whether TF32 arithmetic was actually used by this worker's hardware+stack. Under `auto` it is +informational. Under `numerics = tf32` it compiles through the as-configured pool and **throws** +unless TF32 was observed, so a worker that is capable of TF32 but not using it (for example with +`NVIDIA_TF32_OVERRIDE=0` in its environment) cannot serve a deployment validated on TF32. Leg 2 +runs only under `numerics = f32`: it compiles through the as-configured pool (the real `pin_f32!` +production path) and **throws** unless the result is bitwise-exact f32; a pin that does not hold is +a bug, not a tolerance. Results are logged; the returned fields are `true`/`false`, or `nothing` +for a leg that did not run or was indeterminate. """ function tf32_probe(backend, pool::MemoryPool) A = _probe_sentinel_matrix() + required = pool.numerics == NUMERICS_TF32 tf32_active = nothing try - C = _run_probe_leg(backend, pool, NUMERICS_AUTO) + C = _run_probe_leg(backend, pool, required ? NUMERICS_TF32 : NUMERICS_AUTO) tf32_active = C == A ? false : (all(==(1.0f0), C) ? true : nothing) - tf32_active === nothing && + tf32_active === nothing && !required && @warn "TF32 probe: DEFAULT-precision leg returned neither exact-f32 nor the TF32 signature" platform = pool.platform catch err + required && throw( + ErrorException( + "numerics=tf32 attestation failed: the TF32 probe matmul could not run on " * + "this device ($(pool.platform)): $(sprint(showerror, err))" + ) + ) @warn "TF32 probe: DEFAULT-precision leg failed; skipping detection" exception = (err, catch_backtrace()) end + if required && tf32_active !== true + observed = tf32_active === false ? "full f32" : "neither exact f32 nor the TF32 signature" + throw( + ErrorException( + "numerics=tf32 attestation failed: the probe matmul ran in $observed on this " * + "device ($(pool.platform)), not TF32; check NVIDIA_TF32_OVERRIDE and the XLA flags" + ) + ) + end pinned_exact = nothing if pool.numerics == NUMERICS_F32 # No try/catch: under numerics=f32 the attestation is load-bearing; a pin that cannot be diff --git a/packages/ReactantServer/src/server.jl b/packages/ReactantServer/src/server.jl index 0955d92..5ba0443 100644 --- a/packages/ReactantServer/src/server.jl +++ b/packages/ReactantServer/src/server.jl @@ -133,6 +133,60 @@ function _warn_unenforced_config(cfg::ServerConfig) return nothing end +# The regulated profile deliberately leaves `numerics` alone, so it states the precision actually in +# effect instead, prominently, for the deployment record. `f32` and `tf32` are both attested at +# startup (tf32_probe throws otherwise), so they get a bannered info line naming the precision. Under +# `auto` nothing is guaranteed: the same bundle computes f32 matmuls/convolutions in TF32 on Ampere +# and newer GPUs and in full f32 on older ones, so that case is a bannered warning. +function _report_regulated_precision(rt::RuntimeConfig, backend::AbstractBackend, pool::MemoryPool, probe) + flags = isempty(pool.xla_flags) ? "none" : join([string(k, "=", v) for (k, v) in pool.xla_flags], ", ") + banner = "*"^100 + if rt.numerics == NUMERICS_F32 + @info """ + $banner + REGULATED PROFILE: runtime.numerics = f32 + f32 matmul/convolution precision is pinned to full f32 and was attested bit-exact at startup. + $banner""" platform = pool.platform batch_sizes = rt.batch_sizes xla_flags = flags + return nothing + elseif rt.numerics == NUMERICS_TF32 + @info """ + $banner + REGULATED PROFILE: runtime.numerics = tf32 + TF32 is required and the startup probe confirmed it is active on this worker. XLA and cuBLAS + still choose the kernel per op (StableHLO cannot force TF32 for convolutions), so results + match the validation only on the same GPU class and software stack. + $banner""" platform = pool.platform batch_sizes = rt.batch_sizes xla_flags = flags + return nothing + end + tf32 = probe === nothing ? backend_tf32_capable(backend, pool) : something(probe.tf32_active, backend_tf32_capable(backend, pool)) + effective = tf32 ? + "TF32 (10-bit mantissa inputs, f32 accumulate)" : + "full f32 on this device only; a TF32-capable GPU would use TF32" + @warn """ + $banner + REGULATED PROFILE WITHOUT PINNED NUMERICS: runtime.numerics = auto + f32 matmul/convolution precision in effect on this worker: $effective + Results depend on the GPU generation. Set runtime.numerics to f32 or tf32 to make the precision + a startup requirement. + $banner""" platform = pool.platform batch_sizes = rt.batch_sizes xla_flags = flags + return nothing +end + +# numerics=tf32: NVIDIA_TF32_OVERRIDE=0 disables TF32 inside cuBLAS/cuDNN for the whole process, so +# a worker started with it would compute in a precision the deployment was not validated on. The +# probe would catch it too; checking the environment first gives the operator the actual cause. +function _assert_tf32_not_overridden(env = ENV) + v = get(env, "NVIDIA_TF32_OVERRIDE", nothing) + (v !== nothing && strip(v) == "0") && + throw( + ErrorException( + "runtime.numerics 'tf32' requested but NVIDIA_TF32_OVERRIDE=0 is set in the " * + "environment, which disables TF32 in cuBLAS/cuDNN; unset it or use 'f32'" + ) + ) + return nothing +end + function _bring_up(cfg::ServerConfig, backend::AbstractBackend) _warn_unenforced_config(cfg) # numerics=f32 defense in depth: NVIDIA_TF32_OVERRIDE must be in the environment before the @@ -140,6 +194,7 @@ function _bring_up(cfg::ServerConfig, backend::AbstractBackend) # the op-level precision pin cannot see; it does NOT govern XLA's Triton GEMMs, which the pin # does, so the two mechanisms cover each other's gaps. cfg.runtime.numerics == NUMERICS_F32 && (ENV["NVIDIA_TF32_OVERRIDE"] = "0") + cfg.runtime.numerics == NUMERICS_TF32 && _assert_tf32_not_overridden() pool = resolve_client(backend, cfg.runtime) backend = pool.backend # the CPU fallback may have swapped the backend use_exec_cache = cfg.runtime.executable_cache && supports_executable_cache(backend) @@ -158,12 +213,13 @@ function _bring_up(cfg::ServerConfig, backend::AbstractBackend) ) ) end - # Numerics attestation, GPU only: report whether TF32 is actually in use (auto/tf32) and, under - # f32, prove the precision pin bit-exactly. Runs before any model compile and before the scratch + # Numerics attestation, GPU only: report whether TF32 is actually in use (auto), require it + # (tf32), or prove the f32 precision pin bit-exactly (f32). Runs before any model compile and before the scratch # high-water probe; see tf32_probe for why its transient ~3 MB cannot perturb that measurement. - pool.platform == "cuda" && tf32_probe(backend, pool) + probe = pool.platform == "cuda" ? tf32_probe(backend, pool) : nothing + cfg.runtime.profile == PROFILE_REGULATED && _report_regulated_precision(cfg.runtime, backend, pool, probe) include = isempty(cfg.models_include) ? nothing : cfg.models_include - registry = load_bundles(cfg.model_dirs; include = include) + registry = load_bundles(cfg.model_dirs; include = include, batch_sizes = cfg.runtime.batch_sizes) isempty(registry.by_name) && @warn "no model bundles found" model_dirs = cfg.model_dirs models_include = cfg.models_include # On-demand weight residency is sized as a fraction of the BFC arena (`mem_fraction * device`), # resolved now that the device pool exists. The cache is GPU-only: `arena` is 0 when the device diff --git a/packages/ReactantServer/src/watcher.jl b/packages/ReactantServer/src/watcher.jl index 17c7446..6b88b5d 100644 --- a/packages/ReactantServer/src/watcher.jl +++ b/packages/ReactantServer/src/watcher.jl @@ -134,7 +134,8 @@ function _apply_change!( delete!(w.seen, name) delete!(w.dir_ids, name) else - entry = load_bundle_entry(dir) # named by its directory basename, i.e. `name` + # named by its directory basename, i.e. `name` + entry = load_bundle_entry(dir; batch_sizes = w.cfg.runtime.batch_sizes) if entry isa MetaEntry # Meta bundles need no compilation; register them straight into the meta map. put_meta!(w.scheduler, entry) diff --git a/packages/ReactantServer/test/runtests.jl b/packages/ReactantServer/test/runtests.jl index 8233dde..4cf5b27 100644 --- a/packages/ReactantServer/test/runtests.jl +++ b/packages/ReactantServer/test/runtests.jl @@ -45,6 +45,7 @@ include("grpc_helpers.jl") include("test_watcher.jl") include("test_xla_serialization.jl") include("test_executable_cache.jl") + include("test_batch_sizes.jl") include("test_shared_memory.jl") include("test_kaimon_ext.jl") end diff --git a/packages/ReactantServer/test/test_batch_sizes.jl b/packages/ReactantServer/test/test_batch_sizes.jl new file mode 100644 index 0000000..df943d0 --- /dev/null +++ b/packages/ReactantServer/test/test_batch_sizes.jl @@ -0,0 +1,164 @@ +# runtime.batch_sizes and runtime.xla_flags on the worker: which compiled batch sizes a bundle load +# reads, how the executable cache treats the sizes it skipped, and how XLA flags are checked and +# applied. The compile tests run on Reactant's CPU client. + +using ReactantServer: load_bundle_entry, executable_cache_slots, sync_mlir_hashes!, entry_path, + store_entry, ExecutableCacheSlot, EXEC_CACHE_HASHES_FILE, sha256hex, resolve_xla_flags, + _compile_options, _cache_policy, _coalesce_inputs, _slice_outputs, BATCH_SIZES_ALL, + BATCH_SIZES_LARGEST, ConfigError, BundleError, VariantKey + +# y = x .* w over a (4, n) batch, compiled per batch size n. StableHLO is row-major, so the Julia +# (4, n) tensor is tensor. +function _bs_mlir(n::Int) + return """ + module { + func.func @main(%x: tensor<$(n)x4xf32>, %w: tensor<4xf32>) -> tensor<$(n)x4xf32> { + %0 = stablehlo.broadcast_in_dim %w, dims = [1] : (tensor<4xf32>) -> tensor<$(n)x4xf32> + %1 = stablehlo.multiply %x, %0 : tensor<$(n)x4xf32> + return %1 : tensor<$(n)x4xf32> + } + } + """ +end + +function _write_multi_batch_bundle(root, name; sizes = [1, 4, 8], files = sizes) + dir = joinpath(root, name) + mkpath(dir) + write( + joinpath(dir, "manifest.yaml"), """ + format_version: "2.0" + name: $name + executable_inputs: + - {name: x, dtype: f32, shape: cn, dims: {c: 4}} + executable_outputs: + - {name: y, dtype: f32, shape: cn, dims: {c: 4}} + batching: {compiled_batch_sizes: [$(join(sizes, ", "))]} + """ + ) + for n in files + write(joinpath(dir, "model.b$n.mlir"), stablehlo_artifact(_bs_mlir(n))) + end + SafeTensors.serialize( + joinpath(dir, "weights.safetensors"), Dict("w" => Float32[1, 2, 3, 4]), + Dict("argument_order" => JSON3.write(["w"])) + ) + return dir +end + +_loaded_sizes(entry) = sort!(collect(keys(entry.mlir_bytes[VariantKey()]))) + +@testset "batch_sizes: which modules a bundle load reads" begin + mktempdir() do root + dir = _write_multi_batch_bundle(root, "bs") + @test _loaded_sizes(load_bundle_entry(dir)) == [1, 4, 8] + @test _loaded_sizes(load_bundle_entry(dir; batch_sizes = BATCH_SIZES_LARGEST)) == [8] + + # Under largest, the smaller declared files are not required at all; under all they are. + partial = _write_multi_batch_bundle(root, "partial"; files = [8]) + @test _loaded_sizes(load_bundle_entry(partial; batch_sizes = BATCH_SIZES_LARGEST)) == [8] + @test_throws BundleError load_bundle_entry(partial) + + # The largest declared size must exist. + nolargest = _write_multi_batch_bundle(root, "nolargest"; files = [1, 4]) + @test_throws BundleError load_bundle_entry(nolargest; batch_sizes = BATCH_SIZES_LARGEST) + + # load_bundles passes the mode to every bundle. + reg = ReactantServer.load_bundles([root]; include = ["bs"], batch_sizes = BATCH_SIZES_LARGEST) + @test _loaded_sizes(ReactantServer.get_model(reg, "bs")) == [8] + end +end + +@testset "batch_sizes: largest keeps the cached programs of the sizes it skipped" begin + mktempdir() do root + dir = _write_multi_batch_bundle(root, "bs") + cache = joinpath(dir, ".cache") + # A worker loading every size has cached one program per size. + full = executable_cache_slots(load_bundle_entry(dir)) + programs = Dict( + sz => entry_path(full[(VariantKey(), sz)], "target", "k"^64) for sz in (1, 4, 8) + ) + for p in values(programs) + @test store_entry(p, UInt8[1, 2, 3]) + end + record = ReactantServer._read_hashes(cache) + @test sort!(collect(keys(record))) == ["model.b1.mlir", "model.b4.mlir", "model.b8.mlir"] + + # A worker loading only the largest size, on the same directory, deletes nothing. + slots = executable_cache_slots(load_bundle_entry(dir; batch_sizes = BATCH_SIZES_LARGEST)) + @test collect(keys(slots)) == [(VariantKey(), 8)] + @test all(isfile, values(programs)) + @test ReactantServer._read_hashes(cache) == record + + # A file that really left the bundle (undeclared and gone) is still swept as before. + @test sync_mlir_hashes!(cache, Dict("model.b8.mlir" => record["model.b8.mlir"])) == + ["model.b1.mlir", "model.b4.mlir"] + @test !isfile(programs[1]) && !isfile(programs[4]) && isfile(programs[8]) + end +end + +@testset "xla_flags: checked against DebugOptions" begin + backend = ReactantServer.ReactantBackend() + flags = resolve_xla_flags( + backend, Dict{String, Any}( + "xla_gpu_exclude_nondeterministic_ops" => true, + "xla_gpu_autotune_level" => 2, + "xla_gpu_experimental_autotune_cache_mode" => "AUTOTUNE_CACHE_MODE_READ", + ) + ) + @test first.(flags) == [ + :xla_gpu_autotune_level, :xla_gpu_exclude_nondeterministic_ops, + :xla_gpu_experimental_autotune_cache_mode, + ] + @test flags[1].second === Int32(2) # converted to the proto field's type + @test flags[2].second === true + @test string(Symbol(flags[3].second)) == "AUTOTUNE_CACHE_MODE_READ" + + @test_throws ConfigError resolve_xla_flags(backend, Dict{String, Any}("xla_gpu_no_such_flag" => true)) + @test_throws ConfigError resolve_xla_flags(backend, Dict{String, Any}("xla_gpu_deterministic_ops" => 1)) + @test_throws ConfigError resolve_xla_flags(backend, Dict{String, Any}("xla_gpu_autotune_level" => true)) + @test_throws ConfigError resolve_xla_flags(backend, Dict{String, Any}("xla_gpu_autotune_level" => 2^40)) + @test_throws ConfigError resolve_xla_flags( + backend, Dict{String, Any}("xla_gpu_experimental_autotune_cache_mode" => "NOPE") + ) + + # The mock backend has no XLA to check against; it only converts and sorts. + @test resolve_xla_flags(ReactantServer.MockBackend(), Dict{String, Any}("b" => 1, "a" => true)) == + [:a => true, :b => 1] +end + +@testset "xla_flags + largest: compile, cache key, padded execution (CPU)" begin + backend = ReactantServer.ReactantBackend() + base = ReactantServer.RuntimeConfig(ReactantServer.CPU_BACKEND, 0, 0.9, true, true) + flagged = ReactantServer.RuntimeConfig( + (getfield(base, f) for f in fieldnames(ReactantServer.RuntimeConfig)[1:(end - 3)])..., + ReactantServer.PROFILE_REGULATED, BATCH_SIZES_LARGEST, + Dict{String, Any}("xla_gpu_exclude_nondeterministic_ops" => true), + ) + pool0 = ReactantServer.resolve_client(backend, base) + pool = ReactantServer.resolve_client(backend, flagged) + @test pool.xla_flags == [:xla_gpu_exclude_nondeterministic_ops => true] + + # The flag reaches the compile options, and splits the executable cache key; no flags leaves + # the key exactly as it was. + opts = _compile_options(pool, 0) + @test opts.executable_build_options.debug_options.xla_gpu_exclude_nondeterministic_ops + @test !occursin("xla=", _cache_policy(pool0, false)) + @test endswith(_cache_policy(pool, false), ";xla=xla_gpu_exclude_nondeterministic_ops=true") + + mktempdir() do root + dir = _write_multi_batch_bundle(root, "bs") + entry = load_bundle_entry(dir; batch_sizes = flagged.batch_sizes) + entry.executable = ReactantServer.build_loaded_model(backend, pool, entry; executable_cache = false) + @test ReactantServer._all_batch_sizes(entry.executable) == [8] + + # Three one-row requests pad to the one compiled size and each gets its own row back. + xs = [Float32[i, i, i, i][:, :] for i in 1:3] + pres = [[ReactantServer.NamedTensor("x", x)] for x in xs] + merged = _coalesce_inputs(entry, pres, 3, 8) + @test size(merged[1].data) == (4, 8) + out = ReactantServer.run_model(backend, pool, entry.executable, merged) + for i in 1:3 + @test _slice_outputs(entry, out, i - 1, 1)[1].data == xs[i] .* Float32[1, 2, 3, 4] + end + end +end diff --git a/packages/ReactantServer/test/test_tf32.jl b/packages/ReactantServer/test/test_tf32.jl index daaa3d5..abe9477 100644 --- a/packages/ReactantServer/test/test_tf32.jl +++ b/packages/ReactantServer/test/test_tf32.jl @@ -168,3 +168,44 @@ end res2 = RS.tf32_probe(backend, pool_f32) @test res2.pinned_exact === true end + +@testset "tf32 probe enforces numerics=tf32" begin + backend = RS.ReactantBackend() + pool = RS.resolve_client(backend, RS.RuntimeConfig(RS.CPU_BACKEND, 0, 0.9, true, true)) + # The CPU computes the probe in plain f32, so a tf32 requirement must fail, naming what it saw. + pool_tf32 = RS.MemoryPool( + pool.backend, pool.client, pool.device, pool.platform, pool.ctx, + pool.autotune, RS.NUMERICS_TF32 + ) + err = try + RS.tf32_probe(backend, pool_tf32) + nothing + catch e + e + end + @test err isa ErrorException + @test occursin("numerics=tf32 attestation failed", err.msg) && occursin("full f32", err.msg) + + # NVIDIA_TF32_OVERRIDE=0 is refused up front; any other value (or none) is fine. + @test_throws ErrorException RS._assert_tf32_not_overridden(Dict("NVIDIA_TF32_OVERRIDE" => "0")) + @test RS._assert_tf32_not_overridden(Dict("NVIDIA_TF32_OVERRIDE" => "1")) === nothing + @test RS._assert_tf32_not_overridden(Dict{String, String}()) === nothing +end + +@testset "regulated profile precision report" begin + backend = RS.MockBackend() + pool = RS.MemoryPool(backend, RS.MockClient(), RS.MockDevice(0), "mock", nothing) + base = RS.RuntimeConfig(RS.CPU_BACKEND, 0, 0.9, true, true) + rt(numerics) = RS.RuntimeConfig( + (getfield(base, f) for f in fieldnames(RS.RuntimeConfig)[1:13])..., numerics, + base.executable_cache, RS.PROFILE_REGULATED, RS.BATCH_SIZES_LARGEST, Dict{String, Any}() + ) + # Attested modes are a bannered info line naming the precision; auto is a warning. + @test_logs (:info, r"numerics = tf32\n.*confirmed it is active") RS._report_regulated_precision( + rt(RS.NUMERICS_TF32), backend, pool, (; tf32_active = true, pinned_exact = nothing) + ) + @test_logs (:info, r"numerics = f32\n") RS._report_regulated_precision(rt(RS.NUMERICS_F32), backend, pool, nothing) + @test_logs (:warn, r"WITHOUT PINNED NUMERICS: runtime.numerics = auto") RS._report_regulated_precision( + rt(RS.NUMERICS_AUTO), backend, pool, nothing + ) +end diff --git a/packages/ReactantServerCore/src/ReactantServerCore.jl b/packages/ReactantServerCore/src/ReactantServerCore.jl index ffc706a..8a77505 100644 --- a/packages/ReactantServerCore/src/ReactantServerCore.jl +++ b/packages/ReactantServerCore/src/ReactantServerCore.jl @@ -88,6 +88,8 @@ export ResidencyMode, SELF_MANAGED, EXTERNALLY_MANAGED export ModelControlMode, STATIC, DYNAMIC, EXPLICIT export SchedulingDiscipline, FAIR, FIFO, EDF export NumericsMode, NUMERICS_F32, NUMERICS_AUTO, NUMERICS_TF32 +export BatchSizeMode, BATCH_SIZES_ALL, BATCH_SIZES_LARGEST +export RuntimeProfile, PROFILE_DEFAULT, PROFILE_REGULATED export RuntimeConfig, ModelSchedConfig, SchedulerConfig, EndpointsConfig, GrpcConfig, ServerConfig export build_config, validate_config, apply_env_overrides!, log_effective_config diff --git a/packages/ReactantServerCore/src/config.jl b/packages/ReactantServerCore/src/config.jl index 635fd9a..d2ab8cd 100644 --- a/packages/ReactantServerCore/src/config.jl +++ b/packages/ReactantServerCore/src/config.jl @@ -90,6 +90,38 @@ silently degrading per worker in a mixed fleet. """ @enum NumericsMode NUMERICS_F32 NUMERICS_AUTO NUMERICS_TF32 +""" + BatchSizeMode + +Which of a bundle's compiled batch sizes the worker loads (the `runtime.batch_sizes` knob). +`BATCH_SIZES_ALL` (the default) loads every `model.b{N}.mlir` the manifest declares, so the +scheduler can pick the largest size a dispatch can fill. `BATCH_SIZES_LARGEST` loads only the +largest declared size (per input-shape variant) and never reads the others, so every dispatch +runs the same compiled program, padded with zero rows when fewer requests are queued. That makes +a row's result independent of how many other requests happened to be coalesced with it, at the +cost of paying the full batch's compute on every dispatch. Unbatched bundles (`model.mlir`) are +unaffected. +""" +@enum BatchSizeMode BATCH_SIZES_ALL BATCH_SIZES_LARGEST + +""" + RuntimeProfile + +A named set of runtime defaults (the `runtime.profile` knob). `PROFILE_DEFAULT` changes nothing. +`PROFILE_REGULATED` is for validated deployments that need reproducible results: it defaults +`batch_sizes` to `largest` and adds `xla_gpu_exclude_nondeterministic_ops: true` to `xla_flags`. +Anything set explicitly in the config wins over the profile, so each piece can be turned back off. +The profile does not change `numerics`; instead the worker logs a prominent warning at startup +naming the f32 matmul/convolution precision actually in effect unless `numerics` is `f32`. +""" +@enum RuntimeProfile PROFILE_DEFAULT PROFILE_REGULATED + +# The xla_flags a profile contributes underneath the user's own. +_profile_xla_flags(p::RuntimeProfile) = + p == PROFILE_REGULATED ? Dict{String, Any}("xla_gpu_exclude_nondeterministic_ops" => true) : + Dict{String, Any}() +_profile_batch_sizes(p::RuntimeProfile) = p == PROFILE_REGULATED ? BATCH_SIZES_LARGEST : BATCH_SIZES_ALL + """ RuntimeConfig @@ -121,7 +153,12 @@ are applied to Reactant's compile cache at worker startup, so a container can dr bundle's `.cache/` directory and loads it on later starts instead of recompiling, keyed by the Reactant_jll build, the device, and the MLIR content (a changed `model*.mlir` invalidates its programs; a weights-only update does not); it needs a Reactant that exposes executable -serialization and is otherwise a no-op. +serialization and is otherwise a no-op. `profile` (default `default`) applies a named set of +defaults; see [`RuntimeProfile`](@ref). `batch_sizes` (default `all`, or the profile's) selects +which compiled batch sizes are loaded; see [`BatchSizeMode`](@ref). `xla_flags` (default empty, +plus the profile's) is a mapping of XLA `DebugOptions` field names to values, passed to every +compile, for example `xla_gpu_exclude_nondeterministic_ops: true`; names and value types are +checked against the linked XLA at startup, and the flags are part of the executable cache key. """ struct RuntimeConfig backend::BackendKind @@ -139,8 +176,29 @@ struct RuntimeConfig autotune_cache_dir::String # persistent autotune cache directory; "" = inherit Reactant's LocalPreferences numerics::NumericsMode # f32 matmul/conv precision policy (see NumericsMode) executable_cache::Bool # per-bundle serialized-executable cache + profile::RuntimeProfile # named defaults for batch_sizes and xla_flags + batch_sizes::BatchSizeMode # which compiled batch sizes to load (resolved; profile applied) + xla_flags::Dict{String, Any} # XLA DebugOptions overrides (resolved; profile applied) end +# Fifteen-argument form: the previous full positional layout (through `executable_cache`), with the +# default profile, every batch size, and no XLA flags. +RuntimeConfig( + backend::BackendKind, device_ordinal::Integer, mem_fraction::Real, + preallocate::Bool, allow_cpu_fallback::Bool, residency_mode::ResidencyMode, + shared_host_weights::Bool, shared_host_weights_mode::Integer, + weight_cache_fraction::Real, weight_cache_wiggle_fraction::Real, autotune::Bool, + autotune_cache::Union{Bool, Nothing}, autotune_cache_dir::AbstractString, numerics::NumericsMode, + executable_cache::Bool +) = + RuntimeConfig( + backend, Int(device_ordinal), Float64(mem_fraction), preallocate, allow_cpu_fallback, + residency_mode, shared_host_weights, UInt16(shared_host_weights_mode), + Float64(weight_cache_fraction), Float64(weight_cache_wiggle_fraction), autotune, + autotune_cache, String(autotune_cache_dir), numerics, executable_cache, + PROFILE_DEFAULT, BATCH_SIZES_ALL, Dict{String, Any}() +) + # Five-argument form: device/backend only; residency self-managed, private host weights, and the # on-demand cache off (fraction 0). Used by tests and programmatic construction; the YAML path # (`build_config`) supplies the 1.0 fraction default. @@ -372,6 +430,8 @@ const ENV_PATHS = Tuple{String, Vector{String}, DataType}[ ("RUNTIME_AUTOTUNE_CACHE_DIR", ["runtime", "autotune_cache_dir"], String), ("RUNTIME_NUMERICS", ["runtime", "numerics"], String), ("RUNTIME_EXECUTABLE_CACHE", ["runtime", "executable_cache"], Bool), + ("RUNTIME_PROFILE", ["runtime", "profile"], String), + ("RUNTIME_BATCH_SIZES", ["runtime", "batch_sizes"], String), ("RUNTIME_SHARED_HOST_WEIGHTS", ["runtime", "shared_host_weights"], Bool), ("RUNTIME_SHARED_HOST_WEIGHTS_MODE", ["runtime", "shared_host_weights_mode"], String), ("SCHEDULER_DISCIPLINE", ["scheduler", "discipline"], String), @@ -528,6 +588,39 @@ function _parse_numerics(s) throw(ConfigError("runtime.numerics must be 'f32', 'auto', or 'tf32', got '$s'")) end +function _parse_profile(s) + ls = lowercase(strip(s)) + ls == "default" && return PROFILE_DEFAULT + ls == "regulated" && return PROFILE_REGULATED + throw(ConfigError("runtime.profile must be 'default' or 'regulated', got '$s'")) +end + +function _parse_batch_sizes(s) + ls = lowercase(strip(s)) + ls == "all" && return BATCH_SIZES_ALL + ls == "largest" && return BATCH_SIZES_LARGEST + throw(ConfigError("runtime.batch_sizes must be 'all' or 'largest', got '$s'")) +end + +# runtime.xla_flags: a mapping of XLA DebugOptions field names to scalar values, layered over the +# profile's flags (an explicit entry wins, so `false` turns a profile flag back off). Only the shape +# is checked here; field names and value types are checked against the linked XLA at startup. +function _parse_xla_flags(rt, profile::RuntimeProfile) + flags = _profile_xla_flags(profile) + raw = get(rt, "xla_flags", nothing) + raw === nothing && return flags + raw isa AbstractDict || throw(ConfigError("config 'runtime.xla_flags' must be a mapping")) + for (k, v) in raw + name = String(k) + startswith(name, "xla_") || + throw(ConfigError("runtime.xla_flags.$name: XLA option names start with 'xla_' (no leading '--')")) + (v isa Union{Bool, Real, AbstractString}) || + throw(ConfigError("runtime.xla_flags.$name must be a boolean, number, or string")) + flags[name] = v isa AbstractString ? String(v) : v + end + return flags +end + # Per-model scheduler overrides under scheduler.models. Each entry may set `weight` (relative # compute share, default 1.0), `residency` (initial residency floor), and `max_batch_size` # (coalescing cap, default uncapped). `pin_to_gpu: true` is accepted as a back-compat alias for @@ -569,6 +662,7 @@ function build_config(raw::AbstractDict) rt = _subdict(raw, "runtime") haskey(rt, "residency_mode") && throw(ConfigError(_RESIDENCY_MODE_REMOVED_MSG)) + profile = _parse_profile(_opt(rt, "profile", String, "default")) runtime = RuntimeConfig( _parse_backend(_opt(rt, "backend", String, "cpu")), _opt(rt, "device_ordinal", Int, 0), @@ -585,6 +679,11 @@ function build_config(raw::AbstractDict) _opt(rt, "autotune_cache_dir", String, ""), _parse_numerics(_opt(rt, "numerics", String, "auto")), _opt(rt, "executable_cache", Bool, true), + profile, + haskey(rt, "batch_sizes") ? + _parse_batch_sizes(_coerce(String, rt["batch_sizes"], "runtime.batch_sizes")) : + _profile_batch_sizes(profile), + _parse_xla_flags(rt, profile), ) sc = _subdict(raw, "scheduler") @@ -675,7 +774,7 @@ end # `apply_env_overrides!` is applied on top by `node_server_config`. function log_effective_config(cfg::ServerConfig, applied) - @info "Effective configuration" model_dirs = cfg.model_dirs models_include = cfg.models_include model_control_mode = cfg.model_control_mode model_poll_seconds = cfg.model_poll_seconds cache_dir = cfg.cache_dir backend = cfg.runtime.backend device_ordinal = cfg.runtime.device_ordinal mem_fraction = cfg.runtime.mem_fraction preallocate = cfg.runtime.preallocate allow_cpu_fallback = cfg.runtime.allow_cpu_fallback weight_cache_fraction = cfg.runtime.weight_cache_fraction weight_cache_wiggle_fraction = cfg.runtime.weight_cache_wiggle_fraction autotune = cfg.runtime.autotune autotune_cache = cfg.runtime.autotune_cache autotune_cache_dir = cfg.runtime.autotune_cache_dir numerics = cfg.runtime.numerics executable_cache = cfg.runtime.executable_cache residency_mode = cfg.runtime.residency_mode shared_host_weights = cfg.runtime.shared_host_weights shared_host_weights_mode = string(cfg.runtime.shared_host_weights_mode; base = 8) host = cfg.endpoints.host port = cfg.endpoints.port metrics_port = cfg.endpoints.metrics_port max_concurrent_requests = cfg.endpoints.max_concurrent_requests discipline = cfg.scheduler.discipline ema_halflife_seconds = cfg.scheduler.ema_halflife_seconds recency_penalty_cap = cfg.scheduler.recency_penalty_cap coalescing_discount = cfg.scheduler.coalescing_discount cost_ema_alpha = cfg.scheduler.cost_ema_alpha max_queue_depth = cfg.scheduler.max_queue_depth compaction_interval = cfg.scheduler.compaction_interval scheduler_models = collect(keys(cfg.scheduler.models)) + @info "Effective configuration" model_dirs = cfg.model_dirs models_include = cfg.models_include model_control_mode = cfg.model_control_mode model_poll_seconds = cfg.model_poll_seconds cache_dir = cfg.cache_dir backend = cfg.runtime.backend device_ordinal = cfg.runtime.device_ordinal mem_fraction = cfg.runtime.mem_fraction preallocate = cfg.runtime.preallocate allow_cpu_fallback = cfg.runtime.allow_cpu_fallback weight_cache_fraction = cfg.runtime.weight_cache_fraction weight_cache_wiggle_fraction = cfg.runtime.weight_cache_wiggle_fraction autotune = cfg.runtime.autotune autotune_cache = cfg.runtime.autotune_cache autotune_cache_dir = cfg.runtime.autotune_cache_dir numerics = cfg.runtime.numerics executable_cache = cfg.runtime.executable_cache profile = cfg.runtime.profile batch_sizes = cfg.runtime.batch_sizes xla_flags = sort!(["$k=$v" for (k, v) in cfg.runtime.xla_flags]) residency_mode = cfg.runtime.residency_mode shared_host_weights = cfg.runtime.shared_host_weights shared_host_weights_mode = string(cfg.runtime.shared_host_weights_mode; base = 8) host = cfg.endpoints.host port = cfg.endpoints.port metrics_port = cfg.endpoints.metrics_port max_concurrent_requests = cfg.endpoints.max_concurrent_requests discipline = cfg.scheduler.discipline ema_halflife_seconds = cfg.scheduler.ema_halflife_seconds recency_penalty_cap = cfg.scheduler.recency_penalty_cap coalescing_discount = cfg.scheduler.coalescing_discount cost_ema_alpha = cfg.scheduler.cost_ema_alpha max_queue_depth = cfg.scheduler.max_queue_depth compaction_interval = cfg.scheduler.compaction_interval scheduler_models = collect(keys(cfg.scheduler.models)) isempty(applied) || @info "Configuration overridden by environment" overrides = ["$k=$v" for (k, v) in applied] return nothing end diff --git a/packages/ReactantServerCore/test/test_config.jl b/packages/ReactantServerCore/test/test_config.jl index 12ffadc..4b88938 100644 --- a/packages/ReactantServerCore/test/test_config.jl +++ b/packages/ReactantServerCore/test/test_config.jl @@ -578,3 +578,68 @@ end @test c14.shared_host_weights_mode == 0o660 && c14.numerics == ReactantServer.NUMERICS_F32 end end + +@testset "runtime.profile, batch_sizes, xla_flags" begin + mktempdir() do dir + modeldir = joinpath(dir, "models"); mkpath(modeldir) + rt(body) = load_single_worker(dir, body; model_repo = modeldir)[1].runtime + + # Default profile: every batch size, no XLA flags. + d = rt("runtime:\n backend: cuda") + @test d.profile == ReactantServer.PROFILE_DEFAULT + @test d.batch_sizes == ReactantServer.BATCH_SIZES_ALL + @test isempty(d.xla_flags) + @test rt("runtime:\n backend: cuda\n batch_sizes: largest").batch_sizes == + ReactantServer.BATCH_SIZES_LARGEST + + # The regulated profile supplies largest + the nondeterministic-ops exclusion, and leaves + # numerics alone. + r = rt("runtime:\n backend: cuda\n profile: regulated") + @test r.profile == ReactantServer.PROFILE_REGULATED + @test r.batch_sizes == ReactantServer.BATCH_SIZES_LARGEST + @test r.xla_flags == Dict{String, Any}("xla_gpu_exclude_nondeterministic_ops" => true) + @test r.numerics == ReactantServer.NUMERICS_AUTO + + # Explicit settings win over the profile, piece by piece, and add to its flags. + o = rt( + """ + runtime: + backend: cuda + profile: regulated + batch_sizes: all + xla_flags: + xla_gpu_exclude_nondeterministic_ops: false + xla_gpu_deterministic_ops: true + xla_gpu_autotune_level: 4 + """ + ) + @test o.batch_sizes == ReactantServer.BATCH_SIZES_ALL + @test o.xla_flags == Dict{String, Any}( + "xla_gpu_exclude_nondeterministic_ops" => false, + "xla_gpu_deterministic_ops" => true, "xla_gpu_autotune_level" => 4, + ) + + # Shape errors are ConfigErrors at parse time (names/types are checked by the worker). + @test_throws ReactantServer.ConfigError rt("runtime:\n profile: strict") + @test_throws ReactantServer.ConfigError rt("runtime:\n batch_sizes: smallest") + @test_throws ReactantServer.ConfigError rt("runtime:\n xla_flags: [a, b]") + @test_throws ReactantServer.ConfigError rt("runtime:\n xla_flags:\n --xla_gpu_deterministic_ops: true") + @test_throws ReactantServer.ConfigError rt("runtime:\n xla_flags:\n xla_gpu_foo: [1, 2]") + + # Env overrides for the two scalar knobs. + withenv( + "INFERENCE_SERVER_RUNTIME_PROFILE" => "regulated", + "INFERENCE_SERVER_RUNTIME_BATCH_SIZES" => "all", + ) do + e = rt("runtime:\n backend: cuda") + @test e.profile == ReactantServer.PROFILE_REGULATED + @test e.batch_sizes == ReactantServer.BATCH_SIZES_ALL + @test haskey(e.xla_flags, "xla_gpu_exclude_nondeterministic_ops") + end + + # Positional constructors keep the defaults. + c5 = ReactantServer.RuntimeConfig(ReactantServer.CPU_BACKEND, 0, 0.9, true, true) + @test c5.profile == ReactantServer.PROFILE_DEFAULT && c5.batch_sizes == ReactantServer.BATCH_SIZES_ALL + @test isempty(c5.xla_flags) + end +end