Add regulated profile for deterministic behavior

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`.
 

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` |

packages/ReactantServer/src/bundle.jl

@@ -41,56 +41,72 @@ 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 `<prefix>.b{N}.mlir` (keyed by N) or a single `<prefix>.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
 # `input_shapes`) yields one default variant `Int[]` from `model.mlir`/`model.b{N}.mlir`. A
 # 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,20 +158,21 @@ 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.
 
 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

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).

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
     )