428 lines
18 KiB
Julia
Executable File
428 lines
18 KiB
Julia
Executable File
#!/usr/bin/env julia
|
||
#
|
||
# bench_model.jl: microbenchmark the classifier in isolation, with no server,
|
||
# no queue, and no disk in the way.
|
||
#
|
||
# bin/bench.jl measures the *pipeline*: it reports stage 1 as one number, the
|
||
# wall time of `handle_classify_job`, which is feature reads + inference + a
|
||
# rename + a log line, under whatever thread contention the other three pools are
|
||
# creating. That number is the right one for capacity planning and the wrong one
|
||
# for answering "is the model slow?". This script answers that question by taking
|
||
# the model apart:
|
||
#
|
||
# read_features open, read 16 bytes, seek, read 16 bytes, scale
|
||
# Lux.apply the network itself, on a feature vector already in memory
|
||
# classify both together: what stage 1 actually calls per file
|
||
#
|
||
# Three properties are worth checking beyond the raw per-file cost:
|
||
#
|
||
# * Feature reads should be flat in file size. read_features seeks to the tail
|
||
# rather than slurping, so a 1 GiB file should cost the same as a 1 KiB one.
|
||
# (This is the same claim bin/bench.jl makes about memory, on the CPU axis.)
|
||
# * Batching should be much cheaper per file. A 32x1 matmul wastes most of a
|
||
# BLAS call; if batch-64 inference is many times cheaper per file, that is
|
||
# the headroom a batching stage-1 would buy, worth knowing before building
|
||
# one, since today the pipeline classifies strictly one file at a time.
|
||
# * Inference should scale across threads. `Classifier` is shared read-only by
|
||
# the whole stage-1 pool on the claim that Lux inference is pure. If per-file
|
||
# cost degrades as tasks are added, that claim holds but BLAS threading is
|
||
# fighting the worker pool, and stage-1 workers are contending, not scaling.
|
||
#
|
||
# Usage:
|
||
# julia --project=. -t auto bin/bench_model.jl [options]
|
||
#
|
||
# --model PATH classifier artifact (default: $FS_MODEL_PATH or model/classifier.jld2)
|
||
# --reps N inference calls per timed trial (default: 20000)
|
||
# --trials N timed trials; the minimum is reported (default: 5)
|
||
# --batches LIST batch sizes to sweep, comma-separated (default: 1,8,64,512)
|
||
# --sizes LIST file sizes for the read_features sweep (default: 1k,64k,4m,256m)
|
||
# --no-threads skip the thread-scaling sweep
|
||
# --json PATH also write the results as JSON
|
||
#
|
||
# Reported times are the *minimum* over trials: for a microbenchmark the floor is
|
||
# the signal and everything above it is scheduler and GC noise.
|
||
|
||
using JSON3
|
||
using Random
|
||
using Statistics
|
||
using Printf
|
||
|
||
# The script is run directly, not as part of the package, so pull in exactly the
|
||
# pieces the classifier needs. `Lux`/`JLD2` first, because model.jl and classify.jl both
|
||
# assume the including scope already has them (see the note at the top of model.jl).
|
||
using Lux
|
||
using JLD2
|
||
using LinearAlgebra
|
||
|
||
const SRC = joinpath(dirname(@__DIR__), "src")
|
||
include(joinpath(SRC, "model.jl"))
|
||
include(joinpath(SRC, "classify.jl"))
|
||
|
||
# ---------------------------------------------------------------- option parsing
|
||
|
||
const DEFAULTS = Dict{String,Any}(
|
||
"model" => get(ENV, "FS_MODEL_PATH", "model/classifier.jld2"),
|
||
"reps" => 20_000,
|
||
"trials" => 5,
|
||
"batches" => "1,8,64,512",
|
||
"sizes" => "1k,64k,4m,256m",
|
||
"no-threads" => false,
|
||
"json" => nothing,
|
||
)
|
||
|
||
const FLAGS = ("no-threads",)
|
||
|
||
function parse_size(s::AbstractString)::Int
|
||
m = match(r"^(\d+(?:\.\d+)?)\s*([kKmMgG]?)[bB]?$", strip(s))
|
||
m === nothing && error("bad size: $s (expected e.g. 512, 64k, 8m, 1g)")
|
||
mult = Dict('k' => 1024, 'm' => 1024^2, 'g' => 1024^3)
|
||
scale = isempty(m[2]) ? 1 : mult[lowercase(m[2])[1]]
|
||
return round(Int, parse(Float64, m[1]) * scale)
|
||
end
|
||
|
||
function parse_args(argv)
|
||
opts = copy(DEFAULTS)
|
||
i = 1
|
||
while i <= length(argv)
|
||
a = argv[i]
|
||
startswith(a, "--") || error("unexpected argument: $a")
|
||
key = a[3:end]
|
||
haskey(opts, key) || error("unknown option: $a")
|
||
if key in FLAGS
|
||
opts[key] = true; i += 1; continue
|
||
end
|
||
i + 1 <= length(argv) || error("option --$key needs a value")
|
||
opts[key] = key in ("reps", "trials") ? parse(Int, argv[i+1]) : argv[i+1]
|
||
i += 2
|
||
end
|
||
return opts
|
||
end
|
||
|
||
# ------------------------------------------------------------------- measurement
|
||
|
||
# Every timed loop stores its result here. Without a visible side effect the
|
||
# compiler is free to hoist a pure call out of the loop and we would be timing an
|
||
# empty `for`.
|
||
const SINK = Ref{Any}(nothing)
|
||
|
||
"""
|
||
measure(f, reps; trials) -> (ns_per_op, bytes_per_op)
|
||
|
||
Time `f` over `reps` calls, `trials` times, and report the fastest trial.
|
||
|
||
The first call is thrown away: it pays Julia's JIT compilation, which on a
|
||
function this small is orders of magnitude more than the thing being measured.
|
||
"""
|
||
function measure(f, reps::Int; trials::Int = 5)
|
||
SINK[] = f() # warm up (compile), and keep the result
|
||
best = Inf
|
||
for _ in 1:trials
|
||
GC.gc()
|
||
t0 = time_ns()
|
||
for _ in 1:reps
|
||
SINK[] = f()
|
||
end
|
||
best = min(best, (time_ns() - t0) / reps)
|
||
end
|
||
bytes = @allocated(SINK[] = f()) # one call, after warmup
|
||
return (Float64(best), Float64(bytes))
|
||
end
|
||
|
||
# ------------------------------------------------------------------- formatting
|
||
|
||
function human_time(ns::Real)
|
||
ns < 1_000 && return @sprintf("%.0f ns", ns)
|
||
ns < 1_000_000 && return @sprintf("%.2f µs", ns / 1e3)
|
||
ns < 1e9 && return @sprintf("%.2f ms", ns / 1e6)
|
||
return @sprintf("%.2f s", ns / 1e9)
|
||
end
|
||
|
||
human_bytes(b::Real) = b < 1024 ? @sprintf("%.0f B", b) :
|
||
b < 1024^2 ? @sprintf("%.1f KiB", b / 1024) :
|
||
@sprintf("%.1f MiB", b / 1024^2)
|
||
|
||
rate(ns::Real) = 1e9 / max(ns, 1e-9) # calls per second
|
||
|
||
function human_rate(r::Real)
|
||
r >= 1e6 && return @sprintf("%.2fM/s", r / 1e6)
|
||
r >= 1e3 && return @sprintf("%.1fk/s", r / 1e3)
|
||
return @sprintf("%.0f/s", r)
|
||
end
|
||
|
||
fmt2(x::Real) = @sprintf("%.2f", x)
|
||
|
||
# ------------------------------------------------------------- scaling control
|
||
|
||
"""
|
||
control_kernel(x) -> Float64
|
||
|
||
Pure arithmetic, no allocation, no library call, and deliberately dependent
|
||
(each step needs the last) so the compiler can't vectorize it away, and sized to
|
||
land in the same microsecond neighbourhood as one `Lux.apply`.
|
||
"""
|
||
function control_kernel(x::Float64)
|
||
a = x
|
||
@inbounds for i in 1:600
|
||
a = sqrt(a + i)
|
||
end
|
||
return a
|
||
end
|
||
|
||
"""
|
||
control_scaling(trials) -> Vector
|
||
|
||
Run `control_kernel` over the same task counts as the model sweep. This is the
|
||
machine's own ceiling for perfectly parallel work: if the control scales and the
|
||
model doesn't, the shortfall is the model's, and no amount of `FS_WORKERS` will
|
||
recover it.
|
||
"""
|
||
function control_scaling(trials::Int)
|
||
rows = []
|
||
base = 0.0
|
||
per_task = 200_000
|
||
for k in unique([1; 2; 4; 8; Threads.nthreads()])
|
||
k > Threads.nthreads() && continue
|
||
best = Inf
|
||
for _ in 1:trials
|
||
GC.gc()
|
||
t0 = time_ns()
|
||
@sync for _ in 1:k
|
||
Threads.@spawn begin
|
||
local acc = 0.0
|
||
for i in 1:per_task
|
||
acc += control_kernel(i % 97 + 1.0)
|
||
end
|
||
SINK[] = acc
|
||
end
|
||
end
|
||
best = min(best, Float64(time_ns() - t0))
|
||
end
|
||
r = k * per_task / (best / 1e9)
|
||
k == 1 && (base = r)
|
||
push!(rows, (; tasks = k, ops_per_sec = r, speedup = r / base))
|
||
end
|
||
return rows
|
||
end
|
||
|
||
# ------------------------------------------------------------------------- corpus
|
||
|
||
"""
|
||
Write a file of exactly `size` random bytes, in bounded chunks.
|
||
|
||
Content is random rather than zeros so the classifier sees a realistic input,
|
||
and so the filesystem can't cheat with a sparse file, which would make the tail
|
||
`seek` unrepresentatively fast.
|
||
"""
|
||
function write_file(path::AbstractString, size::Int, rng)
|
||
chunk = 1024 * 1024
|
||
open(path, "w") do io
|
||
remaining = size
|
||
while remaining > 0
|
||
n = min(chunk, remaining)
|
||
write(io, rand(rng, UInt8, n))
|
||
remaining -= n
|
||
end
|
||
end
|
||
return path
|
||
end
|
||
|
||
# ------------------------------------------------------------------------- main
|
||
|
||
function main(argv)
|
||
opts = parse_args(argv)
|
||
modelpath = String(opts["model"])
|
||
isfile(modelpath) || (println(stderr, "model artifact not found: $modelpath"); return 1)
|
||
|
||
reps, trials = opts["reps"], opts["trials"]
|
||
batches = [parse(Int, s) for s in split(String(opts["batches"]), ",")]
|
||
sizes = [parse_size(s) for s in split(String(opts["sizes"]), ",")]
|
||
|
||
clf = load_classifier(modelpath)
|
||
println("model $modelpath")
|
||
println("architecture $(FEATURE_DIM) → 64 → 16 → 2 (Dense/relu, raw logits)")
|
||
println("julia threads $(Threads.nthreads()) BLAS threads $(BLAS.get_num_threads())")
|
||
println("timing min of $trials trials × $reps reps")
|
||
println("=" ^ 72)
|
||
|
||
results = Dict{String,Any}()
|
||
|
||
# --- 1. inference alone, one file at a time: the number the pipeline pays.
|
||
x1 = rand(Float32, FEATURE_DIM, 1)
|
||
infer_ns, infer_bytes = measure(reps; trials) do
|
||
Lux.apply(clf.model, x1, clf.ps, clf.st)
|
||
end
|
||
println()
|
||
println("INFERENCE (Lux.apply, batch 1; features already in memory)")
|
||
println(" per call $(human_time(infer_ns)) $(human_rate(rate(infer_ns)))")
|
||
println(" allocations $(human_bytes(infer_bytes)) per call")
|
||
results["inference_batch1"] = (; ns = infer_ns, bytes = infer_bytes, per_sec = rate(infer_ns))
|
||
|
||
# --- 2. batching: how much of that is per-call overhead rather than math?
|
||
println()
|
||
println("INFERENCE BATCHED (same net, N files per apply)")
|
||
println(" batch per batch per file files/s speedup")
|
||
batch_rows = []
|
||
for b in batches
|
||
xb = rand(Float32, FEATURE_DIM, b)
|
||
ns, _ = measure(max(1, reps ÷ b); trials) do
|
||
Lux.apply(clf.model, xb, clf.ps, clf.st)
|
||
end
|
||
per_file = ns / b
|
||
@printf(" %8d %13s %12s %13s %7.1fx\n",
|
||
b, human_time(ns), human_time(per_file),
|
||
human_rate(rate(per_file)), infer_ns / per_file)
|
||
push!(batch_rows, (; batch = b, ns_per_batch = ns, ns_per_file = per_file,
|
||
files_per_sec = rate(per_file), speedup = infer_ns / per_file))
|
||
end
|
||
results["batched"] = batch_rows
|
||
println(" (a large speedup is headroom a batching stage 1 could claim; the pipeline")
|
||
println(" classifies one file per job today, so it pays the batch-1 row above)")
|
||
|
||
# --- 3. feature reads: should be flat in file size (seek, not slurp).
|
||
println()
|
||
println("FEATURE READS (read_features: 16 head + 16 tail bytes, scaled)")
|
||
println(" file size per call calls/s allocations")
|
||
read_rows = []
|
||
dir = mktempdir(; prefix = "fsmodel-")
|
||
try
|
||
rng = MersenneTwister(1234)
|
||
for sz in sizes
|
||
path = write_file(joinpath(dir, "f-$sz.bin"), sz, rng)
|
||
# Fewer reps for the big files: this touches the page cache, and the
|
||
# point is the shape of the curve, not another digit of precision.
|
||
r = max(200, reps ÷ 20)
|
||
ns, bytes = measure(() -> read_features(path), r; trials)
|
||
@printf(" %13s %14s %14s %14s\n",
|
||
human_bytes(sz), human_time(ns), human_rate(rate(ns)), human_bytes(bytes))
|
||
push!(read_rows, (; size_bytes = sz, ns, bytes, per_sec = rate(ns)))
|
||
end
|
||
finally
|
||
rm(dir; recursive = true, force = true)
|
||
end
|
||
results["read_features"] = read_rows
|
||
flat = length(read_rows) > 1 ?
|
||
maximum(r.ns for r in read_rows) / minimum(r.ns for r in read_rows) : 1.0
|
||
@printf(" spread across a %.0fx size range: %.1fx, %s\n",
|
||
maximum(sizes) / minimum(sizes), flat,
|
||
flat < 3 ? "flat, as designed (it seeks to the tail)" :
|
||
"NOT flat: something is reading more than 32 bytes")
|
||
|
||
# --- 4. classify(): what stage 1 calls, I/O and inference together.
|
||
println()
|
||
println("CLASSIFY (read_features + Lux.apply; one whole stage-1 file)")
|
||
dir2 = mktempdir(; prefix = "fsmodel-")
|
||
classify_ns = 0.0
|
||
try
|
||
path = write_file(joinpath(dir2, "sample.bin"), 64 * 1024, MersenneTwister(7))
|
||
classify_ns, classify_bytes = measure(() -> classify(clf, path), max(200, reps ÷ 20); trials)
|
||
println(" per file $(human_time(classify_ns)) $(human_rate(rate(classify_ns)))")
|
||
println(" allocations $(human_bytes(classify_bytes)) per file")
|
||
@printf(" split %.0f%% feature read, %.0f%% inference\n",
|
||
100 * (classify_ns - infer_ns) / classify_ns, 100 * infer_ns / classify_ns)
|
||
results["classify"] = (; ns = classify_ns, bytes = classify_bytes, per_sec = rate(classify_ns))
|
||
finally
|
||
rm(dir2; recursive = true, force = true)
|
||
end
|
||
|
||
# --- 5. thread scaling: does the shared read-only Classifier actually scale?
|
||
if !opts["no-threads"] && Threads.nthreads() > 1
|
||
println()
|
||
println("THREAD SCALING (concurrent Lux.apply on the one shared Classifier)")
|
||
println(" tasks files/s per file speedup efficiency GC")
|
||
thread_rows = []
|
||
base = 0.0
|
||
for k in unique([1; 2; 4; 8; Threads.nthreads()])
|
||
k > Threads.nthreads() && continue
|
||
# Work per task is held *constant* as tasks are added, so total work
|
||
# scales with `k`. Splitting a fixed total instead would shrink each
|
||
# task as the pool grows until `@spawn`/`@sync` overhead dominated,
|
||
# and the resulting curve would show a collapse that is the
|
||
# measurement's fault rather than the model's.
|
||
per_task = max(reps, 20_000)
|
||
# Each task gets its own input so we measure the model, not cache
|
||
# line ping-pong on a shared buffer.
|
||
xs = [rand(Float32, FEATURE_DIM, 1) for _ in 1:k]
|
||
best, best_gc = Inf, 0.0
|
||
for _ in 1:trials
|
||
GC.gc()
|
||
# Julia's GC stops the world, so it is the one cost that cannot
|
||
# be parallelized away: measuring its share here is what turns a
|
||
# bad efficiency number into a diagnosis (see the note below).
|
||
gc0 = Base.gc_num().total_time
|
||
t0 = time_ns()
|
||
@sync for t in 1:k
|
||
Threads.@spawn begin
|
||
local acc = 0.0f0
|
||
for _ in 1:per_task
|
||
y, _ = Lux.apply(clf.model, xs[t], clf.ps, clf.st)
|
||
acc += y[1] # consume the result
|
||
end
|
||
SINK[] = acc
|
||
end
|
||
end
|
||
elapsed = Float64(time_ns() - t0)
|
||
if elapsed < best
|
||
best = elapsed
|
||
best_gc = Float64(Base.gc_num().total_time - gc0)
|
||
end
|
||
end
|
||
files = k * per_task
|
||
fps = files / (best / 1e9)
|
||
k == 1 && (base = fps)
|
||
@printf(" %8d %13s %11s %7.2fx %10.0f%% %5.0f%%\n",
|
||
k, human_rate(fps), human_time(best / files), fps / base,
|
||
100 * fps / base / k, 100 * best_gc / best)
|
||
push!(thread_rows, (; tasks = k, files_per_sec = fps,
|
||
ns_per_file = best / files, speedup = fps / base,
|
||
gc_fraction = best_gc / best))
|
||
end
|
||
results["thread_scaling"] = thread_rows
|
||
|
||
# A poor scaling curve has two possible authors, the model or the box,
|
||
# and the table alone can't tell them apart. So run the same sweep on a
|
||
# kernel that is pure arithmetic with no allocation and no library
|
||
# underneath: whatever *it* achieves is this machine's ceiling for
|
||
# embarrassingly parallel work, and the gap between the two curves is
|
||
# the part that belongs to Lux.apply.
|
||
ctrl = control_scaling(trials)
|
||
results["control_scaling"] = ctrl
|
||
top = last(ctrl)
|
||
println(" control pure-compute kernel, same sweep: " *
|
||
"$(fmt2(top.speedup))x at $(top.tasks) tasks " *
|
||
"($(round(Int, 100 * top.speedup / top.tasks))% efficiency)")
|
||
model_top = last(thread_rows)
|
||
if model_top.speedup < 0.6 * top.speedup
|
||
println(" → the machine parallelizes; Lux.apply does not. Stage-1")
|
||
println(" workers past ~4 buy little, whatever FS_WORKERS says.")
|
||
else
|
||
println(" → inference tracks the machine's own scaling ceiling.")
|
||
end
|
||
gc_top = maximum(r.gc_fraction for r in thread_rows)
|
||
gc_top > 0.15 && println(" ! GC is $(round(Int, 100 * gc_top))% of the " *
|
||
"worst case: apply allocates per call, and\n" *
|
||
" collection stops every thread.")
|
||
end
|
||
|
||
println()
|
||
println("=" ^ 72)
|
||
println("Stage 1's cost per file in bin/bench.jl is this classify() figure plus a")
|
||
println("rename, a log line, and whatever contention the other three pools create.")
|
||
println("A large gap between the two is pipeline overhead, not the model.")
|
||
|
||
if opts["json"] !== nothing
|
||
results["meta"] = (; model = modelpath, feature_dim = FEATURE_DIM,
|
||
julia_threads = Threads.nthreads(),
|
||
blas_threads = BLAS.get_num_threads(),
|
||
reps, trials)
|
||
open(String(opts["json"]), "w") do io
|
||
JSON3.write(io, results)
|
||
end
|
||
println("\nwrote $(opts["json"])")
|
||
end
|
||
return 0
|
||
end
|
||
|
||
if abspath(PROGRAM_FILE) == @__FILE__
|
||
exit(main(ARGS))
|
||
end
|