Compare commits
19 Commits
842668b2ac
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 4aed713939 | |||
| 630bf5f9a7 | |||
| cc84d70d52 | |||
| 819a4cce7a | |||
| 4a22123001 | |||
| 341b61f806 | |||
| c692d14a2c | |||
| c5d488d9b4 | |||
| 0d8eba05b8 | |||
| e18ac45d70 | |||
| f4e3f5be0b | |||
| 584bad02a7 | |||
| d9f32d9aaf | |||
| 2c8de488a1 | |||
| fac3adbaf6 | |||
| 9fd1bf385b | |||
| e42e8ef8af | |||
| 2a46f5021a | |||
| 1c7d7d6cad |
16
.dockerignore
Normal file
16
.dockerignore
Normal file
@@ -0,0 +1,16 @@
|
|||||||
|
# The pipeline working set (17G here) belongs in a volume, not the image.
|
||||||
|
data/
|
||||||
|
.git/
|
||||||
|
test/
|
||||||
|
# Benchmark / training / calibration harnesses are dev tools, not runtime.
|
||||||
|
bin/bench.jl
|
||||||
|
bin/bench_model.jl
|
||||||
|
bin/bench_stage1.jl
|
||||||
|
bin/bench_stage2.jl
|
||||||
|
bin/cluster_calibrate.jl
|
||||||
|
bin/train.jl
|
||||||
|
bin/send_dir.sh
|
||||||
|
Dockerfile
|
||||||
|
docker-compose.yml
|
||||||
|
.dockerignore
|
||||||
|
*.md
|
||||||
146
Dockerfile
Normal file
146
Dockerfile
Normal file
@@ -0,0 +1,146 @@
|
|||||||
|
# Multi-stage build for the DarkStruct pipeline.
|
||||||
|
#
|
||||||
|
# The runtime image carries only what serving needs: the Julia runtime, a depot
|
||||||
|
# that is already instantiated *and* precompiled (so startup is load-only, no
|
||||||
|
# compilation), exiftool, and optionally github-linguist. Nothing from the build
|
||||||
|
# survives into it: no compilers, no package registry, no gem toolchain.
|
||||||
|
#
|
||||||
|
# docker build -t darkstruct .
|
||||||
|
# docker build --build-arg WITH_LINGUIST=false -t darkstruct . # ~200MB smaller
|
||||||
|
#
|
||||||
|
# github-linguist (stage-4 programming-language detection) is the one heavy
|
||||||
|
# optional dependency: it needs a Ruby toolchain to build rugged. It degrades
|
||||||
|
# gracefully: without it, text files still get natural-language enrichment and
|
||||||
|
# the server logs a warning at startup, so it can be built out.
|
||||||
|
|
||||||
|
ARG JULIA_VERSION=1.12.6
|
||||||
|
ARG DEBIAN_RELEASE=bookworm
|
||||||
|
ARG WITH_LINGUIST=true
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Stage 1: instantiate and precompile the Julia depot.
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
FROM julia:${JULIA_VERSION}-${DEBIAN_RELEASE} AS julia-depot
|
||||||
|
|
||||||
|
ENV JULIA_DEPOT_PATH=/opt/julia-depot \
|
||||||
|
JULIA_PKG_PRECOMPILE_AUTO=0
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
|
||||||
|
# Manifest-first so a source-only edit doesn't re-resolve or rebuild the whole
|
||||||
|
# dependency tree. DarkStruct is its own project's root package, so Pkg insists
|
||||||
|
# on a module file being there; a stub satisfies it while the (slow, rarely
|
||||||
|
# invalidated) dependency layer is built; the real source lands below.
|
||||||
|
COPY Project.toml Manifest.toml ./
|
||||||
|
RUN mkdir -p src \
|
||||||
|
&& printf 'module DarkStruct\nend\n' > src/DarkStruct.jl \
|
||||||
|
&& julia --project=/app -e 'using Pkg; Pkg.instantiate(); Pkg.precompile()'
|
||||||
|
|
||||||
|
# Precompiling DarkStruct itself means the runtime image loads cached code
|
||||||
|
# instead of compiling Lux/Zygote on first request.
|
||||||
|
COPY src/ ./src/
|
||||||
|
RUN julia --project=/app -e 'using Pkg; Pkg.precompile()' \
|
||||||
|
&& julia --project=/app -e 'using DarkStruct'
|
||||||
|
|
||||||
|
# Everything the depot keeps for *resolving* packages is dead weight once the
|
||||||
|
# Manifest is pinned: the registry, git clones, scratch and log dirs.
|
||||||
|
RUN rm -rf /opt/julia-depot/registries \
|
||||||
|
/opt/julia-depot/clones \
|
||||||
|
/opt/julia-depot/scratchspaces \
|
||||||
|
/opt/julia-depot/logs \
|
||||||
|
&& find /opt/julia-depot/packages -type d \
|
||||||
|
\( -name test -o -name docs -o -name .git \) -prune -exec rm -rf {} + \
|
||||||
|
&& find /opt/julia-depot/artifacts -type d -name doc -prune -exec rm -rf {} +
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Stage 2: build github-linguist into a relocatable GEM_HOME (optional).
|
||||||
|
#
|
||||||
|
# Built against Debian's system Ruby, the same package the runtime installs,
|
||||||
|
# so rugged's native extension is ABI-compatible with the interpreter there.
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
FROM debian:${DEBIAN_RELEASE}-slim AS linguist
|
||||||
|
ARG WITH_LINGUIST
|
||||||
|
|
||||||
|
ENV GEM_HOME=/opt/gems
|
||||||
|
|
||||||
|
RUN set -eux; \
|
||||||
|
mkdir -p "$GEM_HOME"; \
|
||||||
|
if [ "$WITH_LINGUIST" = "true" ]; then \
|
||||||
|
apt-get update; \
|
||||||
|
apt-get install -y --no-install-recommends \
|
||||||
|
ruby ruby-dev build-essential cmake pkg-config \
|
||||||
|
libicu-dev zlib1g-dev libssl-dev; \
|
||||||
|
gem install --no-document github-linguist; \
|
||||||
|
# Drop sources, docs and test fixtures from the installed gems; only the
|
||||||
|
# built .so and the .rb load path matter at runtime.
|
||||||
|
rm -rf "$GEM_HOME"/cache "$GEM_HOME"/doc; \
|
||||||
|
find "$GEM_HOME"/gems -type d \( -name test -o -name spec -o -name samples \) \
|
||||||
|
-prune -exec rm -rf {} +; \
|
||||||
|
find "$GEM_HOME"/extensions -type d -name tmp -prune -exec rm -rf {} +; \
|
||||||
|
rm -rf /var/lib/apt/lists/*; \
|
||||||
|
fi
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Stage 3: runtime.
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
FROM debian:${DEBIAN_RELEASE}-slim AS runtime
|
||||||
|
ARG WITH_LINGUIST
|
||||||
|
|
||||||
|
# exiftool is a hard startup prerequisite (stage-2 enrichment fails fast without
|
||||||
|
# it). Ruby + libicu are only needed to *run* the linguist gem built above.
|
||||||
|
RUN set -eux; \
|
||||||
|
apt-get update; \
|
||||||
|
apt-get install -y --no-install-recommends \
|
||||||
|
libimage-exiftool-perl ca-certificates curl; \
|
||||||
|
if [ "$WITH_LINGUIST" = "true" ]; then \
|
||||||
|
apt-get install -y --no-install-recommends ruby libicu72; \
|
||||||
|
fi; \
|
||||||
|
rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
RUN useradd --system --uid 10001 --create-home --home-dir /app app
|
||||||
|
|
||||||
|
# Depot ownership matters: Julia rewrites a precompile cache it considers stale,
|
||||||
|
# and a read-only depot would turn that into a startup crash instead of a
|
||||||
|
# recompile. --chown does it in the copy, so it costs no extra layer.
|
||||||
|
COPY --from=julia-depot /usr/local/julia /usr/local/julia
|
||||||
|
COPY --from=julia-depot --chown=app:app /opt/julia-depot /opt/julia-depot
|
||||||
|
COPY --from=linguist /opt/gems /opt/gems
|
||||||
|
|
||||||
|
ENV PATH=/usr/local/julia/bin:/opt/gems/bin:$PATH \
|
||||||
|
GEM_HOME=/opt/gems \
|
||||||
|
GEM_PATH=/opt/gems \
|
||||||
|
JULIA_DEPOT_PATH=/opt/julia-depot \
|
||||||
|
JULIA_NUM_THREADS=auto \
|
||||||
|
JULIA_PKG_OFFLINE=true
|
||||||
|
|
||||||
|
WORKDIR /app
|
||||||
|
COPY Project.toml Manifest.toml ./
|
||||||
|
COPY src/ ./src/
|
||||||
|
COPY bin/server.jl bin/cluster_sweep.jl ./bin/
|
||||||
|
COPY model/classifier.jld2 ./model/
|
||||||
|
|
||||||
|
# Every stage directory hangs off /data so one volume covers the whole pipeline.
|
||||||
|
# 0.0.0.0 rather than the 127.0.0.1 default: the port is only reachable from
|
||||||
|
# outside if it is published.
|
||||||
|
ENV FS_HOST=0.0.0.0 \
|
||||||
|
FS_PORT=8080 \
|
||||||
|
FS_SPOOL_DIR=/data/spool \
|
||||||
|
FS_BINARY_DIR=/data/binary \
|
||||||
|
FS_DONE_DIR=/data/done \
|
||||||
|
FS_TEXT_DONE_DIR=/data/text_done \
|
||||||
|
FS_FAILED_DIR=/data/failed \
|
||||||
|
FS_CLUSTER_DIR=/data/binary \
|
||||||
|
FS_CLUSTER_CATALOG=/data/catalog.json \
|
||||||
|
FS_NOMINATED_DIR=/data/nominated \
|
||||||
|
FS_MODEL_PATH=/app/model/classifier.jld2
|
||||||
|
|
||||||
|
RUN mkdir -p /data && chown app:app /data
|
||||||
|
USER app
|
||||||
|
VOLUME ["/data"]
|
||||||
|
|
||||||
|
EXPOSE 8080
|
||||||
|
HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --retries=3 \
|
||||||
|
CMD curl -fsS "http://127.0.0.1:${FS_PORT}/health" || exit 1
|
||||||
|
|
||||||
|
ENTRYPOINT ["julia", "--project=/app", "-t", "auto"]
|
||||||
|
CMD ["bin/server.jl"]
|
||||||
@@ -2,7 +2,7 @@
|
|||||||
|
|
||||||
julia_version = "1.12.6"
|
julia_version = "1.12.6"
|
||||||
manifest_format = "2.0"
|
manifest_format = "2.0"
|
||||||
project_hash = "1d7ce552eaac13c97732dccac45c97d40527ec58"
|
project_hash = "d7c84379a20829ec867546b4a887bc3c10df8311"
|
||||||
|
|
||||||
[[deps.ADTypes]]
|
[[deps.ADTypes]]
|
||||||
git-tree-sha1 = "d9aaef7c63466eee4de23b4d9dad03629df54bea"
|
git-tree-sha1 = "d9aaef7c63466eee4de23b4d9dad03629df54bea"
|
||||||
@@ -15,6 +15,12 @@ weakdeps = ["ChainRulesCore", "ConstructionBase", "EnzymeCore"]
|
|||||||
ADTypesConstructionBaseExt = "ConstructionBase"
|
ADTypesConstructionBaseExt = "ConstructionBase"
|
||||||
ADTypesEnzymeCoreExt = "EnzymeCore"
|
ADTypesEnzymeCoreExt = "EnzymeCore"
|
||||||
|
|
||||||
|
[[deps.AMQPClient]]
|
||||||
|
deps = ["Logging", "MbedTLS", "Sockets"]
|
||||||
|
git-tree-sha1 = "508457ed7a2afb432590247dc363fffc51f242fc"
|
||||||
|
uuid = "79c8b4cd-a41a-55fa-907c-fab5288e1383"
|
||||||
|
version = "0.5.1"
|
||||||
|
|
||||||
[[deps.AbstractFFTs]]
|
[[deps.AbstractFFTs]]
|
||||||
deps = ["LinearAlgebra"]
|
deps = ["LinearAlgebra"]
|
||||||
git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef"
|
git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef"
|
||||||
@@ -308,8 +314,8 @@ weakdeps = ["HTTP"]
|
|||||||
[deps.FileIO.extensions]
|
[deps.FileIO.extensions]
|
||||||
HTTPExt = "HTTP"
|
HTTPExt = "HTTP"
|
||||||
|
|
||||||
[[deps.FileServer]]
|
[[deps.DarkStruct]]
|
||||||
deps = ["HTTP", "JSON3", "Logging", "Oxygen", "UUIDs"]
|
deps = ["HTTP", "JLD2", "JSON3", "Languages", "Logging", "Lux", "Optimisers", "Oxygen", "Random", "UUIDs", "Zygote"]
|
||||||
path = "."
|
path = "."
|
||||||
uuid = "b3f1c2d4-5e6a-4b7c-8d9e-0f1a2b3c4d5e"
|
uuid = "b3f1c2d4-5e6a-4b7c-8d9e-0f1a2b3c4d5e"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
@@ -468,6 +474,12 @@ weakdeps = ["Serialization"]
|
|||||||
[deps.LRUCache.extensions]
|
[deps.LRUCache.extensions]
|
||||||
SerializationExt = ["Serialization"]
|
SerializationExt = ["Serialization"]
|
||||||
|
|
||||||
|
[[deps.Languages]]
|
||||||
|
deps = ["InteractiveUtils", "JSON", "RelocatableFolders"]
|
||||||
|
git-tree-sha1 = "023ac3b12f82da68ed2556c71a134a03e1a11343"
|
||||||
|
uuid = "8ef0a80b-9436-5d2c-a485-80b904378c43"
|
||||||
|
version = "0.4.7"
|
||||||
|
|
||||||
[[deps.LibCURL]]
|
[[deps.LibCURL]]
|
||||||
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
|
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
|
||||||
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
|
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
|
||||||
|
|||||||
15
Project.toml
15
Project.toml
@@ -1,26 +1,39 @@
|
|||||||
name = "FileServer"
|
name = "DarkStruct"
|
||||||
uuid = "b3f1c2d4-5e6a-4b7c-8d9e-0f1a2b3c4d5e"
|
uuid = "b3f1c2d4-5e6a-4b7c-8d9e-0f1a2b3c4d5e"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
authors = ["wardjm@gmail.com"]
|
authors = ["wardjm@gmail.com"]
|
||||||
|
|
||||||
[deps]
|
[deps]
|
||||||
|
AMQPClient = "79c8b4cd-a41a-55fa-907c-fab5288e1383"
|
||||||
HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
|
HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
|
||||||
JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
|
JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
|
||||||
JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
|
JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
|
||||||
|
Languages = "8ef0a80b-9436-5d2c-a485-80b904378c43"
|
||||||
Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
|
Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
|
||||||
Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
|
Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
|
||||||
Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
|
Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
|
||||||
Oxygen = "df9a0d86-3283-4920-82dc-4555fc0d1d8b"
|
Oxygen = "df9a0d86-3283-4920-82dc-4555fc0d1d8b"
|
||||||
|
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
|
||||||
UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
|
UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
|
||||||
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
|
Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
|
||||||
|
|
||||||
[compat]
|
[compat]
|
||||||
|
AMQPClient = "0.5.1"
|
||||||
HTTP = "1.11.0"
|
HTTP = "1.11.0"
|
||||||
JLD2 = "0.6.4"
|
JLD2 = "0.6.4"
|
||||||
JSON3 = "1.14.3"
|
JSON3 = "1.14.3"
|
||||||
|
Languages = "0.4.7"
|
||||||
Logging = "1.11.0"
|
Logging = "1.11.0"
|
||||||
Lux = "1.31.4"
|
Lux = "1.31.4"
|
||||||
Optimisers = "0.4.7"
|
Optimisers = "0.4.7"
|
||||||
Oxygen = "1.10.2"
|
Oxygen = "1.10.2"
|
||||||
|
Random = "1.11.0"
|
||||||
UUIDs = "1.11.0"
|
UUIDs = "1.11.0"
|
||||||
Zygote = "0.7.11"
|
Zygote = "0.7.11"
|
||||||
|
|
||||||
|
[extras]
|
||||||
|
JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
|
||||||
|
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
|
||||||
|
|
||||||
|
[targets]
|
||||||
|
test = ["Test", "JSON3"]
|
||||||
|
|||||||
784
bin/bench.jl
Executable file
784
bin/bench.jl
Executable file
@@ -0,0 +1,784 @@
|
|||||||
|
#!/usr/bin/env julia
|
||||||
|
#
|
||||||
|
# bench.jl: measure end-to-end throughput and server memory for a running DarkStruct.
|
||||||
|
#
|
||||||
|
# Two things make this pipeline awkward to benchmark with off-the-shelf tools
|
||||||
|
# (ab/hey/wrk), and both shape what this script does:
|
||||||
|
#
|
||||||
|
# 1. HTTP latency is not throughput. /upload returns 202 as soon as the bytes
|
||||||
|
# are spooled and a reference is enqueued, and all four stages run afterwards.
|
||||||
|
# So real throughput is the *arrival rate at the terminal sinks*
|
||||||
|
# (done/, text_done/, binary/, failed/), not the response rate. We upload a
|
||||||
|
# corpus, then poll the sinks until the file count stops moving.
|
||||||
|
#
|
||||||
|
# 2. End-to-end throughput doesn't name the slow stage. The four stages run
|
||||||
|
# concurrently behind their own queues, so the pipeline's rate is the
|
||||||
|
# slowest stage's rate and the others are invisible. Directory polling can't
|
||||||
|
# recover them either: every in-flight file sits in spool/ whatever stage it
|
||||||
|
# is at, since stages route by enqueueing rather than by moving bytes. So the
|
||||||
|
# server keeps per-stage counters
|
||||||
|
# (src/stats.jl) and we scrape GET /stats before and after: the deltas give
|
||||||
|
# each stage's throughput, mean service time, and worker utilization, and
|
||||||
|
# utilization is what actually names the bottleneck (see `stage_report`).
|
||||||
|
#
|
||||||
|
# 3. Memory should be flat in file size, and that claim needs checking on two
|
||||||
|
# axes. The workers read bounded prefixes (16+16 bytes to classify, 8 KB to
|
||||||
|
# sniff, 64 KB to language-detect), and intake streams each upload from the
|
||||||
|
# socket to the spool file a chunk at a time (FS_UPLOAD_CHUNK_BYTES, see
|
||||||
|
# src/multipart.jl). So peak RSS should track *concurrency*, not file size:
|
||||||
|
# sweeping --size at fixed --concurrency should be a flat line, and that is
|
||||||
|
# the regression this measures. We sample the server's RSS throughout and
|
||||||
|
# report the high-water mark.
|
||||||
|
#
|
||||||
|
# (Before intake was streamed it buffered each upload whole, several times
|
||||||
|
# over, and a 256 MiB upload grew RSS by ~700 MiB. If a --size sweep ever
|
||||||
|
# slopes upward again, something has started buffering.)
|
||||||
|
#
|
||||||
|
# The harness talks to the server only over HTTP (/upload, /health, /stats) and
|
||||||
|
# reads the on-disk sink layout; it must run on the same machine (sink dirs and
|
||||||
|
# /proc). If /stats is missing (an older build) everything else still works and
|
||||||
|
# the per-stage section is skipped.
|
||||||
|
#
|
||||||
|
# Usage:
|
||||||
|
# julia --project=. -t auto bin/bench.jl [options]
|
||||||
|
#
|
||||||
|
# --url URL server base URL (default: $FS_URL or http://127.0.0.1:8080)
|
||||||
|
# --files N number of files to upload (default: 200)
|
||||||
|
# --size S size of each generated file, e.g. 4k, 512k, 8m, 1g (default: 64k)
|
||||||
|
# --concurrency J uploads in flight at once (default: 8)
|
||||||
|
# --kind K binary | text | mixed: what to generate (default: binary)
|
||||||
|
# --corpus DIR upload an existing directory instead of generating
|
||||||
|
# (the only way to exercise stage 2: point it at real known files)
|
||||||
|
# --keep-corpus don't delete the generated corpus on exit
|
||||||
|
# --pid PID server pid for memory sampling (default: autodetect)
|
||||||
|
# --no-mem skip memory sampling entirely
|
||||||
|
# --no-stats skip the per-stage /stats scrape and its table
|
||||||
|
# --sample-ms MS sink/RSS sampling interval (default: 200)
|
||||||
|
# --timeout SEC give up after this long with no drain progress (default: 120)
|
||||||
|
# --json PATH also write the results as JSON
|
||||||
|
# --force skip the corpus-size safety check
|
||||||
|
#
|
||||||
|
# Examples:
|
||||||
|
# # throughput: many small files, high concurrency
|
||||||
|
# julia --project=. -t auto bin/bench.jl --files 2000 --size 8k --concurrency 32
|
||||||
|
#
|
||||||
|
# # memory: flat in file size? sweep --size with concurrency pinned
|
||||||
|
# julia --project=. -t auto bin/bench.jl --files 4 --size 256m --concurrency 1
|
||||||
|
# julia --project=. -t auto bin/bench.jl --files 2 --size 1g --concurrency 1
|
||||||
|
# julia --project=. -t auto bin/bench.jl --files 8 --size 1g --concurrency 8
|
||||||
|
#
|
||||||
|
# # stage 2 (exiftool) with real known files
|
||||||
|
# julia --project=. -t auto bin/bench.jl --corpus ../training_set --concurrency 16
|
||||||
|
|
||||||
|
using HTTP
|
||||||
|
using JSON3
|
||||||
|
using Random
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------- option parsing
|
||||||
|
|
||||||
|
const DEFAULTS = Dict{String,Any}(
|
||||||
|
"url" => get(ENV, "FS_URL", "http://127.0.0.1:8080"),
|
||||||
|
"files" => 200,
|
||||||
|
"size" => 64 * 1024,
|
||||||
|
"concurrency" => 8,
|
||||||
|
"kind" => "binary",
|
||||||
|
"corpus" => nothing,
|
||||||
|
"keep-corpus" => false,
|
||||||
|
"pid" => nothing,
|
||||||
|
"no-mem" => false,
|
||||||
|
"no-stats" => false,
|
||||||
|
"sample-ms" => 200,
|
||||||
|
"timeout" => 120,
|
||||||
|
"json" => nothing,
|
||||||
|
"force" => false,
|
||||||
|
)
|
||||||
|
|
||||||
|
const FLAGS = ("keep-corpus", "no-mem", "no-stats", "force")
|
||||||
|
|
||||||
|
# Approximate RSS of a freshly started server (Lux + the loaded classifier + the
|
||||||
|
# language detector, measured on Julia 1.12 / -t auto). Only used to notice that
|
||||||
|
# a baseline is inflated by a previous run's un-returned GC memory, so a rough
|
||||||
|
# figure is enough.
|
||||||
|
const FRESH_RSS_HINT = 950 * 1024^2
|
||||||
|
|
||||||
|
"Parse `4k`/`8M`/`1g`/`4096` into a byte count."
|
||||||
|
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)::Dict{String,Any}
|
||||||
|
opts = copy(DEFAULTS)
|
||||||
|
i = 1
|
||||||
|
while i <= length(argv)
|
||||||
|
a = argv[i]
|
||||||
|
startswith(a, "--") || error("unexpected argument: $a (see the header of $(PROGRAM_FILE))")
|
||||||
|
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")
|
||||||
|
val = argv[i+1]
|
||||||
|
opts[key] = key == "size" ? parse_size(val) :
|
||||||
|
key in ("files", "concurrency", "sample-ms") ? parse(Int, val) :
|
||||||
|
key == "timeout" ? parse(Float64, val) :
|
||||||
|
key == "pid" ? parse(Int, val) :
|
||||||
|
val
|
||||||
|
i += 2
|
||||||
|
end
|
||||||
|
opts["kind"] in ("binary", "text", "mixed") ||
|
||||||
|
error("--kind must be binary, text or mixed (got $(opts["kind"]))")
|
||||||
|
opts["files"] >= 1 || error("--files must be >= 1")
|
||||||
|
opts["concurrency"] >= 1 || error("--concurrency must be >= 1")
|
||||||
|
return opts
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------------- dirs
|
||||||
|
#
|
||||||
|
# Resolved from the same environment variables src/config.jl reads, so a server
|
||||||
|
# started with custom dirs is benchmarked correctly. Kept as a standalone table
|
||||||
|
# rather than `using DarkStruct` so the harness doesn't pay to load Lux.
|
||||||
|
|
||||||
|
sinkdirs() = (
|
||||||
|
done = get(ENV, "FS_DONE_DIR", "data/done"),
|
||||||
|
text_done = get(ENV, "FS_TEXT_DONE_DIR", "data/text_done"),
|
||||||
|
binary = get(ENV, "FS_BINARY_DIR", "data/binary"),
|
||||||
|
failed = get(ENV, "FS_FAILED_DIR", "data/failed"),
|
||||||
|
)
|
||||||
|
|
||||||
|
# One directory, not four. Files no longer move between stages: spool/ holds
|
||||||
|
# every in-flight file at every stage, and which stage it has reached lives in
|
||||||
|
# the queue holding its reference (src/worker.jl header). So this depth is
|
||||||
|
# "files in flight", full stop; per-stage depth comes from /stats, which is the
|
||||||
|
# only place that can see it at all.
|
||||||
|
stagedirs() = (spool = get(ENV, "FS_SPOOL_DIR", "data/spool"),)
|
||||||
|
|
||||||
|
"Count work items in `dir`, ignoring the .meta.json sidecars stages 2/4 write."
|
||||||
|
function count_files(dir::AbstractString)::Int
|
||||||
|
isdir(dir) || return 0
|
||||||
|
n = 0
|
||||||
|
for name in readdir(dir)
|
||||||
|
endswith(name, ".meta.json") && continue
|
||||||
|
isfile(joinpath(dir, name)) && (n += 1)
|
||||||
|
end
|
||||||
|
return n
|
||||||
|
end
|
||||||
|
|
||||||
|
counts(dirs) = NamedTuple{keys(dirs)}(map(count_files, values(dirs)))
|
||||||
|
total(c) = sum(values(c))
|
||||||
|
deltas(now_, base) = NamedTuple{keys(now_)}(map(-, values(now_), values(base)))
|
||||||
|
|
||||||
|
# ----------------------------------------------------------------------- memory
|
||||||
|
|
||||||
|
"Read (VmRSS, VmHWM) in bytes for `pid`, or `nothing` if unreadable."
|
||||||
|
function read_rss(pid::Int)
|
||||||
|
rss = hwm = nothing
|
||||||
|
try
|
||||||
|
for line in eachline("/proc/$pid/status")
|
||||||
|
if startswith(line, "VmRSS:")
|
||||||
|
rss = parse(Int, split(line)[2]) * 1024
|
||||||
|
elseif startswith(line, "VmHWM:")
|
||||||
|
hwm = parse(Int, split(line)[2]) * 1024
|
||||||
|
end
|
||||||
|
end
|
||||||
|
catch
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
return (rss === nothing || hwm === nothing) ? nothing : (rss, hwm)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Find the running server process, or `nothing`.
|
||||||
|
|
||||||
|
`pgrep -f` matches against the whole command line, which catches more than the
|
||||||
|
server: any shell launched with the command in its own argv (`sh -c 'julia …
|
||||||
|
bin/server.jl > log'`, a `setsid`/`nohup` wrapper, even the terminal running the
|
||||||
|
benchmark) matches the same pattern. Sampling one of those reports a few MiB of
|
||||||
|
shell as the server's memory: a wrong answer that looks plausible, which is the
|
||||||
|
worst kind.
|
||||||
|
|
||||||
|
So candidates are filtered by what each process *is* (`/proc/<pid>/comm`, the
|
||||||
|
executable name) rather than by what its arguments say. No pattern over argv can
|
||||||
|
make that distinction.
|
||||||
|
"""
|
||||||
|
function detect_pid()
|
||||||
|
out = try
|
||||||
|
readchomp(`pgrep -f "bin/server.jl"`)
|
||||||
|
catch
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
candidates = parse.(Int, split(out))
|
||||||
|
pids = filter(candidates) do pid
|
||||||
|
comm = try
|
||||||
|
readchomp("/proc/$pid/comm")
|
||||||
|
catch
|
||||||
|
return false # exited between pgrep and here
|
||||||
|
end
|
||||||
|
startswith(comm, "julia")
|
||||||
|
end
|
||||||
|
isempty(pids) && return nothing
|
||||||
|
length(pids) > 1 && @warn "multiple server processes matched; sampling the first" pids
|
||||||
|
return first(pids)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Reset the kernel's peak-RSS counter so VmHWM reflects only this run.
|
||||||
|
|
||||||
|
Without it, VmHWM carries the high-water mark from startup (model load) or from
|
||||||
|
an earlier benchmark, which would silently dominate a small run's result.
|
||||||
|
Requires the server to run as the same user; on failure we say so and fall back
|
||||||
|
to sampled VmRSS, which can miss a spike between samples.
|
||||||
|
"""
|
||||||
|
function reset_peak_rss(pid::Int)::Bool
|
||||||
|
try
|
||||||
|
write("/proc/$pid/clear_refs", "5")
|
||||||
|
return true
|
||||||
|
catch
|
||||||
|
return false
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------- per-stage counters
|
||||||
|
#
|
||||||
|
# The server exposes monotonic counters at GET /stats (src/stats.jl). Rates are
|
||||||
|
# ours to compute: scrape once before the run and once after, subtract, divide by
|
||||||
|
# the elapsed *server* clock so a slow scrape doesn't distort the window.
|
||||||
|
|
||||||
|
"Fetch and parse GET /stats, or `nothing` if the server doesn't serve it."
|
||||||
|
function scrape_stats(url::String)
|
||||||
|
try
|
||||||
|
resp = HTTP.get(string(rstrip(url, '/'), "/stats");
|
||||||
|
status_exception = false, retry = false, readtimeout = 5)
|
||||||
|
resp.status == 200 || return nothing
|
||||||
|
return JSON3.read(String(resp.body))
|
||||||
|
catch
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"One stage's activity between two scrapes."
|
||||||
|
struct StageDelta
|
||||||
|
stage::Int
|
||||||
|
name::String
|
||||||
|
workers::Int
|
||||||
|
completed::Int
|
||||||
|
failed::Int
|
||||||
|
bytes::Int
|
||||||
|
busy::Float64 # summed handler seconds across all workers in the pool
|
||||||
|
blocked::Float64 # of `busy`, seconds parked on a full downstream queue
|
||||||
|
peak_depth::Int # deepest its queue got, from the sampler
|
||||||
|
capacity::Int
|
||||||
|
end
|
||||||
|
|
||||||
|
files_per_sec(d::StageDelta, window) = d.completed / max(window, 1e-9)
|
||||||
|
mib_per_sec(d::StageDelta, window) = d.bytes / max(window, 1e-9) / 1024^2
|
||||||
|
"Mean wall time one file spends in one worker of this stage."
|
||||||
|
service_ms(d::StageDelta) = d.completed == 0 ? NaN :
|
||||||
|
(d.busy - d.blocked) / d.completed * 1000
|
||||||
|
"""
|
||||||
|
Fraction of the pool's capacity spent doing this stage's own work.
|
||||||
|
|
||||||
|
Blocked time is subtracted first: a stage parked on a full downstream queue is
|
||||||
|
waiting, not working, and leaving it in would light up every stage upstream of a
|
||||||
|
jam as though each were the jam.
|
||||||
|
"""
|
||||||
|
utilization(d::StageDelta, window) =
|
||||||
|
(d.busy - d.blocked) / max(window * d.workers, 1e-9)
|
||||||
|
blocked_share(d::StageDelta) = d.busy <= 0 ? 0.0 : d.blocked / d.busy
|
||||||
|
|
||||||
|
"Subtract two scrapes into per-stage deltas, folding in sampled peak depths."
|
||||||
|
function stage_deltas(before, after, peak_depth::Dict{Int,Int})
|
||||||
|
out = StageDelta[]
|
||||||
|
for (b, a) in zip(before.stages, after.stages)
|
||||||
|
push!(out, StageDelta(a.stage, String(a.name), a.workers,
|
||||||
|
a.completed - b.completed,
|
||||||
|
a.failed - b.failed,
|
||||||
|
a.bytes - b.bytes,
|
||||||
|
a.busy_seconds - b.busy_seconds,
|
||||||
|
a.blocked_seconds - b.blocked_seconds,
|
||||||
|
get(peak_depth, Int(a.stage), 0),
|
||||||
|
a.queue_capacity))
|
||||||
|
end
|
||||||
|
return out
|
||||||
|
end
|
||||||
|
|
||||||
|
"JSON has no NaN. A stage that completed nothing has no service time, and `null`
|
||||||
|
is the honest way to say that; writing NaN just makes JSON3 throw."
|
||||||
|
json_num(x::Real) = isfinite(x) ? x : nothing
|
||||||
|
|
||||||
|
pad(s, n) = rpad(string(s), n)
|
||||||
|
lpad_(s, n) = lpad(string(s), n)
|
||||||
|
|
||||||
|
"""
|
||||||
|
Print the per-stage table and say which stage is the bottleneck.
|
||||||
|
|
||||||
|
The verdict reads utilization, not throughput: in a pipeline every stage
|
||||||
|
completes the same files, so at steady state they all report nearly the same
|
||||||
|
files/s regardless of which one is the constraint. What separates them is how
|
||||||
|
hard each pool had to work to keep up: the bottleneck is pinned near 1.0 while
|
||||||
|
its neighbours idle.
|
||||||
|
"""
|
||||||
|
function stage_report(deltas::Vector{StageDelta}, window::Float64)
|
||||||
|
println("PER-STAGE (server counters, delta over the end-to-end window)")
|
||||||
|
println(" stage files/s MiB/s svc ms util blocked peak queue failed")
|
||||||
|
for d in deltas
|
||||||
|
svc = service_ms(d)
|
||||||
|
println(" $(d.stage) $(pad(d.name, 10)) " *
|
||||||
|
lpad_(fmt(files_per_sec(d, window)), 8) * " " *
|
||||||
|
lpad_(fmt(mib_per_sec(d, window)), 8) * " " *
|
||||||
|
lpad_(isnan(svc) ? "—" : fmt(svc, 1), 8) * " " *
|
||||||
|
lpad_(fmt(utilization(d, window)), 6) * " " *
|
||||||
|
lpad_(fmt(blocked_share(d) * 100, 0) * "%", 7) * " " *
|
||||||
|
lpad_("$(d.peak_depth)/$(d.capacity)", 11) * " " *
|
||||||
|
lpad_(d.failed, 7))
|
||||||
|
end
|
||||||
|
|
||||||
|
# Per-stage files/s are not comparable across rows and saying so costs one
|
||||||
|
# line: the stages process different subsets (stage 2 only known files,
|
||||||
|
# stage 4 only text), so a low rate can mean "little work arrived here"
|
||||||
|
# rather than "slow". Utilization is the column that compares.
|
||||||
|
println(" (files/s counts only files routed to that stage; svc is per-file wall time in")
|
||||||
|
println(" one worker; util = (busy − blocked) / (window × workers))")
|
||||||
|
|
||||||
|
worked = filter(d -> d.completed > 0, deltas)
|
||||||
|
if isempty(worked)
|
||||||
|
println(" (no stage completed a file in this window)")
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
top = argmax(d -> utilization(d, window), worked)
|
||||||
|
println(" bottleneck stage $(top.stage) ($(top.name)) at " *
|
||||||
|
"$(fmt(utilization(top, window) * 100, 0))% utilization of " *
|
||||||
|
"$(top.workers) worker(s)")
|
||||||
|
if utilization(top, window) < 0.5
|
||||||
|
println(" ...but no stage is near saturated: the pipeline is " *
|
||||||
|
"waiting on intake,\n not on itself. Raise --concurrency " *
|
||||||
|
"or --files to load it properly.")
|
||||||
|
end
|
||||||
|
for d in worked
|
||||||
|
blocked_share(d) > 0.25 && println(" ! stage $(d.stage) ($(d.name)) spent " *
|
||||||
|
"$(fmt(blocked_share(d) * 100, 0))% of its time parked on a full downstream " *
|
||||||
|
"queue.\n It is being held up by the stage after it, not doing that work itself.")
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------------ corpus
|
||||||
|
|
||||||
|
const WORDS = split("the quick brown fox jumps over a lazy dog while parsing " *
|
||||||
|
"headers and spooling bytes onto disk for later enrichment " *
|
||||||
|
"because throughput matters more than latency here")
|
||||||
|
|
||||||
|
"Write one file of exactly `size` bytes, in bounded chunks so the generator
|
||||||
|
itself never holds a whole 1 GB file in memory."
|
||||||
|
function write_file(path::AbstractString, size::Int, kind::Symbol, rng)
|
||||||
|
chunk = 1024 * 1024
|
||||||
|
open(path, "w") do io
|
||||||
|
remaining = size
|
||||||
|
while remaining > 0
|
||||||
|
n = min(chunk, remaining)
|
||||||
|
if kind === :binary
|
||||||
|
write(io, rand(rng, UInt8, n))
|
||||||
|
else
|
||||||
|
buf = IOBuffer()
|
||||||
|
while buf.size < n
|
||||||
|
print(buf, rand(rng, WORDS), rand(rng) < 0.06 ? ".\n" : " ")
|
||||||
|
end
|
||||||
|
write(io, take!(buf)[1:n])
|
||||||
|
end
|
||||||
|
remaining -= n
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"Generate the corpus and return (dir, paths, total_bytes)."
|
||||||
|
function make_corpus(opts)
|
||||||
|
n, size, kind = opts["files"], opts["size"], opts["kind"]
|
||||||
|
totalbytes = n * size
|
||||||
|
if totalbytes > 16 * 1024^3 && !opts["force"]
|
||||||
|
error("corpus would be $(human(totalbytes)) on disk; pass --force if that's intended")
|
||||||
|
end
|
||||||
|
dir = mktempdir(; prefix = "fsbench-")
|
||||||
|
rng = MersenneTwister(1234)
|
||||||
|
paths = String[]
|
||||||
|
for i in 1:n
|
||||||
|
k = kind == "mixed" ? (isodd(i) ? :binary : :text) :
|
||||||
|
kind == "text" ? :text : :binary
|
||||||
|
ext = k === :text ? "txt" : "bin"
|
||||||
|
path = joinpath(dir, "bench-$(lpad(i, 6, '0')).$ext")
|
||||||
|
write_file(path, size, k, rng)
|
||||||
|
push!(paths, path)
|
||||||
|
end
|
||||||
|
return dir, paths, totalbytes
|
||||||
|
end
|
||||||
|
|
||||||
|
function existing_corpus(dir::AbstractString)
|
||||||
|
isdir(dir) || error("--corpus is not a directory: $dir")
|
||||||
|
paths = sort(filter(isfile, readdir(dir; join = true)))
|
||||||
|
isempty(paths) && error("--corpus directory is empty: $dir")
|
||||||
|
return paths, sum(filesize, paths)
|
||||||
|
end
|
||||||
|
|
||||||
|
# ----------------------------------------------------------------------- upload
|
||||||
|
|
||||||
|
struct Upload
|
||||||
|
status::Int # HTTP status, or 0 if the request threw
|
||||||
|
accepted::Int # jobs the server actually queued (from the 202/503 body)
|
||||||
|
seconds::Float64
|
||||||
|
end
|
||||||
|
|
||||||
|
"POST one file as multipart/form-data and report what the server accepted."
|
||||||
|
function upload_one(url::String, path::String)::Upload
|
||||||
|
t0 = time()
|
||||||
|
try
|
||||||
|
form = HTTP.Form(["file" => HTTP.Multipart(basename(path), open(path, "r"),
|
||||||
|
"application/octet-stream")])
|
||||||
|
resp = HTTP.post(url, [], form; status_exception = false, retry = false)
|
||||||
|
# 202 and 503 both carry an `accepted` array: a partially-accepted batch
|
||||||
|
# still queued those jobs, and they will show up in the sinks.
|
||||||
|
acc = try
|
||||||
|
length(JSON3.read(String(resp.body)).accepted)
|
||||||
|
catch
|
||||||
|
resp.status == 202 ? 1 : 0
|
||||||
|
end
|
||||||
|
return Upload(resp.status, acc, time() - t0)
|
||||||
|
catch e
|
||||||
|
e isa InterruptException && rethrow()
|
||||||
|
@warn "upload failed" file = basename(path) exception = e
|
||||||
|
return Upload(0, 0, time() - t0)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Upload every path, at most `concurrency` in flight.
|
||||||
|
|
||||||
|
A bounded set of worker tasks pulling from a shared index keeps exactly
|
||||||
|
`concurrency` requests in flight for the whole run, unlike batching, where each
|
||||||
|
batch stalls on its slowest (largest) file and the real concurrency sags.
|
||||||
|
"""
|
||||||
|
function upload_all(url::String, paths::Vector{String}, concurrency::Int)
|
||||||
|
endpoint = string(rstrip(url, '/'), "/upload")
|
||||||
|
results = Vector{Upload}(undef, length(paths))
|
||||||
|
next = Threads.Atomic{Int}(1)
|
||||||
|
@sync for _ in 1:min(concurrency, length(paths))
|
||||||
|
Threads.@spawn while true
|
||||||
|
i = Threads.atomic_add!(next, 1)
|
||||||
|
i > length(paths) && break
|
||||||
|
results[i] = upload_one(endpoint, paths[i])
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return results
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------- formatting
|
||||||
|
|
||||||
|
function human(bytes::Real)
|
||||||
|
b = Float64(bytes)
|
||||||
|
for unit in ("B", "KiB", "MiB", "GiB", "TiB")
|
||||||
|
(abs(b) < 1024 || unit == "TiB") && return "$(round(b; digits = 2)) $unit"
|
||||||
|
b /= 1024
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
fmt(x::Real, digits::Int = 2) = string(round(Float64(x); digits = digits))
|
||||||
|
|
||||||
|
function percentile(sorted::Vector{Float64}, p::Float64)
|
||||||
|
isempty(sorted) && return NaN
|
||||||
|
idx = clamp(ceil(Int, p * length(sorted)), 1, length(sorted))
|
||||||
|
return sorted[idx]
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------------- main
|
||||||
|
|
||||||
|
function main(argv)
|
||||||
|
opts = parse_args(argv)
|
||||||
|
url = string(opts["url"])
|
||||||
|
|
||||||
|
# Fail fast and clearly if there's no server, rather than reporting a run of zeros.
|
||||||
|
try
|
||||||
|
HTTP.get(string(rstrip(url, '/'), "/health"); retry = false, readtimeout = 5)
|
||||||
|
catch e
|
||||||
|
println(stderr, "cannot reach $url/health. Is the server running?")
|
||||||
|
println(stderr, " start it with: julia --project=. -t auto bin/server.jl")
|
||||||
|
return 1
|
||||||
|
end
|
||||||
|
|
||||||
|
sinks, stages = sinkdirs(), stagedirs()
|
||||||
|
|
||||||
|
# Leftovers mid-pipeline would land in the sinks during our window and be
|
||||||
|
# counted as our throughput, so say so up front rather than quietly skewing.
|
||||||
|
pending = total(counts(stages))
|
||||||
|
pending > 0 && @warn "pipeline is not idle: $pending file(s) still in spool/; " *
|
||||||
|
"throughput will include their completions"
|
||||||
|
|
||||||
|
# --- corpus
|
||||||
|
generated = opts["corpus"] === nothing
|
||||||
|
corpusdir, paths, corpusbytes = if generated
|
||||||
|
print("generating corpus: $(opts["files"]) × $(human(opts["size"])) ($(opts["kind"]))… ")
|
||||||
|
t = time()
|
||||||
|
d, p, b = make_corpus(opts)
|
||||||
|
println("done in $(fmt(time() - t))s → $d")
|
||||||
|
d, p, b
|
||||||
|
else
|
||||||
|
p, b = existing_corpus(String(opts["corpus"]))
|
||||||
|
println("using corpus: $(length(p)) file(s), $(human(b)) from $(opts["corpus"])")
|
||||||
|
String(opts["corpus"]), p, b
|
||||||
|
end
|
||||||
|
|
||||||
|
try
|
||||||
|
pid = opts["no-mem"] ? nothing : something(opts["pid"], detect_pid(), Some(nothing))
|
||||||
|
if pid === nothing && !opts["no-mem"]
|
||||||
|
@warn "could not find the server process; skipping memory (pass --pid PID)"
|
||||||
|
end
|
||||||
|
|
||||||
|
baseline_rss = nothing
|
||||||
|
peak_reset = false
|
||||||
|
if pid !== nothing
|
||||||
|
r = read_rss(pid)
|
||||||
|
r === nothing && (@warn "cannot read /proc/$pid/status; skipping memory"; pid = nothing)
|
||||||
|
if pid !== nothing
|
||||||
|
peak_reset = reset_peak_rss(pid)
|
||||||
|
peak_reset || @warn "could not reset the peak-RSS counter (/proc/$pid/clear_refs); " *
|
||||||
|
"reporting sampled RSS only"
|
||||||
|
# Baseline is read *after* the reset, not before. VmHWM restarts
|
||||||
|
# from whatever RSS is at the moment of the reset, so a baseline
|
||||||
|
# sampled earlier is measured against a different origin, and if
|
||||||
|
# the GC hands memory back in between, the run reports negative
|
||||||
|
# growth, which is nonsense on its face.
|
||||||
|
r2 = read_rss(pid)
|
||||||
|
baseline_rss = r2 === nothing ? r[1] : r2[1]
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
if Threads.nthreads() == 1 && opts["size"] > 64 * 1024^2
|
||||||
|
@warn "running with 1 thread and large files: the sampler shares a thread with " *
|
||||||
|
"blocking file reads, so the RSS curve will be coarse. Prefer -t auto."
|
||||||
|
end
|
||||||
|
|
||||||
|
base_sinks = counts(sinks)
|
||||||
|
interval = opts["sample-ms"] / 1000
|
||||||
|
|
||||||
|
# --- per-stage counters: the "before" half of the delta.
|
||||||
|
stats_before = opts["no-stats"] ? nothing : scrape_stats(url)
|
||||||
|
if stats_before === nothing && !opts["no-stats"]
|
||||||
|
@warn "no /stats endpoint on this server; skipping the per-stage table " *
|
||||||
|
"(the server predates src/stats.jl)"
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- sampler: RSS curve, spool depth, and queue depths.
|
||||||
|
stop = Threads.Atomic{Bool}(false)
|
||||||
|
rss_samples = Float64[]
|
||||||
|
depth_max = Dict(k => 0 for k in keys(stages))
|
||||||
|
queue_peak = Dict{Int,Int}()
|
||||||
|
sampler = Threads.@spawn begin
|
||||||
|
while !stop[]
|
||||||
|
if pid !== nothing
|
||||||
|
r = read_rss(pid)
|
||||||
|
r !== nothing && push!(rss_samples, Float64(r[1]))
|
||||||
|
end
|
||||||
|
d = counts(stages)
|
||||||
|
for k in keys(d)
|
||||||
|
depth_max[k] = max(depth_max[k], getfield(d, k))
|
||||||
|
end
|
||||||
|
# Queue depth, unlike directory depth, can't be missed by a slow
|
||||||
|
# sample in the same way, since a file's *reference* sits in the
|
||||||
|
# queue for the whole time it waits, so this is the depth the stage
|
||||||
|
# table reports.
|
||||||
|
if stats_before !== nothing
|
||||||
|
s = scrape_stats(url)
|
||||||
|
s === nothing || for st in s.stages
|
||||||
|
k = Int(st.stage)
|
||||||
|
queue_peak[k] = max(get(queue_peak, k, 0), Int(st.queue_depth))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
sleep(interval)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- intake
|
||||||
|
println("uploading $(length(paths)) file(s) at concurrency $(opts["concurrency"])…")
|
||||||
|
t_start = time()
|
||||||
|
ups = upload_all(url, paths, opts["concurrency"])
|
||||||
|
t_intake_end = time()
|
||||||
|
|
||||||
|
accepted = sum(u.accepted for u in ups)
|
||||||
|
n_202 = count(u -> u.status == 202, ups)
|
||||||
|
n_503 = count(u -> u.status == 503, ups)
|
||||||
|
n_err = count(u -> !(u.status in (202, 503)), ups)
|
||||||
|
intake_secs = t_intake_end - t_start
|
||||||
|
lat = sort([u.seconds for u in ups])
|
||||||
|
|
||||||
|
println(" intake: $accepted job(s) accepted in $(fmt(intake_secs))s " *
|
||||||
|
"($(fmt(accepted / max(intake_secs, 1e-9))) files/s, " *
|
||||||
|
"$(fmt(corpusbytes / max(intake_secs, 1e-9) / 1024^2)) MiB/s)")
|
||||||
|
n_503 > 0 && println(" backpressure: $n_503 request(s) got 503 (intake queue full)")
|
||||||
|
n_err > 0 && println(" errors: $n_err request(s) failed or returned an unexpected status")
|
||||||
|
|
||||||
|
# --- drain: poll the terminal sinks until they stop moving.
|
||||||
|
println("draining (polling sinks every $(opts["sample-ms"])ms)…")
|
||||||
|
completed = 0
|
||||||
|
t_last_progress = time()
|
||||||
|
t_last_completion = t_intake_end
|
||||||
|
timed_out = false
|
||||||
|
while completed < accepted
|
||||||
|
sleep(interval)
|
||||||
|
c = total(deltas(counts(sinks), base_sinks))
|
||||||
|
if c > completed
|
||||||
|
completed = c
|
||||||
|
t_last_completion = time()
|
||||||
|
t_last_progress = t_last_completion
|
||||||
|
elseif time() - t_last_progress > opts["timeout"]
|
||||||
|
timed_out = true
|
||||||
|
break
|
||||||
|
end
|
||||||
|
end
|
||||||
|
# Scrape before stopping the sampler, so the window closes as near the
|
||||||
|
# last completion as we can manage.
|
||||||
|
stats_after = stats_before === nothing ? nothing : scrape_stats(url)
|
||||||
|
stop[] = true
|
||||||
|
wait(sampler)
|
||||||
|
|
||||||
|
sink_delta = deltas(counts(sinks), base_sinks)
|
||||||
|
e2e_secs = t_last_completion - t_start
|
||||||
|
|
||||||
|
# The stage window is the server's own clock across the two scrapes, not
|
||||||
|
# e2e_secs: it starts a scrape earlier and ends a scrape later, and using
|
||||||
|
# our wall time against its counters would misattribute the difference.
|
||||||
|
stage_window, stage_delta = if stats_after === nothing
|
||||||
|
(0.0, StageDelta[])
|
||||||
|
else
|
||||||
|
(Float64(stats_after.now - stats_before.now),
|
||||||
|
stage_deltas(stats_before, stats_after, queue_peak))
|
||||||
|
end
|
||||||
|
|
||||||
|
final_rss = pid === nothing ? nothing : read_rss(pid)
|
||||||
|
peak_rss = if final_rss !== nothing && peak_reset
|
||||||
|
final_rss[2] # kernel VmHWM: catches spikes between samples
|
||||||
|
elseif !isempty(rss_samples)
|
||||||
|
maximum(rss_samples)
|
||||||
|
else
|
||||||
|
nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- report
|
||||||
|
println()
|
||||||
|
println("=" ^ 68)
|
||||||
|
println("corpus $(length(paths)) file(s), $(human(corpusbytes)) total, " *
|
||||||
|
"$(human(corpusbytes / length(paths))) avg")
|
||||||
|
println("concurrency $(opts["concurrency"])")
|
||||||
|
println()
|
||||||
|
println("INTAKE (HTTP 202: bytes spooled, not processed)")
|
||||||
|
println(" accepted $accepted of $(length(paths)) [202: $n_202, 503: $n_503, error: $n_err]")
|
||||||
|
println(" wall $(fmt(intake_secs))s")
|
||||||
|
println(" rate $(fmt(accepted / max(intake_secs, 1e-9))) files/s, " *
|
||||||
|
"$(fmt(corpusbytes / max(intake_secs, 1e-9) / 1024^2)) MiB/s")
|
||||||
|
println(" latency p50 $(fmt(percentile(lat, 0.5) * 1000, 1))ms " *
|
||||||
|
"p95 $(fmt(percentile(lat, 0.95) * 1000, 1))ms " *
|
||||||
|
"max $(fmt(percentile(lat, 1.0) * 1000, 1))ms")
|
||||||
|
println()
|
||||||
|
println("END-TO-END (files reaching a terminal sink)")
|
||||||
|
println(" completed $completed of $accepted accepted" * (timed_out ? " ** TIMED OUT **" : ""))
|
||||||
|
println(" wall $(fmt(e2e_secs))s (first upload → last completion)")
|
||||||
|
println(" throughput $(fmt(completed / max(e2e_secs, 1e-9))) files/s, " *
|
||||||
|
"$(fmt(corpusbytes / max(e2e_secs, 1e-9) / 1024^2)) MiB/s")
|
||||||
|
println(" sinks done $(sink_delta.done) text_done $(sink_delta.text_done) " *
|
||||||
|
"binary $(sink_delta.binary) failed $(sink_delta.failed)")
|
||||||
|
println(" peak dir depth spool $(depth_max[:spool]) known $(depth_max[:known]) " *
|
||||||
|
"unknown $(depth_max[:unknown]) text $(depth_max[:text])")
|
||||||
|
println()
|
||||||
|
if !isempty(stage_delta)
|
||||||
|
stage_report(stage_delta, stage_window)
|
||||||
|
println()
|
||||||
|
end
|
||||||
|
if peak_rss !== nothing
|
||||||
|
println("SERVER MEMORY (pid $pid)")
|
||||||
|
println(" baseline RSS $(human(baseline_rss))")
|
||||||
|
println(" peak RSS $(human(peak_rss))" *
|
||||||
|
(peak_reset ? " (kernel VmHWM, reset at start)" : " (sampled; may miss spikes)"))
|
||||||
|
growth = peak_rss - baseline_rss
|
||||||
|
if growth <= 0
|
||||||
|
# RSS never got back to where it started, so the run's own cost
|
||||||
|
# is below the noise floor of the server settling after startup.
|
||||||
|
# Printing a negative "growth" would invite reading a memory
|
||||||
|
# *saving* into what is really "too small to measure here".
|
||||||
|
println(" growth none measurable (peak never exceeded the baseline)")
|
||||||
|
println(" The baseline was still falling when we sampled it. Give the")
|
||||||
|
println(" server ~30s to settle after startup for a comparable figure.")
|
||||||
|
else
|
||||||
|
println(" growth $(human(growth))")
|
||||||
|
println(" per in-flight $(human(growth / opts["concurrency"])) " *
|
||||||
|
"at $(human(corpusbytes / length(paths))) avg file size")
|
||||||
|
println(" (should not grow with file size: intake streams to disk)")
|
||||||
|
end
|
||||||
|
# Julia's GC returns memory to the OS lazily, so a second run on the
|
||||||
|
# same process starts from an inflated baseline and under-reports
|
||||||
|
# growth. Absolute peak is the number to trust across runs.
|
||||||
|
baseline_rss > 1.3 * FRESH_RSS_HINT &&
|
||||||
|
println(" ! baseline is well above a fresh start ($(human(FRESH_RSS_HINT))): the GC " *
|
||||||
|
"has not\n returned memory from earlier work. Compare " *
|
||||||
|
"absolute peak, or restart\n the server for a clean growth figure.")
|
||||||
|
else
|
||||||
|
println("SERVER MEMORY not sampled")
|
||||||
|
end
|
||||||
|
println("=" ^ 68)
|
||||||
|
|
||||||
|
sink_delta.failed > 0 &&
|
||||||
|
println("\nnote: $(sink_delta.failed) file(s) landed in $(sinks.failed); check the server log.")
|
||||||
|
timed_out &&
|
||||||
|
println("\nnote: drain stalled with $(accepted - completed) file(s) outstanding. " *
|
||||||
|
"Check the server log and spool/; raise --timeout if the pipeline is just slow.")
|
||||||
|
|
||||||
|
if opts["json"] !== nothing
|
||||||
|
result = (
|
||||||
|
url, concurrency = opts["concurrency"], kind = opts["kind"],
|
||||||
|
files = length(paths), corpus_bytes = corpusbytes,
|
||||||
|
avg_file_bytes = corpusbytes / length(paths),
|
||||||
|
intake = (; accepted, n_202, n_503, n_err, seconds = intake_secs,
|
||||||
|
files_per_sec = accepted / max(intake_secs, 1e-9),
|
||||||
|
p50_ms = json_num(percentile(lat, 0.5) * 1000),
|
||||||
|
p95_ms = json_num(percentile(lat, 0.95) * 1000),
|
||||||
|
max_ms = json_num(percentile(lat, 1.0) * 1000)),
|
||||||
|
end_to_end = (; completed, seconds = e2e_secs, timed_out,
|
||||||
|
files_per_sec = completed / max(e2e_secs, 1e-9),
|
||||||
|
mib_per_sec = corpusbytes / max(e2e_secs, 1e-9) / 1024^2,
|
||||||
|
sinks = sink_delta, peak_stage_depth = depth_max),
|
||||||
|
stages = [(; stage = d.stage, name = d.name, workers = d.workers,
|
||||||
|
completed = d.completed, failed = d.failed, bytes = d.bytes,
|
||||||
|
busy_seconds = d.busy, blocked_seconds = d.blocked,
|
||||||
|
window_seconds = stage_window,
|
||||||
|
files_per_sec = files_per_sec(d, stage_window),
|
||||||
|
mib_per_sec = mib_per_sec(d, stage_window),
|
||||||
|
service_ms = json_num(service_ms(d)),
|
||||||
|
utilization = utilization(d, stage_window),
|
||||||
|
blocked_share = blocked_share(d),
|
||||||
|
peak_queue_depth = d.peak_depth,
|
||||||
|
queue_capacity = d.capacity) for d in stage_delta],
|
||||||
|
memory = (; pid, baseline_rss, peak_rss, peak_is_kernel_hwm = peak_reset,
|
||||||
|
growth = peak_rss === nothing ? nothing : peak_rss - baseline_rss,
|
||||||
|
samples = rss_samples),
|
||||||
|
)
|
||||||
|
open(String(opts["json"]), "w") do io
|
||||||
|
JSON3.write(io, result)
|
||||||
|
end
|
||||||
|
println("\nwrote $(opts["json"])")
|
||||||
|
end
|
||||||
|
|
||||||
|
return timed_out ? 1 : 0
|
||||||
|
finally
|
||||||
|
if generated && !opts["keep-corpus"]
|
||||||
|
rm(corpusdir; recursive = true, force = true)
|
||||||
|
elseif generated
|
||||||
|
println("\nkept corpus: $corpusdir")
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
if abspath(PROGRAM_FILE) == @__FILE__
|
||||||
|
exit(main(ARGS))
|
||||||
|
end
|
||||||
427
bin/bench_model.jl
Executable file
427
bin/bench_model.jl
Executable file
@@ -0,0 +1,427 @@
|
|||||||
|
#!/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
|
||||||
535
bin/bench_stage1.jl
Normal file
535
bin/bench_stage1.jl
Normal file
@@ -0,0 +1,535 @@
|
|||||||
|
#!/usr/bin/env julia
|
||||||
|
#
|
||||||
|
# bench_stage1.jl: take stage 1 apart and find the slowest component.
|
||||||
|
#
|
||||||
|
# bin/bench.jl reports stage 1 as a single number (the wall time of
|
||||||
|
# `handle_classify_job` under whole-pipeline contention) and bin/bench_model.jl
|
||||||
|
# takes the *classifier* apart. Neither answers "which part of stage 1 costs the
|
||||||
|
# most?", because stage 1 is more than the model. Per file it does:
|
||||||
|
#
|
||||||
|
# classify filesize + two 16-byte reads + a 32x1 forward pass
|
||||||
|
# read_features open, read head, seek, read tail, scale to Float32
|
||||||
|
# Lux.apply the network on a feature vector already in memory
|
||||||
|
# enqueue push a Job reference onto the downstream bounded queue
|
||||||
|
# logging two @info lines ("classified file", "routed to ...")
|
||||||
|
#
|
||||||
|
# This script times each of those in isolation, then times the real
|
||||||
|
# `handle_classify_job` end to end so the parts can be checked against the whole.
|
||||||
|
# The two knobs that most change the answer get their own sweeps:
|
||||||
|
#
|
||||||
|
# * Logger. The server runs `FlushLogger(ConsoleLogger(stderr))`, which formats
|
||||||
|
# and flushes every message. Under redirect (a log file, journald) that is a
|
||||||
|
# syscall per line, two lines per file, on the hot path. We time the handler
|
||||||
|
# under a null logger, a formatting-but-discarding logger, and the real
|
||||||
|
# flushing-to-file logger, so the cost of logging is a subtraction, not a
|
||||||
|
# guess. This sweep is what demoted stage 1's per-file lines to `@debug`
|
||||||
|
# (see the note in src/worker.jl); the standalone `logging (...)` rows below
|
||||||
|
# still price a *formatted* line, i.e. what those lines cost when switched
|
||||||
|
# back on with `JULIA_DEBUG=DarkStruct`, while the handler rows show what the
|
||||||
|
# stage pays with them off.
|
||||||
|
# * Concurrency. Components that own a lock (the queue's condition, the
|
||||||
|
# logger's stream) don't scale, and the ranking at one worker need not be the
|
||||||
|
# ranking at sixteen. The `--threads` sweep runs the full handler across
|
||||||
|
# worker counts.
|
||||||
|
#
|
||||||
|
# Usage:
|
||||||
|
# julia --project=. -t auto bin/bench_stage1.jl [options]
|
||||||
|
#
|
||||||
|
# --files N files per timed pass for whole-corpus benchmarks (default: 2000)
|
||||||
|
# --reps N calls per timed pass for per-call benchmarks (default: 20000)
|
||||||
|
# --trials N timed passes; the minimum is reported (default: 5)
|
||||||
|
# --size SPEC corpus file size (default: 64k)
|
||||||
|
# --dir PATH working directory for the corpus (default: a temp dir under data/)
|
||||||
|
# --model PATH classifier artifact (default: $FS_MODEL_PATH or model/classifier.jld2)
|
||||||
|
# --threads LIST worker counts for the concurrency sweep (default: 1,2,4,8,nthreads)
|
||||||
|
# --no-threads skip the concurrency sweep
|
||||||
|
# --json PATH also write the results as JSON
|
||||||
|
#
|
||||||
|
# Reported times are the *minimum* over trials: the floor is the signal and
|
||||||
|
# everything above it is scheduler, page-cache and GC noise.
|
||||||
|
|
||||||
|
using DarkStruct
|
||||||
|
using Lux
|
||||||
|
using JSON3
|
||||||
|
using Logging
|
||||||
|
using Printf
|
||||||
|
using Random
|
||||||
|
|
||||||
|
const FS = DarkStruct
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------- option parsing
|
||||||
|
|
||||||
|
const DEFAULTS = Dict{String,Any}(
|
||||||
|
"files" => 2000,
|
||||||
|
"reps" => 20_000,
|
||||||
|
"trials" => 5,
|
||||||
|
"size" => "64k",
|
||||||
|
"dir" => nothing,
|
||||||
|
"model" => get(ENV, "FS_MODEL_PATH", "model/classifier.jld2"),
|
||||||
|
"threads" => nothing,
|
||||||
|
"no-threads" => false,
|
||||||
|
"json" => nothing,
|
||||||
|
)
|
||||||
|
|
||||||
|
const FLAGS = ("no-threads",)
|
||||||
|
const INTS = ("files", "reps", "trials")
|
||||||
|
|
||||||
|
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 INTS ? 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)
|
||||||
|
|
||||||
|
"""
|
||||||
|
best_of(pass, prepare; trials) -> ns_per_op
|
||||||
|
|
||||||
|
Run `pass()` `trials` times and report the fastest, in nanoseconds per operation
|
||||||
|
(`pass` returns the number of operations it performed). `prepare()` runs before
|
||||||
|
each pass and is *not* timed: that is where a benchmark resets whatever its
|
||||||
|
last pass consumed (today: the queues). `pass` comes first so callers can pass it as a `do` block.
|
||||||
|
|
||||||
|
The first pass is thrown away: it pays Julia's JIT compilation, which on calls
|
||||||
|
this small is orders of magnitude more than the thing being measured.
|
||||||
|
"""
|
||||||
|
function best_of(pass, prepare; trials::Int)
|
||||||
|
best = Inf
|
||||||
|
for t in 0:trials
|
||||||
|
prepare()
|
||||||
|
GC.gc()
|
||||||
|
t0 = time_ns()
|
||||||
|
n = pass()
|
||||||
|
dt = Float64(time_ns() - t0)
|
||||||
|
t == 0 && continue # warm-up: compiled, not measured
|
||||||
|
best = min(best, dt / n)
|
||||||
|
end
|
||||||
|
return best
|
||||||
|
end
|
||||||
|
|
||||||
|
noop() = nothing
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------- 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
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
rate(ns::Real) = 1e9 / max(ns, 1e-9)
|
||||||
|
|
||||||
|
rule(n = 78) = println("-" ^ n)
|
||||||
|
|
||||||
|
function header(title)
|
||||||
|
println()
|
||||||
|
println(title)
|
||||||
|
rule()
|
||||||
|
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
|
||||||
|
|
||||||
|
"""
|
||||||
|
make_corpus(cfg, n, size, rng) -> Vector{Job}
|
||||||
|
|
||||||
|
Create `n` spooled files and the `Job` references a stage-1 worker would dequeue
|
||||||
|
for them: the exact input `handle_classify_job` sees.
|
||||||
|
"""
|
||||||
|
function make_corpus(cfg::FS.Config, n::Int, size::Int, rng)
|
||||||
|
jobs = FS.Job[]
|
||||||
|
for i in 1:n
|
||||||
|
id, path = FS.spool_path(cfg, @sprintf("bench-%06d.bin", i))
|
||||||
|
write_file(path, size, rng)
|
||||||
|
push!(jobs, FS.Job(id, basename(path), path, size, time()))
|
||||||
|
end
|
||||||
|
return jobs
|
||||||
|
end
|
||||||
|
|
||||||
|
# There is no `respool!` step any more. Stage 1 used to consume its input by
|
||||||
|
# renaming each file into known/ or unknown/, so every repeated pass had to put
|
||||||
|
# the corpus back first. It now routes by enqueueing alone and leaves the file in
|
||||||
|
# spool/, so the corpus is reusable as-is and the only per-pass reset is draining
|
||||||
|
# the queues.
|
||||||
|
|
||||||
|
"Drain a queue without blocking, so the next pass starts from empty."
|
||||||
|
function drain!(q::FS.ChannelQueue)
|
||||||
|
while length(q) > 0
|
||||||
|
FS.dequeue!(q)
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------- loggers
|
||||||
|
|
||||||
|
"""
|
||||||
|
with_logger_named(name, path, f)
|
||||||
|
|
||||||
|
Run `f` under one of the three loggers the cost of logging is bracketed by:
|
||||||
|
|
||||||
|
* `:null` `NullLogger`: the `@info` macro's own overhead, nothing else.
|
||||||
|
* `:format` `ConsoleLogger` to `devnull`: message formatting and key/value
|
||||||
|
interpolation, but no I/O.
|
||||||
|
* `:flush` `FlushLogger(ConsoleLogger(io))` to a real file: what
|
||||||
|
`DarkStruct.run` installs, under the redirect it was written for.
|
||||||
|
* `:debug` the same, at `Debug` level: the stage's per-file lines are
|
||||||
|
`@debug`, so this is the equivalent of running the server with
|
||||||
|
`JULIA_DEBUG=DarkStruct` and the only setting under which they
|
||||||
|
are emitted at all.
|
||||||
|
"""
|
||||||
|
function with_logger_named(f, which::Symbol, path::AbstractString)
|
||||||
|
if which === :null
|
||||||
|
return with_logger(f, NullLogger())
|
||||||
|
elseif which === :format
|
||||||
|
return with_logger(f, ConsoleLogger(devnull))
|
||||||
|
elseif which === :flush || which === :debug
|
||||||
|
level = which === :debug ? Logging.Debug : Logging.Info
|
||||||
|
return open(path, "w") do io
|
||||||
|
with_logger(f, FS.FlushLogger(ConsoleLogger(io, level)))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
error("unknown logger: $which")
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------ components
|
||||||
|
|
||||||
|
"""
|
||||||
|
component_rows(cfg, clf, jobs, opts) -> Vector
|
||||||
|
|
||||||
|
Time each piece of stage 1 on its own. Per-call pieces (`filesize`,
|
||||||
|
`read_features`, `Lux.apply`, `classify`, the log lines) run `reps` times over
|
||||||
|
the corpus; whole-corpus pieces (`enqueue_blocking!`, the full handler) run once
|
||||||
|
per corpus file, with an untimed queue drain between passes.
|
||||||
|
"""
|
||||||
|
function component_rows(cfg::FS.Config, clf::FS.Classifier, jobs::Vector{FS.Job}, opts)
|
||||||
|
reps, trials = opts["reps"], opts["trials"]
|
||||||
|
nfiles = length(jobs)
|
||||||
|
paths = [j.path for j in jobs]
|
||||||
|
rows = []
|
||||||
|
|
||||||
|
# A feature vector already in memory, so the inference row measures the net
|
||||||
|
# and not the disk read in front of it.
|
||||||
|
feats = FS.read_features(paths[1])
|
||||||
|
x = reshape(feats, FS.FEATURE_DIM, 1)
|
||||||
|
|
||||||
|
logger = ConsoleLogger(devnull) # components other than the log rows: quiet
|
||||||
|
|
||||||
|
# --- filesize: the stat() read_features does before touching the bytes
|
||||||
|
push!(rows, (; name = "filesize (stat)", part = "classify",
|
||||||
|
ns = best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:reps
|
||||||
|
SINK[] = filesize(paths[(i - 1) % nfiles + 1])
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end))
|
||||||
|
|
||||||
|
# --- read_features: open + head read + seek + tail read + scale
|
||||||
|
push!(rows, (; name = "read_features", part = "classify",
|
||||||
|
ns = best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:reps
|
||||||
|
SINK[] = FS.read_features(paths[(i - 1) % nfiles + 1])
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end))
|
||||||
|
|
||||||
|
# --- Lux.apply: the network alone, features already in memory
|
||||||
|
push!(rows, (; name = "Lux.apply (1x32)", part = "classify",
|
||||||
|
ns = best_of(noop; trials) do
|
||||||
|
for _ in 1:reps
|
||||||
|
SINK[] = Lux.apply(clf.model, x, clf.ps, clf.st)
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end))
|
||||||
|
|
||||||
|
# --- classify: read_features + apply + argmax, what the handler calls
|
||||||
|
push!(rows, (; name = "classify (total)", part = "classify",
|
||||||
|
ns = best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:reps
|
||||||
|
SINK[] = FS.classify(clf, paths[(i - 1) % nfiles + 1])
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end))
|
||||||
|
|
||||||
|
# Stage 1 used to rename each file into known/ or unknown/ before enqueueing
|
||||||
|
# it, and that rename was measured here as its own line item. It is gone:
|
||||||
|
# routing is the enqueue alone, and a file does not move until it is
|
||||||
|
# committed to a terminal sink (src/worker.jl header). So `route` below is
|
||||||
|
# now just the queue handoff.
|
||||||
|
|
||||||
|
# --- enqueue: lock, push, notify on an uncontended, non-full queue
|
||||||
|
q = FS.ChannelQueue(nfiles + 1)
|
||||||
|
stats = FS.StageStats()
|
||||||
|
push!(rows, (; name = "enqueue_blocking!", part = "route",
|
||||||
|
ns = best_of(() -> drain!(q); trials) do
|
||||||
|
@inbounds for job in jobs
|
||||||
|
SINK[] = FS.enqueue_blocking!(q, job, stats;
|
||||||
|
retry_seconds = FS.ROUTE_ENQUEUE_RETRY_SECONDS)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end))
|
||||||
|
drain!(q)
|
||||||
|
|
||||||
|
# --- the two @info lines, under each of the three loggers
|
||||||
|
job1 = jobs[1]
|
||||||
|
logfile = joinpath(cfg.spool_dir, "..", "bench_stage1.log")
|
||||||
|
for (which, label) in ((:null, "logging (NullLogger)"),
|
||||||
|
(:format, "logging (format only)"),
|
||||||
|
(:flush, "logging (flush→file)"))
|
||||||
|
ns = with_logger_named(which, logfile) do
|
||||||
|
best_of(noop; trials) do
|
||||||
|
for _ in 1:reps
|
||||||
|
@info "classified file" worker=1 id=job1.id name=job1.original_name size=job1.size classification=:unknown
|
||||||
|
@info "routed to content triage" worker=1 id=job1.id dest=job1.path
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end
|
||||||
|
end
|
||||||
|
push!(rows, (; name = label, part = "log", ns))
|
||||||
|
end
|
||||||
|
rm(logfile; force = true)
|
||||||
|
|
||||||
|
return rows, logger
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
handler_rows(cfg, jobs, opts) -> Vector
|
||||||
|
|
||||||
|
Time the real `handle_classify_job` end to end under each logger. The difference
|
||||||
|
between the rows is the cost logging adds to a file; the `:flush` row is what the
|
||||||
|
running server actually pays.
|
||||||
|
"""
|
||||||
|
function handler_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
|
||||||
|
trials = opts["trials"]
|
||||||
|
nfiles = length(jobs)
|
||||||
|
logfile = joinpath(dirname(cfg.spool_dir), "bench_stage1.log")
|
||||||
|
known = FS.ChannelQueue(nfiles + 1)
|
||||||
|
unknown = FS.ChannelQueue(nfiles + 1)
|
||||||
|
stats = FS.StageStats()
|
||||||
|
rows = []
|
||||||
|
for (which, label) in ((:null, "handle_classify_job (NullLogger)"),
|
||||||
|
(:format, "handle_classify_job (format only)"),
|
||||||
|
(:flush, "handle_classify_job (flush→file)"),
|
||||||
|
(:debug, "handle_classify_job (JULIA_DEBUG)"))
|
||||||
|
ns = with_logger_named(which, logfile) do
|
||||||
|
best_of(() -> (drain!(known); drain!(unknown)); trials) do
|
||||||
|
@inbounds for job in jobs
|
||||||
|
FS.handle_classify_job(job, cfg, 1, known, unknown, stats)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end
|
||||||
|
end
|
||||||
|
push!(rows, (; name = label, part = "total", ns))
|
||||||
|
end
|
||||||
|
drain!(known); drain!(unknown)
|
||||||
|
rm(logfile; force = true)
|
||||||
|
return rows
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
thread_rows(cfg, jobs, opts) -> Vector
|
||||||
|
|
||||||
|
Run the full handler across worker counts, under the server's real logger. A
|
||||||
|
component that owns a lock (the queue's condition variable, the logger's
|
||||||
|
stream) stops scaling here even though it looked cheap single-threaded, so this
|
||||||
|
is where the single-thread ranking gets checked against the deployed one.
|
||||||
|
"""
|
||||||
|
function thread_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
|
||||||
|
trials = opts["trials"]
|
||||||
|
nfiles = length(jobs)
|
||||||
|
counts = opts["threads"] === nothing ?
|
||||||
|
unique([1; 2; 4; 8; Threads.nthreads()]) :
|
||||||
|
[parse(Int, s) for s in split(String(opts["threads"]), ",")]
|
||||||
|
counts = sort(unique(filter(k -> 1 <= k <= Threads.nthreads(), counts)))
|
||||||
|
|
||||||
|
logfile = joinpath(dirname(cfg.spool_dir), "bench_stage1.log")
|
||||||
|
known = FS.ChannelQueue(nfiles + 1)
|
||||||
|
unknown = FS.ChannelQueue(nfiles + 1)
|
||||||
|
stats = FS.StageStats()
|
||||||
|
rows = []
|
||||||
|
base = 0.0
|
||||||
|
for k in counts
|
||||||
|
ns = with_logger_named(:flush, logfile) do
|
||||||
|
best_of(() -> (drain!(known); drain!(unknown)); trials) do
|
||||||
|
# Static split: each task takes a contiguous slice, so the only
|
||||||
|
# sharing between workers is the state the server also shares:
|
||||||
|
# the classifier, the queues, the logger, the filesystem.
|
||||||
|
chunk = cld(nfiles, k)
|
||||||
|
@sync for t in 1:k
|
||||||
|
lo = (t - 1) * chunk + 1
|
||||||
|
hi = min(t * chunk, nfiles)
|
||||||
|
lo > hi && continue
|
||||||
|
Threads.@spawn begin
|
||||||
|
@inbounds for i in lo:hi
|
||||||
|
FS.handle_classify_job(jobs[i], cfg, t, known, unknown, stats)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end
|
||||||
|
end
|
||||||
|
r = rate(ns) # files/sec aggregate (ns is already per file, wall-clock)
|
||||||
|
k == counts[1] && (base = r)
|
||||||
|
push!(rows, (; workers = k, ns, files_per_sec = r, speedup = r / base))
|
||||||
|
end
|
||||||
|
drain!(known); drain!(unknown)
|
||||||
|
rm(logfile; force = true)
|
||||||
|
return rows
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------- reporting
|
||||||
|
|
||||||
|
function print_components(rows, total_ns)
|
||||||
|
@printf("%-34s %-9s %12s %10s %9s\n", "component", "part", "per file", "rate", "% total")
|
||||||
|
rule()
|
||||||
|
for r in rows
|
||||||
|
@printf("%-34s %-9s %12s %10s %8.1f%%\n", r.name, r.part, human_time(r.ns),
|
||||||
|
human_rate(rate(r.ns)), 100 * r.ns / total_ns)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
function print_threads(rows)
|
||||||
|
@printf("%-9s %12s %12s %9s\n", "workers", "per file", "throughput", "speedup")
|
||||||
|
rule()
|
||||||
|
for r in rows
|
||||||
|
@printf("%-9d %12s %12s %8.2fx\n", r.workers, human_time(r.ns),
|
||||||
|
human_rate(r.files_per_sec), r.speedup)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------------- main
|
||||||
|
|
||||||
|
function main(argv)
|
||||||
|
opts = parse_args(argv)
|
||||||
|
modelpath = String(opts["model"])
|
||||||
|
isfile(modelpath) || (println(stderr, "model artifact not found: $modelpath"); return 1)
|
||||||
|
|
||||||
|
size = parse_size(String(opts["size"]))
|
||||||
|
nfiles, trials = opts["files"], opts["trials"]
|
||||||
|
|
||||||
|
root = opts["dir"] === nothing ?
|
||||||
|
mktempdir(pwd(); prefix = "bench_stage1_") : String(opts["dir"])
|
||||||
|
owned = opts["dir"] === nothing
|
||||||
|
cfg = FS.Config(
|
||||||
|
spool_dir = joinpath(root, "spool"),
|
||||||
|
failed_dir = joinpath(root, "failed"),
|
||||||
|
model_path = modelpath,
|
||||||
|
)
|
||||||
|
for d in (cfg.spool_dir, cfg.failed_dir)
|
||||||
|
mkpath(d)
|
||||||
|
end
|
||||||
|
|
||||||
|
clf = FS.load_classifier(modelpath)
|
||||||
|
FS.CLASSIFIER[] = clf # handle_classify_job reads the global, as in the server
|
||||||
|
|
||||||
|
println("stage-1 component benchmark")
|
||||||
|
rule()
|
||||||
|
@printf("%-22s %s\n", "julia threads", Threads.nthreads())
|
||||||
|
@printf("%-22s %s\n", "model", modelpath)
|
||||||
|
@printf("%-22s %s\n", "corpus", "$(nfiles) files x $(size) B in $(root)")
|
||||||
|
@printf("%-22s %s\n", "reps / trials", "$(opts["reps"]) / $(trials)")
|
||||||
|
|
||||||
|
rng = MersenneTwister(0x5741524d)
|
||||||
|
jobs = make_corpus(cfg, nfiles, size, rng)
|
||||||
|
|
||||||
|
try
|
||||||
|
comps, _ = component_rows(cfg, clf, jobs, opts)
|
||||||
|
handlers = handler_rows(cfg, jobs, opts)
|
||||||
|
|
||||||
|
# The denominator is the handler as the server actually runs it: the real
|
||||||
|
# flushing logger at Info level, one worker. Percentages are shares of
|
||||||
|
# that, so they are directly comparable and the parts can be checked
|
||||||
|
# against the whole. The JULIA_DEBUG row is deliberately *not* the
|
||||||
|
# baseline. It is the opt-in configuration, and letting it set the scale
|
||||||
|
# would make every other component look free.
|
||||||
|
total = only(r.ns for r in handlers if r.name == "handle_classify_job (flush→file)")
|
||||||
|
|
||||||
|
header("Components (single worker)")
|
||||||
|
print_components(comps, total)
|
||||||
|
|
||||||
|
header("Whole handler, by logger")
|
||||||
|
print_components(handlers, total)
|
||||||
|
|
||||||
|
# Stage 1's own log lines are `@debug`, so what the deployed handler pays
|
||||||
|
# for them is the disabled-macro cost, not a formatted line.
|
||||||
|
accounted = sum(r.ns for r in comps if r.name in
|
||||||
|
("classify (total)", "enqueue_blocking!", "logging (NullLogger)"))
|
||||||
|
println()
|
||||||
|
@printf("accounted: %s of %s (%.0f%%); unaccounted overhead %s\n",
|
||||||
|
human_time(accounted), human_time(total), 100 * accounted / total,
|
||||||
|
human_time(max(total - accounted, 0)))
|
||||||
|
|
||||||
|
threads = nothing
|
||||||
|
if !opts["no-threads"] && Threads.nthreads() > 1
|
||||||
|
threads = thread_rows(cfg, jobs, opts)
|
||||||
|
header("Full handler across workers (server logger)")
|
||||||
|
print_threads(threads)
|
||||||
|
end
|
||||||
|
|
||||||
|
if opts["json"] !== nothing
|
||||||
|
open(String(opts["json"]), "w") do io
|
||||||
|
JSON3.write(io, (;
|
||||||
|
julia_threads = Threads.nthreads(),
|
||||||
|
file_size = size, files = nfiles, reps = opts["reps"], trials,
|
||||||
|
components = comps, handlers, threads,
|
||||||
|
))
|
||||||
|
end
|
||||||
|
println("\nwrote ", opts["json"])
|
||||||
|
end
|
||||||
|
finally
|
||||||
|
owned && rm(root; recursive = true, force = true)
|
||||||
|
end
|
||||||
|
return 0
|
||||||
|
end
|
||||||
|
|
||||||
|
exit(main(ARGS))
|
||||||
700
bin/bench_stage2.jl
Normal file
700
bin/bench_stage2.jl
Normal file
@@ -0,0 +1,700 @@
|
|||||||
|
#!/usr/bin/env julia
|
||||||
|
#
|
||||||
|
# bench_stage2.jl: take stage 2 apart and find the slowest component.
|
||||||
|
#
|
||||||
|
# bin/bench.jl reports stage 2 as a single number (its throughput and worker
|
||||||
|
# utilization under whole-pipeline contention). It doesn't say *which part* of
|
||||||
|
# the stage costs the most, and stage 2 is the one stage whose cost is dominated
|
||||||
|
# by something outside Julia entirely: it forks `exiftool`, a Perl program, once
|
||||||
|
# per file. Per file the stage does:
|
||||||
|
#
|
||||||
|
# build_metadata
|
||||||
|
# run_exiftool fork/exec exiftool -json -G -n, capture stdout
|
||||||
|
# run_with_timeout the watchdog wrapper around the subprocess
|
||||||
|
# JSON3.read parse the dump
|
||||||
|
# normalize_metadata coalesce ~20 tag names into the sidecar schema
|
||||||
|
# finalize_known! (commit_enriched!)
|
||||||
|
# JSON3.write serialize the sidecar payload
|
||||||
|
# write + fsync durably persist the sidecar bytes to a temp name
|
||||||
|
# mv + fsync_dir commit the sidecar, then persist the rename itself
|
||||||
|
# move_to rename spool/<f> -> done/<f>, the commit point
|
||||||
|
# logging one @info line ("enriched")
|
||||||
|
#
|
||||||
|
# This script times each of those in isolation, then times the real
|
||||||
|
# `handle_known_job` end to end so the parts can be checked against the whole.
|
||||||
|
#
|
||||||
|
# Three things here that the stage-1 benchmark has no equivalent of:
|
||||||
|
#
|
||||||
|
# * The corpus must be real files. exiftool's cost depends on what it finds;
|
||||||
|
# random bytes exit early and would understate the stage by a lot. The
|
||||||
|
# default corpus is `data/done`, files that already went through stage 2 on
|
||||||
|
# this machine, copied back into a scratch spool/ dir.
|
||||||
|
# * Two rows price the *alternatives* to one-fork-per-file, because if the
|
||||||
|
# fork dominates then the only fixes are to stop paying it per file:
|
||||||
|
# `exiftool (batched Nx)` runs the whole corpus through one process, and
|
||||||
|
# `exiftool (-stay_open)` keeps a single process alive and feeds it one file
|
||||||
|
# at a time over a pipe: the shape a streaming pipeline could actually use.
|
||||||
|
# Both are measured, not assumed.
|
||||||
|
# * `run_with_timeout` gets its own row *next to* a bare `Base.run` of the same
|
||||||
|
# command. The difference is what the watchdog costs, and its polling loop
|
||||||
|
# (`sleep(0.1)`) is a suspicious enough design to want measured rather than
|
||||||
|
# reasoned about.
|
||||||
|
#
|
||||||
|
# The `--threads` sweep runs the full handler across worker counts: subprocess
|
||||||
|
# spawning contends on things (the kernel's fork path, page cache, the logger's
|
||||||
|
# stream) that a single-threaded ranking can't reveal.
|
||||||
|
#
|
||||||
|
# Usage:
|
||||||
|
# julia --project=. -t auto bin/bench_stage2.jl [options]
|
||||||
|
#
|
||||||
|
# --files N corpus files per timed pass (default: 48). The concurrency
|
||||||
|
# sweep wants more than the component rows do. With a single
|
||||||
|
# 2 s file in the corpus, 48 files can't show more than ~4x no
|
||||||
|
# matter how many workers run, so pass --files 150 when the
|
||||||
|
# question is scaling.
|
||||||
|
# --reps N calls per timed pass for cheap, non-consuming benchmarks (default: 2000)
|
||||||
|
# --trials N timed passes; the minimum is reported (default: 3)
|
||||||
|
# --corpus PATH directory of real files to draw the corpus from (default: data/done)
|
||||||
|
# --dir PATH working directory for the corpus (default: a temp dir under data/)
|
||||||
|
# --timeout SEC exiftool timeout, as Config.exiftool_timeout (default: 30)
|
||||||
|
# --threads LIST worker counts for the concurrency sweep (default: 1,2,4,8,nthreads)
|
||||||
|
# --no-threads skip the concurrency sweep
|
||||||
|
# --no-stay-open skip the persistent-exiftool probe
|
||||||
|
# --json PATH also write the results as JSON
|
||||||
|
#
|
||||||
|
# Reported times are the *minimum* over trials: the floor is the signal and
|
||||||
|
# everything above it is scheduler, page-cache and GC noise.
|
||||||
|
|
||||||
|
using DarkStruct
|
||||||
|
using JSON3
|
||||||
|
using Logging
|
||||||
|
using Printf
|
||||||
|
using Random
|
||||||
|
|
||||||
|
const FS = DarkStruct
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------- option parsing
|
||||||
|
|
||||||
|
const DEFAULTS = Dict{String,Any}(
|
||||||
|
"files" => 48,
|
||||||
|
"reps" => 2000,
|
||||||
|
"trials" => 3,
|
||||||
|
"corpus" => "data/done",
|
||||||
|
"dir" => nothing,
|
||||||
|
"timeout" => 30,
|
||||||
|
"threads" => nothing,
|
||||||
|
"no-threads" => false,
|
||||||
|
"no-stay-open" => false,
|
||||||
|
"json" => nothing,
|
||||||
|
)
|
||||||
|
|
||||||
|
const FLAGS = ("no-threads", "no-stay-open")
|
||||||
|
const INTS = ("files", "reps", "trials", "timeout")
|
||||||
|
|
||||||
|
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 INTS ? 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)
|
||||||
|
|
||||||
|
"""
|
||||||
|
best_of(pass, prepare; trials) -> ns_per_op
|
||||||
|
|
||||||
|
Run `pass()` `trials` times and report the fastest, in nanoseconds per operation
|
||||||
|
(`pass` returns the number of operations it performed). `prepare()` runs before
|
||||||
|
each pass and is *not* timed: that is where a consuming benchmark puts the file
|
||||||
|
back where it started. `pass` comes first so callers can pass it as a `do` block.
|
||||||
|
|
||||||
|
The first pass is thrown away: it pays Julia's JIT compilation, which on calls
|
||||||
|
this small is orders of magnitude more than the thing being measured.
|
||||||
|
"""
|
||||||
|
function best_of(pass, prepare; trials::Int)
|
||||||
|
best = Inf
|
||||||
|
for t in 0:trials
|
||||||
|
prepare()
|
||||||
|
GC.gc()
|
||||||
|
t0 = time_ns()
|
||||||
|
n = pass()
|
||||||
|
dt = Float64(time_ns() - t0)
|
||||||
|
t == 0 && continue # warm-up: compiled, not measured
|
||||||
|
best = min(best, dt / n)
|
||||||
|
end
|
||||||
|
return best
|
||||||
|
end
|
||||||
|
|
||||||
|
noop() = nothing
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------- 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
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
rate(ns::Real) = 1e9 / max(ns, 1e-9)
|
||||||
|
|
||||||
|
rule(n = 84) = println("-" ^ n)
|
||||||
|
|
||||||
|
function header(title)
|
||||||
|
println()
|
||||||
|
println(title)
|
||||||
|
rule()
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------------- corpus
|
||||||
|
|
||||||
|
"""
|
||||||
|
make_corpus(cfg, corpus_dir, n) -> Vector{Job}
|
||||||
|
|
||||||
|
Copy up to `n` real files from `corpus_dir` into `spool/` and build the `Job`
|
||||||
|
references a stage-2 worker would dequeue for them: the exact input
|
||||||
|
`handle_known_job` sees.
|
||||||
|
|
||||||
|
Real files, not generated ones: exiftool's cost is a function of what it can
|
||||||
|
parse, and a file of random bytes bails out early enough to understate the stage
|
||||||
|
by an order of magnitude. `.meta.json` sidecars are skipped, since they are stage-2
|
||||||
|
*output*, and enriching them would measure the wrong population.
|
||||||
|
"""
|
||||||
|
function make_corpus(cfg::FS.Config, corpus_dir::AbstractString, n::Int)
|
||||||
|
isdir(corpus_dir) || error("corpus dir not found: $corpus_dir")
|
||||||
|
names = filter(readdir(corpus_dir)) do f
|
||||||
|
!endswith(f, ".meta.json") && isfile(joinpath(corpus_dir, f))
|
||||||
|
end
|
||||||
|
isempty(names) && error("no usable files in corpus dir: $corpus_dir")
|
||||||
|
sort!(names) # deterministic selection across runs
|
||||||
|
length(names) > n && (names = names[1:n])
|
||||||
|
|
||||||
|
jobs = FS.Job[]
|
||||||
|
for (i, name) in enumerate(names)
|
||||||
|
src = joinpath(corpus_dir, name)
|
||||||
|
# Give it a fresh id/spool-style filename so nothing collides with the
|
||||||
|
# corpus the file came from.
|
||||||
|
id, spooled = FS.spool_path(cfg, @sprintf("s2-%04d-%s", i, basename(name)))
|
||||||
|
cp(src, spooled; force = true)
|
||||||
|
# Staged in spool/, not a known/ dir: stage 1 routes by enqueueing and
|
||||||
|
# leaves the bytes where intake put them, so this is the exact on-disk
|
||||||
|
# state a stage-2 worker dequeues into (src/worker.jl header).
|
||||||
|
push!(jobs, FS.Job(id, basename(name), spooled, filesize(spooled), time()))
|
||||||
|
end
|
||||||
|
return jobs
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
respool!(cfg, jobs)
|
||||||
|
|
||||||
|
Put every corpus file back in `spool/`, wherever the last pass left it (done/ or
|
||||||
|
already home), and delete any sidecar it produced. This is the untimed `prepare`
|
||||||
|
step for benchmarks that consume their input by committing it. Stage 2 still
|
||||||
|
moves, because its move is the terminal commit, not an inter-stage hop.
|
||||||
|
"""
|
||||||
|
function respool!(cfg::FS.Config, jobs::Vector{FS.Job})
|
||||||
|
for job in jobs
|
||||||
|
base = basename(job.path)
|
||||||
|
for dir in (cfg.done_dir, cfg.failed_dir)
|
||||||
|
sidecar = joinpath(dir, string(base, ".meta.json"))
|
||||||
|
rm(sidecar; force = true)
|
||||||
|
rm(string(sidecar, ".tmp"); force = true)
|
||||||
|
end
|
||||||
|
isfile(job.path) && continue
|
||||||
|
for dir in (cfg.done_dir, cfg.failed_dir)
|
||||||
|
candidate = joinpath(dir, base)
|
||||||
|
if isfile(candidate)
|
||||||
|
mv(candidate, job.path; force = true)
|
||||||
|
break
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------- loggers
|
||||||
|
|
||||||
|
"""
|
||||||
|
with_logger_named(f, which, path)
|
||||||
|
|
||||||
|
Run `f` under one of the loggers the cost of logging is bracketed by:
|
||||||
|
|
||||||
|
* `:null` `NullLogger`: the `@info` macro's own overhead, nothing else.
|
||||||
|
* `:format` `ConsoleLogger` to `devnull`: message formatting and key/value
|
||||||
|
interpolation, but no I/O.
|
||||||
|
* `:flush` `FlushLogger(ConsoleLogger(io))` to a real file: what
|
||||||
|
`DarkStruct.run` installs, under the redirect it was written
|
||||||
|
for. Stage 2's per-file line is `@info`, not `@debug`, so this
|
||||||
|
row is what the deployed server actually pays.
|
||||||
|
"""
|
||||||
|
function with_logger_named(f, which::Symbol, path::AbstractString)
|
||||||
|
if which === :null
|
||||||
|
return with_logger(f, NullLogger())
|
||||||
|
elseif which === :format
|
||||||
|
return with_logger(f, ConsoleLogger(devnull))
|
||||||
|
elseif which === :flush
|
||||||
|
return open(path, "w") do io
|
||||||
|
with_logger(f, FS.FlushLogger(ConsoleLogger(io, Logging.Info)))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
error("unknown logger: $which")
|
||||||
|
end
|
||||||
|
|
||||||
|
# --------------------------------------------------- exiftool spawn alternatives
|
||||||
|
|
||||||
|
"""
|
||||||
|
capture(cmd) -> Vector{UInt8}
|
||||||
|
|
||||||
|
Run `cmd` and return its stdout, tolerating a non-zero exit the way
|
||||||
|
`run_with_timeout` does. `read(cmd, String)` would throw instead, and a real
|
||||||
|
corpus makes that a question of when, not whether: exiftool exits 1 on a file
|
||||||
|
whose type it can't recognize, which in this pipeline is a routine outcome (it
|
||||||
|
yields a degraded sidecar, not a failure). This is `run_with_timeout` minus the
|
||||||
|
watchdog, so the gap between the two rows prices the watchdog exactly.
|
||||||
|
"""
|
||||||
|
function capture(cmd::Cmd)
|
||||||
|
out = IOBuffer()
|
||||||
|
proc = Base.run(pipeline(cmd; stdout = out, stderr = devnull); wait = false)
|
||||||
|
wait(proc)
|
||||||
|
return take!(out)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
batched_ns(paths, trials) -> ns_per_file
|
||||||
|
|
||||||
|
Run the whole corpus through *one* `exiftool` process and divide by the file
|
||||||
|
count. This is the floor for "what does exiftool cost if you stop paying the
|
||||||
|
interpreter startup per file": the fork, the Perl boot and the module loads are
|
||||||
|
paid once for the batch instead of once per file.
|
||||||
|
"""
|
||||||
|
function batched_ns(paths::Vector{String}, trials::Int)
|
||||||
|
return best_of(noop; trials) do
|
||||||
|
SINK[] = capture(`exiftool -json -G -n $paths`)
|
||||||
|
length(paths)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
stay_open_ns(paths, trials) -> ns_per_file (or nothing if unsupported)
|
||||||
|
|
||||||
|
Feed files one at a time to a single long-lived `exiftool -stay_open True -@ -`
|
||||||
|
process over a pipe, reading its `{ready}` sentinel after each. Unlike the
|
||||||
|
batched row this preserves the pipeline's actual shape (one file in, one result
|
||||||
|
out, arriving whenever it arrives) so it prices the realistic fix rather than
|
||||||
|
an unrealistic one.
|
||||||
|
"""
|
||||||
|
function stay_open_ns(paths::Vector{String}, trials::Int)
|
||||||
|
inp, outp = Pipe(), Pipe()
|
||||||
|
proc = Base.run(pipeline(`exiftool -stay_open True -@ -`;
|
||||||
|
stdin = inp, stdout = outp, stderr = devnull); wait = false)
|
||||||
|
close(inp.out); close(outp.in)
|
||||||
|
|
||||||
|
ask(path) = begin
|
||||||
|
write(inp, "-json\n-G\n-n\n", path, "\n-execute\n")
|
||||||
|
flush(inp)
|
||||||
|
readuntil(outp, "{ready}")
|
||||||
|
end
|
||||||
|
try
|
||||||
|
ask(paths[1]) # pay the one-time process startup untimed
|
||||||
|
return best_of(noop; trials) do
|
||||||
|
for p in paths
|
||||||
|
SINK[] = ask(p)
|
||||||
|
end
|
||||||
|
length(paths)
|
||||||
|
end
|
||||||
|
finally
|
||||||
|
try
|
||||||
|
write(inp, "-stay_open\nFalse\n"); flush(inp); close(inp)
|
||||||
|
wait(proc)
|
||||||
|
catch
|
||||||
|
kill(proc, Base.SIGKILL)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------ components
|
||||||
|
|
||||||
|
"""
|
||||||
|
component_rows(cfg, jobs, opts) -> Vector
|
||||||
|
|
||||||
|
Time each piece of stage 2 on its own. The subprocess rows run once per corpus
|
||||||
|
file (they cost milliseconds and don't need repetition); the in-memory and
|
||||||
|
filesystem rows run `reps` times; the committing rows run once per corpus file
|
||||||
|
with an untimed reset between passes.
|
||||||
|
"""
|
||||||
|
function component_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
|
||||||
|
reps, trials = opts["reps"], opts["trials"]
|
||||||
|
timeout = opts["timeout"]
|
||||||
|
nfiles = length(jobs)
|
||||||
|
paths = String[j.path for j in jobs]
|
||||||
|
rows = []
|
||||||
|
add!(name, part, ns) = push!(rows, (; name, part, ns))
|
||||||
|
|
||||||
|
# --- the bare interpreter: fork/exec + Perl boot, reading no file at all.
|
||||||
|
# Everything the real call does beyond this is actual work.
|
||||||
|
add!("exiftool -ver (spawn)", "extract", best_of(noop; trials) do
|
||||||
|
for _ in 1:nfiles
|
||||||
|
SINK[] = capture(`exiftool -ver`)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- the real command, run directly: spawn + parse the file, no watchdog.
|
||||||
|
add!("exiftool -json (raw run)", "extract", best_of(noop; trials) do
|
||||||
|
@inbounds for p in paths
|
||||||
|
SINK[] = capture(`exiftool -json -G -n $p`)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- the same command through the watchdog wrapper the stage actually uses.
|
||||||
|
# The gap to the row above is what the timeout costs.
|
||||||
|
add!("run_with_timeout", "extract", best_of(noop; trials) do
|
||||||
|
@inbounds for p in paths
|
||||||
|
SINK[] = FS.run_with_timeout(`exiftool -json -G -n $p`, timeout)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- run_exiftool: the wrapper plus JSON3.read plus the group-stripped Dict.
|
||||||
|
add!("run_exiftool (total)", "extract", best_of(noop; trials) do
|
||||||
|
@inbounds for p in paths
|
||||||
|
SINK[] = FS.run_exiftool(p, timeout)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- what a fork-free exiftool would cost, two ways (see the docstrings).
|
||||||
|
add!("exiftool (batched $(nfiles)x)", "alt", batched_ns(paths, trials))
|
||||||
|
if !opts["no-stay-open"]
|
||||||
|
try
|
||||||
|
add!("exiftool (-stay_open)", "alt", stay_open_ns(paths, trials))
|
||||||
|
catch e
|
||||||
|
@warn "persistent-exiftool probe failed; skipping" exception = e
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- parsing alone, from bytes already captured: isolates JSON3 from the fork.
|
||||||
|
raw = [FS.run_with_timeout(`exiftool -json -G -n $p`, timeout) for p in paths]
|
||||||
|
valid = [b for b in raw if b !== nothing]
|
||||||
|
if !isempty(valid)
|
||||||
|
add!("JSON3.read (parse dump)", "extract", best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:reps
|
||||||
|
SINK[] = JSON3.read(String(copy(valid[(i - 1) % length(valid) + 1])))
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end)
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- normalize_metadata: the ~20 tag coalesces, on a tag map already in memory.
|
||||||
|
bytags = [FS.run_exiftool(p, timeout) for p in paths]
|
||||||
|
good = [(j, b) for (j, b) in zip(jobs, bytags) if b !== nothing]
|
||||||
|
isempty(good) && error("exiftool produced no parseable output for any corpus file")
|
||||||
|
add!("normalize_metadata", "extract", best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:reps
|
||||||
|
j, b = good[(i - 1) % length(good) + 1]
|
||||||
|
SINK[] = FS.normalize_metadata(j, b)
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end)
|
||||||
|
|
||||||
|
# The sidecar payloads, built once, untimed: the commit rows below measure
|
||||||
|
# committing, not extracting.
|
||||||
|
metas = [FS.normalize_metadata(j, b) for (j, b) in good]
|
||||||
|
|
||||||
|
# --- serializing the sidecar (the raw dump makes this bigger than it looks).
|
||||||
|
add!("JSON3.write (sidecar)", "commit", best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:reps
|
||||||
|
SINK[] = JSON3.write(metas[(i - 1) % length(metas) + 1])
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- sidecar bytes: open + write + flush + fsync, to a temp name.
|
||||||
|
# Non-consuming: same path rewritten each rep, as commit_enriched! does.
|
||||||
|
tmp = joinpath(cfg.done_dir, "bench_stage2_sidecar.tmp")
|
||||||
|
blobs = [JSON3.write(m) for m in metas]
|
||||||
|
# One pass over the real sidecar population, not `reps` of them. Two reasons,
|
||||||
|
# and the first is a correctness trap: thousands of back-to-back fsyncs
|
||||||
|
# saturate the device's write cache and each one starts waiting on the
|
||||||
|
# queue, which reported this row at 16 ms/file: eight times the whole
|
||||||
|
# `commit_enriched!` that contains it. The real stage fsyncs once per file
|
||||||
|
# with ~160 ms of exiftool between, and never queues that way. Second, real
|
||||||
|
# sidecars vary hugely in size (a zip's raw dump dwarfs a jpeg's), so the
|
||||||
|
# honest per-file number is one pass over all of them, not a cycle.
|
||||||
|
nio = length(blobs)
|
||||||
|
add!("write + fsync (sidecar)", "commit", best_of(noop; trials) do
|
||||||
|
@inbounds for i in 1:nio
|
||||||
|
open(tmp, "w") do io
|
||||||
|
write(io, blobs[(i - 1) % length(blobs) + 1])
|
||||||
|
flush(io)
|
||||||
|
FS.fsync_fd(fd(io))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
nio
|
||||||
|
end)
|
||||||
|
rm(tmp; force = true)
|
||||||
|
|
||||||
|
# --- fsync_dir: persisting the rename itself, once per file in the real path.
|
||||||
|
add!("fsync_dir (done/)", "commit", best_of(noop; trials) do
|
||||||
|
for _ in 1:nio
|
||||||
|
SINK[] = FS.fsync_dir(cfg.done_dir)
|
||||||
|
end
|
||||||
|
nio
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- move_to: the rename spool/<f> -> done/<f>. Consuming: reset each pass.
|
||||||
|
add!("move_to (rename)", "commit", best_of(() -> respool!(cfg, jobs); trials) do
|
||||||
|
@inbounds for job in jobs
|
||||||
|
SINK[] = FS.move_to(cfg.done_dir, job)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end)
|
||||||
|
|
||||||
|
# --- commit_enriched!: the whole sidecar-first commit, extraction excluded.
|
||||||
|
committable = [j for (j, _) in good]
|
||||||
|
add!("commit_enriched! (total)", "commit",
|
||||||
|
best_of(() -> respool!(cfg, jobs); trials) do
|
||||||
|
@inbounds for (k, job) in enumerate(committable)
|
||||||
|
SINK[] = FS.commit_enriched!(cfg.done_dir, job, metas[k])
|
||||||
|
end
|
||||||
|
length(committable)
|
||||||
|
end)
|
||||||
|
respool!(cfg, jobs)
|
||||||
|
|
||||||
|
# --- the one @info line, under each logger.
|
||||||
|
job1, meta1 = good[1][1], metas[1]
|
||||||
|
logfile = joinpath(dirname(cfg.spool_dir), "bench_stage2.log")
|
||||||
|
for (which, label) in ((:null, "logging (NullLogger)"),
|
||||||
|
(:format, "logging (format only)"),
|
||||||
|
(:flush, "logging (flush→file)"))
|
||||||
|
ns = with_logger_named(which, logfile) do
|
||||||
|
best_of(noop; trials) do
|
||||||
|
for _ in 1:reps
|
||||||
|
@info "enriched" worker=1 id=job1.id dest=job1.path sidecar="x.meta.json" file_type=meta1.file_type created_by=meta1.created_by degraded=(meta1.error !== nothing)
|
||||||
|
end
|
||||||
|
reps
|
||||||
|
end
|
||||||
|
end
|
||||||
|
add!(label, "log", ns)
|
||||||
|
end
|
||||||
|
rm(logfile; force = true)
|
||||||
|
|
||||||
|
return rows
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
handler_rows(cfg, jobs, opts) -> Vector
|
||||||
|
|
||||||
|
Time the real `handle_known_job` end to end under each logger. The difference
|
||||||
|
between the rows is the cost logging adds to a file; the `:flush` row is what the
|
||||||
|
running server actually pays.
|
||||||
|
"""
|
||||||
|
function handler_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
|
||||||
|
trials = opts["trials"]
|
||||||
|
nfiles = length(jobs)
|
||||||
|
logfile = joinpath(dirname(cfg.spool_dir), "bench_stage2.log")
|
||||||
|
rows = []
|
||||||
|
for (which, label) in ((:null, "handle_known_job (NullLogger)"),
|
||||||
|
(:format, "handle_known_job (format only)"),
|
||||||
|
(:flush, "handle_known_job (flush→file)"))
|
||||||
|
ns = with_logger_named(which, logfile) do
|
||||||
|
best_of(() -> respool!(cfg, jobs); trials) do
|
||||||
|
@inbounds for job in jobs
|
||||||
|
FS.handle_known_job(job, cfg, 1)
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end
|
||||||
|
end
|
||||||
|
push!(rows, (; name = label, part = "total", ns))
|
||||||
|
end
|
||||||
|
respool!(cfg, jobs)
|
||||||
|
rm(logfile; force = true)
|
||||||
|
return rows
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
thread_rows(cfg, jobs, opts) -> Vector
|
||||||
|
|
||||||
|
Run the full handler across worker counts, under the server's real logger. Stage
|
||||||
|
2 spends most of its time in a child process, so this is the sweep that matters
|
||||||
|
most: whether the stage scales is a question about the kernel's fork path and the
|
||||||
|
machine's cores, not about Julia.
|
||||||
|
|
||||||
|
Workers pull from a shared atomic counter rather than taking a contiguous slice.
|
||||||
|
That matches the server (its pool pulls from one queue), and it matters here in a
|
||||||
|
way it doesn't for stage 1: per-file exiftool time spans two orders of magnitude
|
||||||
|
on a real corpus (a single 2 s archive among 48 files), so a static split leaves
|
||||||
|
whichever worker drew it running alone while the rest idle, and the sweep would
|
||||||
|
report a scaling ceiling that is really just load imbalance.
|
||||||
|
"""
|
||||||
|
function thread_rows(cfg::FS.Config, jobs::Vector{FS.Job}, opts)
|
||||||
|
trials = opts["trials"]
|
||||||
|
nfiles = length(jobs)
|
||||||
|
counts = opts["threads"] === nothing ?
|
||||||
|
unique([1; 2; 4; 8; Threads.nthreads()]) :
|
||||||
|
[parse(Int, s) for s in split(String(opts["threads"]), ",")]
|
||||||
|
counts = sort(unique(filter(k -> 1 <= k <= Threads.nthreads(), counts)))
|
||||||
|
|
||||||
|
logfile = joinpath(dirname(cfg.spool_dir), "bench_stage2.log")
|
||||||
|
rows = []
|
||||||
|
base = 0.0
|
||||||
|
for k in counts
|
||||||
|
ns = with_logger_named(:flush, logfile) do
|
||||||
|
next = Threads.Atomic{Int}(1)
|
||||||
|
best_of(() -> (respool!(cfg, jobs); next[] = 1); trials) do
|
||||||
|
# Shared counter, not a contiguous slice: every worker takes the
|
||||||
|
# next unclaimed file the moment it frees up, exactly as the
|
||||||
|
# server's pool takes the next job off the known queue.
|
||||||
|
@sync for t in 1:k
|
||||||
|
Threads.@spawn begin
|
||||||
|
while true
|
||||||
|
i = Threads.atomic_add!(next, 1)
|
||||||
|
i > nfiles && break
|
||||||
|
@inbounds FS.handle_known_job(jobs[i], cfg, t)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
nfiles
|
||||||
|
end
|
||||||
|
end
|
||||||
|
r = rate(ns) # files/sec aggregate (ns is already per file, wall-clock)
|
||||||
|
k == counts[1] && (base = r)
|
||||||
|
push!(rows, (; workers = k, ns, files_per_sec = r, speedup = r / base))
|
||||||
|
end
|
||||||
|
respool!(cfg, jobs)
|
||||||
|
rm(logfile; force = true)
|
||||||
|
return rows
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------- reporting
|
||||||
|
|
||||||
|
function print_components(rows, total_ns)
|
||||||
|
@printf("%-34s %-9s %12s %10s %9s\n", "component", "part", "per file", "rate", "% total")
|
||||||
|
rule()
|
||||||
|
for r in rows
|
||||||
|
@printf("%-34s %-9s %12s %10s %8.1f%%\n", r.name, r.part, human_time(r.ns),
|
||||||
|
human_rate(rate(r.ns)), 100 * r.ns / total_ns)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
function print_threads(rows)
|
||||||
|
@printf("%-9s %12s %12s %9s\n", "workers", "per file", "throughput", "speedup")
|
||||||
|
rule()
|
||||||
|
for r in rows
|
||||||
|
@printf("%-9d %12s %12s %8.2fx\n", r.workers, human_time(r.ns),
|
||||||
|
human_rate(r.files_per_sec), r.speedup)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------------- main
|
||||||
|
|
||||||
|
function main(argv)
|
||||||
|
opts = parse_args(argv)
|
||||||
|
try
|
||||||
|
FS.assert_exiftool()
|
||||||
|
catch e
|
||||||
|
println(stderr, sprint(showerror, e)); return 1
|
||||||
|
end
|
||||||
|
|
||||||
|
root = opts["dir"] === nothing ?
|
||||||
|
mktempdir(pwd(); prefix = "bench_stage2_") : String(opts["dir"])
|
||||||
|
owned = opts["dir"] === nothing
|
||||||
|
cfg = FS.Config(
|
||||||
|
spool_dir = joinpath(root, "spool"),
|
||||||
|
done_dir = joinpath(root, "done"),
|
||||||
|
failed_dir = joinpath(root, "failed"),
|
||||||
|
exiftool_timeout = opts["timeout"],
|
||||||
|
)
|
||||||
|
for d in (cfg.spool_dir, cfg.done_dir, cfg.failed_dir)
|
||||||
|
mkpath(d)
|
||||||
|
end
|
||||||
|
|
||||||
|
jobs = try
|
||||||
|
make_corpus(cfg, String(opts["corpus"]), opts["files"])
|
||||||
|
catch e
|
||||||
|
owned && rm(root; recursive = true, force = true)
|
||||||
|
println(stderr, sprint(showerror, e)); return 1
|
||||||
|
end
|
||||||
|
bytes = sum(j.size for j in jobs)
|
||||||
|
|
||||||
|
println("stage-2 component benchmark")
|
||||||
|
rule()
|
||||||
|
@printf("%-22s %s\n", "julia threads", Threads.nthreads())
|
||||||
|
@printf("%-22s %s\n", "exiftool", strip(read(`exiftool -ver`, String)))
|
||||||
|
@printf("%-22s %s\n", "corpus", "$(length(jobs)) files ($(round(bytes / 1024^2; digits=1)) MiB) from $(opts["corpus"])")
|
||||||
|
@printf("%-22s %s\n", "scratch", root)
|
||||||
|
@printf("%-22s %s\n", "reps / trials", "$(opts["reps"]) / $(opts["trials"])")
|
||||||
|
@printf("%-22s %s\n", "exiftool timeout", "$(opts["timeout"]) s")
|
||||||
|
|
||||||
|
try
|
||||||
|
comps = component_rows(cfg, jobs, opts)
|
||||||
|
handlers = handler_rows(cfg, jobs, opts)
|
||||||
|
|
||||||
|
# The denominator is the handler as the server actually runs it: the real
|
||||||
|
# flushing logger, one worker. Percentages are shares of that, so they are
|
||||||
|
# directly comparable and the parts can be checked against the whole.
|
||||||
|
total = only(r.ns for r in handlers if r.name == "handle_known_job (flush→file)")
|
||||||
|
|
||||||
|
header("Components (single worker)")
|
||||||
|
print_components(comps, total)
|
||||||
|
|
||||||
|
header("Whole handler, by logger")
|
||||||
|
print_components(handlers, total)
|
||||||
|
|
||||||
|
pick(name) = only(r.ns for r in comps if r.name == name)
|
||||||
|
accounted = pick("run_exiftool (total)") + pick("commit_enriched! (total)") +
|
||||||
|
pick("logging (flush→file)")
|
||||||
|
println()
|
||||||
|
@printf("accounted: %s of %s (%.0f%%); unaccounted overhead %s\n",
|
||||||
|
human_time(accounted), human_time(total), 100 * accounted / total,
|
||||||
|
human_time(max(total - accounted, 0)))
|
||||||
|
|
||||||
|
threads = nothing
|
||||||
|
if !opts["no-threads"] && Threads.nthreads() > 1
|
||||||
|
threads = thread_rows(cfg, jobs, opts)
|
||||||
|
header("Full handler across workers (server logger)")
|
||||||
|
print_threads(threads)
|
||||||
|
end
|
||||||
|
|
||||||
|
if opts["json"] !== nothing
|
||||||
|
open(String(opts["json"]), "w") do io
|
||||||
|
JSON3.write(io, (;
|
||||||
|
julia_threads = Threads.nthreads(),
|
||||||
|
files = length(jobs), corpus_bytes = bytes,
|
||||||
|
reps = opts["reps"], trials = opts["trials"],
|
||||||
|
exiftool_timeout = opts["timeout"],
|
||||||
|
components = comps, handlers, threads,
|
||||||
|
))
|
||||||
|
end
|
||||||
|
println("\nwrote ", opts["json"])
|
||||||
|
end
|
||||||
|
finally
|
||||||
|
owned && rm(root; recursive = true, force = true)
|
||||||
|
end
|
||||||
|
return 0
|
||||||
|
end
|
||||||
|
|
||||||
|
exit(main(ARGS))
|
||||||
217
bin/cluster_calibrate.jl
Normal file
217
bin/cluster_calibrate.jl
Normal file
@@ -0,0 +1,217 @@
|
|||||||
|
#!/usr/bin/env julia
|
||||||
|
#
|
||||||
|
# Phase-A calibration for stage-5 header clustering (model/DESIGN_clustering.md
|
||||||
|
# §7). Runs labeled known files through the exact clustering pipeline, scores the
|
||||||
|
# recovered partition against magic-collapsed ground truth with ARI / V-measure,
|
||||||
|
# grid-tunes (α, β, bg_mass, n), and cross-checks the winning config against a
|
||||||
|
# model-free NCD (gzip) baseline (§8). The settings printed here are the ones the
|
||||||
|
# machine rediscovers known formats at; copy the winner into config.jl.
|
||||||
|
#
|
||||||
|
# julia --project=. bin/cluster_calibrate.jl [training_set_dir]
|
||||||
|
#
|
||||||
|
# Defaults to ../training_set. Prints a report; writes nothing.
|
||||||
|
|
||||||
|
using Random
|
||||||
|
using Printf
|
||||||
|
|
||||||
|
include(joinpath(@__DIR__, "..", "src", "cluster.jl"))
|
||||||
|
|
||||||
|
# --- ground truth: magic-collapsed classes, NOT extensions (DESIGN §7.2) -----
|
||||||
|
|
||||||
|
"""
|
||||||
|
truth_label(path) -> String
|
||||||
|
|
||||||
|
The magic-collapsed format class of a file, read from its actual bytes (so
|
||||||
|
docx≡zip and the whole ELF family merge, exactly the answer we want the
|
||||||
|
clustering to reproduce). `tar` is detected by the `ustar` magic at offset 257,
|
||||||
|
outside the model's front window, so tars are the accepted blind spot that
|
||||||
|
scatters to background.
|
||||||
|
"""
|
||||||
|
function truth_label(path::AbstractString)
|
||||||
|
b = zeros(UInt8, 262)
|
||||||
|
open(path) do io
|
||||||
|
chunk = read(io, 262)
|
||||||
|
copyto!(b, 1, chunk, 1, length(chunk))
|
||||||
|
end
|
||||||
|
b[1] == 0x1f && b[2] == 0x8b && return "gzip"
|
||||||
|
b[1] == 0x50 && b[2] == 0x4b && return "pkzip"
|
||||||
|
b[1] == 0x25 && b[2] == 0x50 && b[3] == 0x44 && b[4] == 0x46 && return "pdf"
|
||||||
|
b[1] == 0xff && b[2] == 0xd8 && b[3] == 0xff && return "jpeg"
|
||||||
|
b[1] == 0x7f && b[2] == 0x45 && b[3] == 0x4c && b[4] == 0x46 && return "elf"
|
||||||
|
(b[258] == 0x75 && b[259] == 0x73 && b[260] == 0x74 && b[261] == 0x61 && b[262] == 0x72) && return "tar"
|
||||||
|
return "other"
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- NCD (Normalized Compression Distance) baseline, model-free (DESIGN §8) ---
|
||||||
|
|
||||||
|
"gzip-compressed size of a byte buffer, via the gzip CLI (no CodecZlib dep)."
|
||||||
|
function gz_size(bytes::Vector{UInt8})
|
||||||
|
out = IOBuffer()
|
||||||
|
open(pipeline(`gzip -c`; stdout=out); write=true) do io
|
||||||
|
write(io, bytes)
|
||||||
|
end
|
||||||
|
return length(take!(out))
|
||||||
|
end
|
||||||
|
|
||||||
|
"NCD(x,y) = (C(xy) - min(C(x),C(y))) / max(C(x),C(y)); 0 = identical, ~1 = unrelated."
|
||||||
|
function ncd(xb, yb, cx, cy)
|
||||||
|
cxy = gz_size(vcat(xb, yb))
|
||||||
|
return (cxy - min(cx, cy)) / max(cx, cy)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
ncd_1nn_purity(paths, truth; head_bytes) -> Float64
|
||||||
|
|
||||||
|
Fraction of files whose NCD-nearest neighbour shares its true label: a cheap,
|
||||||
|
O(N²) sanity read on how well raw gzip-similarity alone separates formats on the
|
||||||
|
same input. The Bayesian clusters should broadly agree; a big gap is a red flag
|
||||||
|
(DESIGN §10.3). Uses each file's first `head_bytes` so the giant files don't
|
||||||
|
dominate compression time.
|
||||||
|
"""
|
||||||
|
function ncd_1nn_purity(paths::Vector{String}, truth::Vector{String}; head_bytes::Int=4096)
|
||||||
|
bufs = map(paths) do p
|
||||||
|
open(io -> read(io, head_bytes), p)
|
||||||
|
end
|
||||||
|
csz = gz_size.(bufs)
|
||||||
|
N = length(paths)
|
||||||
|
correct = 0
|
||||||
|
for i in 1:N
|
||||||
|
best_j = 0; best_d = Inf
|
||||||
|
for j in 1:N
|
||||||
|
i == j && continue
|
||||||
|
d = ncd(bufs[i], bufs[j], csz[i], csz[j])
|
||||||
|
if d < best_d
|
||||||
|
best_d = d; best_j = j
|
||||||
|
end
|
||||||
|
end
|
||||||
|
best_j != 0 && truth[best_j] == truth[i] && (correct += 1)
|
||||||
|
end
|
||||||
|
return correct / N
|
||||||
|
end
|
||||||
|
|
||||||
|
"1-NN label purity of a *cluster* assignment vs truth (same yardstick as NCD's)."
|
||||||
|
function cluster_1nn_purity(pred::Vector{Int}, truth::Vector{String})
|
||||||
|
# For each file, its 'nearest neighbour' is any other file in the same
|
||||||
|
# cluster; purity = P(a random same-cluster neighbour shares the true label).
|
||||||
|
groups = Dict{Int,Vector{Int}}()
|
||||||
|
for (i, k) in enumerate(pred)
|
||||||
|
push!(get!(groups, k, Int[]), i)
|
||||||
|
end
|
||||||
|
correct = 0; total = 0
|
||||||
|
for (_, idxs) in groups
|
||||||
|
length(idxs) < 2 && continue
|
||||||
|
for i in idxs
|
||||||
|
same = count(j -> j != i && truth[j] == truth[i], idxs)
|
||||||
|
total += 1
|
||||||
|
same > 0 && (correct += 1)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return total == 0 ? 0.0 : correct / total
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- data ---------------------------------------------------------------------
|
||||||
|
|
||||||
|
function load_corpus(dir::AbstractString)
|
||||||
|
paths = String[]
|
||||||
|
for name in readdir(dir; join=true)
|
||||||
|
isfile(name) && push!(paths, name)
|
||||||
|
end
|
||||||
|
truth = truth_label.(paths)
|
||||||
|
return paths, truth
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- grid search --------------------------------------------------------------
|
||||||
|
|
||||||
|
function evaluate(X, truth; α, β, bg_mass, sweeps, restarts, seed)
|
||||||
|
r = gibbs_cluster(X; α=α, β=β, bg_mass=bg_mass, sweeps=sweeps,
|
||||||
|
restarts=restarts, rng=MersenneTwister(seed))
|
||||||
|
pred = r.assignments
|
||||||
|
ari = adjusted_rand_index(truth, pred)
|
||||||
|
keep = truth .!= "tar"
|
||||||
|
ari_notar = adjusted_rand_index(truth[keep], pred[keep])
|
||||||
|
v, h, comp = v_measure(truth, pred)
|
||||||
|
return (; ari, ari_notar, v, h, comp, k=length(r.clusters),
|
||||||
|
bg=count(==(0), pred), result=r)
|
||||||
|
end
|
||||||
|
|
||||||
|
function main()
|
||||||
|
dir = length(ARGS) >= 1 ? ARGS[1] : joinpath(@__DIR__, "..", "..", "training_set")
|
||||||
|
isdir(dir) || error("training set dir not found: $dir")
|
||||||
|
paths, truth = load_corpus(dir)
|
||||||
|
classes = sort(unique(truth))
|
||||||
|
counts = [(c, count(==(c), truth)) for c in classes]
|
||||||
|
@printf("corpus: %d files from %s\n", length(paths), dir)
|
||||||
|
println("magic-collapsed truth classes: ", join(["$c=$n" for (c, n) in counts], " "))
|
||||||
|
println()
|
||||||
|
|
||||||
|
sweeps = 150
|
||||||
|
restarts = 6
|
||||||
|
seed = 20260703
|
||||||
|
|
||||||
|
# Grid. n is expensive to re-featurize, so loop it outermost. Ranges are
|
||||||
|
# centred where the coarse sweep found the optimum: small β (peaked
|
||||||
|
# per-position priors) is what separates formats whose headers differ in only
|
||||||
|
# a few magic bytes; large β over-merges. bg_mass barely moves the result
|
||||||
|
# here (almost nothing lands in background on this corpus), so it is fixed.
|
||||||
|
αs = [1.0, 2.0]
|
||||||
|
βs = [0.05, 0.08, 0.1, 0.15, 0.2]
|
||||||
|
bgs = [5.0]
|
||||||
|
ns = [32, 64]
|
||||||
|
|
||||||
|
println("grid search (sweeps=$sweeps, restarts=$restarts):")
|
||||||
|
@printf(" %-4s %-5s %-5s %-6s | %-6s %-8s %-6s %-6s %-6s %-4s %-4s\n",
|
||||||
|
"n", "alpha", "beta", "bgmss", "ARI", "ARI-tar", "V", "homog", "compl", "k", "bg")
|
||||||
|
results = Vector{Any}()
|
||||||
|
for n in ns
|
||||||
|
X = header_matrix(paths; n=n)
|
||||||
|
for α in αs, β in βs, bg in bgs
|
||||||
|
e = evaluate(X, truth; α=α, β=β, bg_mass=bg, sweeps=sweeps, restarts=restarts, seed=seed)
|
||||||
|
push!(results, (; n, α, β, bg, e))
|
||||||
|
@printf(" %-4d %-5.1f %-5.2f %-6.1f | %-6.3f %-8.3f %-6.3f %-6.3f %-6.3f %-4d %-4d\n",
|
||||||
|
n, α, β, bg, e.ari, e.ari_notar, e.v, e.h, e.comp, e.k, e.bg)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# Rank by ARI-excluding-tar (tar is the accepted blind spot; scoring it would
|
||||||
|
# penalise the correct answer of scattering tars to background; DESIGN §7.2).
|
||||||
|
sort!(results; by=r -> r.e.ari_notar, rev=true)
|
||||||
|
best = results[1]
|
||||||
|
println()
|
||||||
|
@printf("BEST (by ARI excl. tar): n=%d α=%.1f β=%.2f bg_mass=%.1f\n",
|
||||||
|
best.n, best.α, best.β, best.bg)
|
||||||
|
@printf(" ARI=%.3f ARI(excl tar)=%.3f V=%.3f homogeneity=%.3f completeness=%.3f clusters=%d background=%d\n",
|
||||||
|
best.e.ari, best.e.ari_notar, best.e.v, best.e.h, best.e.comp, best.e.k, best.e.bg)
|
||||||
|
|
||||||
|
# Per-cluster composition of the winning partition, and promotion nominations.
|
||||||
|
pred = best.e.result.assignments
|
||||||
|
println("\nwinning partition, cluster composition (truth breakdown):")
|
||||||
|
for (id, c) in sort(collect(best.e.result.clusters); by=x -> -x[2].members)
|
||||||
|
members = [truth[i] for i in eachindex(pred) if pred[i] == id]
|
||||||
|
comp = sort([(l, count(==(l), members)) for l in unique(members)]; by=x -> -x[2])
|
||||||
|
sig = signature(c)
|
||||||
|
promo = is_promotable(c, sig; min_members=20, min_magic=3) ? " ✓NOMINATE" : ""
|
||||||
|
@printf(" cluster %-4d n=%-3d magic=%-2d %s%s\n",
|
||||||
|
id, c.members, magic_positions(sig),
|
||||||
|
join(["$l:$k" for (l, k) in comp], " "), promo)
|
||||||
|
end
|
||||||
|
nbg = count(==(0), pred)
|
||||||
|
bg_truth = [truth[i] for i in eachindex(pred) if pred[i] == 0]
|
||||||
|
bgc = sort([(l, count(==(l), bg_truth)) for l in unique(bg_truth)]; by=x -> -x[2])
|
||||||
|
@printf(" background n=%-3d %s\n", nbg, join(["$l:$k" for (l, k) in bgc], " "))
|
||||||
|
|
||||||
|
# NCD baseline cross-check on a subsample (O(N²), so keep it small).
|
||||||
|
println("\nNCD (gzip) baseline cross-check:")
|
||||||
|
subn = min(150, length(paths))
|
||||||
|
sub = shuffle(MersenneTwister(seed), collect(1:length(paths)))[1:subn]
|
||||||
|
subpaths = paths[sub]; subtruth = truth[sub]
|
||||||
|
ncd_pur = ncd_1nn_purity(subpaths, subtruth)
|
||||||
|
Xsub = header_matrix(subpaths; n=best.n)
|
||||||
|
rsub = gibbs_cluster(Xsub; α=best.α, β=best.β, bg_mass=best.bg,
|
||||||
|
sweeps=sweeps, restarts=restarts, rng=MersenneTwister(seed))
|
||||||
|
bay_pur = cluster_1nn_purity(rsub.assignments, subtruth)
|
||||||
|
@printf(" subsample=%d NCD 1-NN label purity=%.3f Bayesian same-cluster purity=%.3f\n",
|
||||||
|
subn, ncd_pur, bay_pur)
|
||||||
|
println(" (both high ⇒ header-byte signal agrees with model-free gzip similarity; DESIGN §10.3)")
|
||||||
|
end
|
||||||
|
|
||||||
|
main()
|
||||||
19
bin/cluster_sweep.jl
Normal file
19
bin/cluster_sweep.jl
Normal file
@@ -0,0 +1,19 @@
|
|||||||
|
#!/usr/bin/env julia
|
||||||
|
#
|
||||||
|
# Stage-5 phase-B runner (model/DESIGN_clustering.md §9): the single-owner,
|
||||||
|
# periodic/cron process that sweeps `binary/`, folds new files into the durable
|
||||||
|
# format catalog by sequential CRP-predictive assignment, and (re)writes
|
||||||
|
# promotion nominations. Run it single-threaded on a schedule: it is the ONLY
|
||||||
|
# writer of the catalog, so no locking is needed.
|
||||||
|
#
|
||||||
|
# julia --project=. bin/cluster_sweep.jl # incremental live sweep
|
||||||
|
# julia --project=. bin/cluster_sweep.jl --compact # offline Gibbs (seed / recompact)
|
||||||
|
#
|
||||||
|
# On a fresh catalog (nothing processed yet) the incremental sweep would send
|
||||||
|
# every file to background (there are no clusters to match), so the first run
|
||||||
|
# auto-promotes to a compaction pass to seed the catalog. Configure via the
|
||||||
|
# FS_CLUSTER_* / FS_NOMINATED_DIR env vars (see src/config.jl).
|
||||||
|
|
||||||
|
using DarkStruct
|
||||||
|
|
||||||
|
DarkStruct.cluster_sweep_cli(ARGS)
|
||||||
@@ -1,6 +1,6 @@
|
|||||||
#!/usr/bin/env bash
|
#!/usr/bin/env bash
|
||||||
#
|
#
|
||||||
# send_dir.sh — upload every file in a directory to a running FileServer.
|
# send_dir.sh — upload every file in a directory to a running DarkStruct.
|
||||||
#
|
#
|
||||||
# Each file is POSTed to /upload as multipart/form-data. By default one file per
|
# Each file is POSTed to /upload as multipart/form-data. By default one file per
|
||||||
# request (so you get per-file feedback); use -b to batch all files into a
|
# request (so you get per-file feedback); use -b to batch all files into a
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
#!/usr/bin/env julia
|
#!/usr/bin/env julia
|
||||||
# Entry point. Run with: julia --project -t auto bin/server.jl
|
# Entry point. Run with: julia --project -t auto bin/server.jl
|
||||||
using FileServer
|
using DarkStruct
|
||||||
FileServer.run()
|
DarkStruct.run()
|
||||||
|
|||||||
55
docker-compose.rabbitmq.yml
Normal file
55
docker-compose.rabbitmq.yml
Normal file
@@ -0,0 +1,55 @@
|
|||||||
|
# Overlay that swaps the in-process queues for durable RabbitMQ ones.
|
||||||
|
#
|
||||||
|
# docker compose up # channels
|
||||||
|
# docker compose -f docker-compose.yml -f docker-compose.rabbitmq.yml up
|
||||||
|
#
|
||||||
|
# The base file is left alone deliberately: `docker compose up` still brings up
|
||||||
|
# exactly what it always did, so neither backend is the awkward one to run.
|
||||||
|
#
|
||||||
|
# What the broker changes: a job stays unacked until its handler commits, so a
|
||||||
|
# restart resumes each file at the stage it had reached instead of re-driving
|
||||||
|
# everything in spool/ from stage 1. See "Queue backends" in the README for the
|
||||||
|
# guarantees, and for the deliberate holes.
|
||||||
|
|
||||||
|
services:
|
||||||
|
rabbitmq:
|
||||||
|
image: rabbitmq:4-management
|
||||||
|
environment:
|
||||||
|
# The built-in guest/guest only authenticates over loopback, so it cannot
|
||||||
|
# work across the compose network. Override these for anything real; they
|
||||||
|
# are here so the stack comes up with one command, not as a credential.
|
||||||
|
RABBITMQ_DEFAULT_USER: "darkstruct"
|
||||||
|
RABBITMQ_DEFAULT_PASS: "darkstruct"
|
||||||
|
ports:
|
||||||
|
# 5672 is published so you can run the server on the host
|
||||||
|
# (`julia --project=. -t auto bin/server.jl`) against this broker; 15672 is
|
||||||
|
# the management UI, which is most of the reason to run a broker you can
|
||||||
|
# look at. Neither needs publishing for the composed server to work.
|
||||||
|
- "5672:5672"
|
||||||
|
- "15672:15672"
|
||||||
|
volumes:
|
||||||
|
# Durable queues and persistent messages are only as durable as what they
|
||||||
|
# are written to: without this, `docker compose down` discards the backlog
|
||||||
|
# the acks exist to protect.
|
||||||
|
- rabbitmq-data:/var/lib/rabbitmq
|
||||||
|
healthcheck:
|
||||||
|
test: ["CMD", "rabbitmq-diagnostics", "-q", "ping"]
|
||||||
|
interval: 10s
|
||||||
|
timeout: 5s
|
||||||
|
retries: 12
|
||||||
|
start_period: 30s
|
||||||
|
restart: unless-stopped
|
||||||
|
|
||||||
|
darkstruct:
|
||||||
|
environment:
|
||||||
|
FS_QUEUE_BACKEND: "rabbitmq"
|
||||||
|
FS_AMQP_URL: "amqp://darkstruct:darkstruct@rabbitmq:5672/"
|
||||||
|
depends_on:
|
||||||
|
rabbitmq:
|
||||||
|
# The server retries a broker that isn't up yet (AMQP_CONNECT_RETRY_SECONDS),
|
||||||
|
# so this gate is belt-and-braces rather than load-bearing — but it keeps
|
||||||
|
# a cold start quiet instead of a screenful of retry lines.
|
||||||
|
condition: service_healthy
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
rabbitmq-data:
|
||||||
27
docker-compose.yml
Normal file
27
docker-compose.yml
Normal file
@@ -0,0 +1,27 @@
|
|||||||
|
services:
|
||||||
|
darkstruct:
|
||||||
|
build:
|
||||||
|
context: .
|
||||||
|
args:
|
||||||
|
# false drops the Ruby toolchain and github-linguist (~160MB smaller).
|
||||||
|
# Stage-4 text files then get natural-language enrichment but no
|
||||||
|
# programming language; the server warns at startup and carries on.
|
||||||
|
WITH_LINGUIST: "true"
|
||||||
|
image: darkstruct:latest
|
||||||
|
ports:
|
||||||
|
- "8080:8080"
|
||||||
|
volumes:
|
||||||
|
# Every pipeline stage (spool, done, text_done, binary, failed, nominated)
|
||||||
|
# lives under /data, so one volume holds the whole in-flight working set.
|
||||||
|
- fs-data:/data
|
||||||
|
environment:
|
||||||
|
# The four pools each default to nthreads(), which oversubscribes the
|
||||||
|
# threads Julia actually has. Size them for the host instead.
|
||||||
|
FS_WORKERS: "4"
|
||||||
|
FS_KNOWN_WORKERS: "4"
|
||||||
|
FS_UNKNOWN_WORKERS: "2"
|
||||||
|
FS_TEXT_WORKERS: "2"
|
||||||
|
restart: unless-stopped
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
fs-data:
|
||||||
314
model/DESIGN_clustering.md
Normal file
314
model/DESIGN_clustering.md
Normal file
@@ -0,0 +1,314 @@
|
|||||||
|
# Stage-5: Unknown-format discovery by Bayesian header clustering
|
||||||
|
|
||||||
|
Status: phases A and B implemented and calibrated (`src/cluster.jl` +
|
||||||
|
`src/catalog.jl`, `bin/cluster_calibrate.jl` + `bin/cluster_sweep.jl`, tests in
|
||||||
|
`test/runtests.jl`). Phase A (offline Gibbs) is calibrated; phase B's durable
|
||||||
|
single-owner catalog, incremental sweep, and nomination writer are now built on
|
||||||
|
top of the `assign_file` scoring core. This is the product of a design
|
||||||
|
interview. It captures the decisions and, just as importantly, the assumptions
|
||||||
|
we *rejected*, so they don't get silently reintroduced. §11 records
|
||||||
|
what building it actually taught us, including three assumptions in this document
|
||||||
|
that the data corrected.
|
||||||
|
|
||||||
|
## 1. Goal
|
||||||
|
|
||||||
|
Discover recurring new file formats hiding in the `binary/` bucket (the
|
||||||
|
`:unknown` sink from `classify.jl` → stage-3 triage). A genuinely novel format is
|
||||||
|
a plausible proxy for a genuinely novel producing application, but we do not
|
||||||
|
try to identify producers directly (see §3). The output is a **growing catalog of
|
||||||
|
discovered formats**, each with a magic-byte signature that can be promoted into
|
||||||
|
the classifier's fast path.
|
||||||
|
|
||||||
|
Task shape (settled): **unsupervised clustering with an unknown number of
|
||||||
|
clusters.** Not pairwise "same producer" scoring, not classification against a
|
||||||
|
fixed label set.
|
||||||
|
|
||||||
|
## 2. Two phases: build (A) then run (B)
|
||||||
|
|
||||||
|
**(A) Batch, offline, the science.** Cluster the accumulated pile from scratch.
|
||||||
|
Its job is *not* to be the catalog; it is to (i) prove the header-byte signal
|
||||||
|
actually separates formats, cross-checked against an NCD baseline (§8), and
|
||||||
|
(ii) calibrate hyperparameters against known formats (§7). Ship this first,
|
||||||
|
because it de-risks (B). If (A)'s clusters are garbage, (B)'s machinery is wasted.
|
||||||
|
|
||||||
|
**(B) Online, live, the catalog.** The target deliverable. A persistent catalog
|
||||||
|
where each discovered format has a durable, frozen ID and stored sufficient
|
||||||
|
statistics. New unknown files are scored against existing clusters; only
|
||||||
|
genuinely novel ones spawn a new entry. Clusters that accumulate enough evidence
|
||||||
|
are nominated for promotion into the classifier (§6).
|
||||||
|
|
||||||
|
## 3. What we are and are NOT clustering
|
||||||
|
|
||||||
|
We cluster by file format, not by producer. The first-*n* header bytes are
|
||||||
|
format-mandated and producer-invariant: every valid PNG shares the same magic
|
||||||
|
regardless of which program wrote it; a PDF's producer string lives deep inside
|
||||||
|
the file, not in the header. Producer identity, where recoverable at all, is
|
||||||
|
`exiftool`'s job (stage 2), not this stage's.
|
||||||
|
|
||||||
|
Corollary already visible in `../training_set`: extension labels are not
|
||||||
|
header-format labels. `docx` *is* a PK zip; `so`/`o`/`elf`/`out` are all ELF.
|
||||||
|
Merging those is correct, not error (see §7).
|
||||||
|
|
||||||
|
## 4. Model: DP mixture of per-position categoricals
|
||||||
|
|
||||||
|
A cluster is a product of independent per-position categorical distributions
|
||||||
|
over the first *n* header bytes. Position *i* carries a distribution `θᵢ` over a
|
||||||
|
**257-symbol alphabet**: byte values `0–255`, plus symbol `256 = "past EOF"`.
|
||||||
|
|
||||||
|
- Invariant positions (magic bytes) learn a spiked `θᵢ`; variable positions
|
||||||
|
(lengths, timestamps) learn a flat one. A cluster's signature = the vector of
|
||||||
|
modal symbols + per-position peakedness. That signature **is a magic-number
|
||||||
|
template**, which is the entire reason for the categorical choice.
|
||||||
|
- `257` alphabet handles short files honestly: a format that is always 20 bytes
|
||||||
|
produces a spiked "past-EOF" at positions 20–31, which is real, discriminative
|
||||||
|
signal. No zero-padding (would collide `0x00` padding with real `0x00` bytes).
|
||||||
|
|
||||||
|
**Priors:** Dirichlet on each `θᵢ` (conjugate to Categorical); **Dirichlet
|
||||||
|
process (CRP)** over cluster assignments → unknown *k* falls out natively.
|
||||||
|
|
||||||
|
**Why categorical, not Euclidean.** Bytes are categorical, not ordinal: `0x89`
|
||||||
|
and `0x88` are not "close," `0x00` and `0xFF` are not "far." k-means / Gaussian
|
||||||
|
mixtures over scaled bytes assert a metric that does not exist in header space.
|
||||||
|
**Do not reuse `model.jl`'s `[0,1]` byte scaling here.** That scaling is correct
|
||||||
|
for the Lux net and wrong for this model. We need the raw `0–255` byte as a
|
||||||
|
categorical index.
|
||||||
|
|
||||||
|
### 4a. Background component (high-entropy handling)
|
||||||
|
|
||||||
|
Add a fixed, non-adaptive uniform component (each position uniform over 257)
|
||||||
|
as the "junk drawer." Compressed/encrypted/structureless blobs are ~uniform after
|
||||||
|
any magic and would otherwise either (i) mint a singleton per file or (ii)
|
||||||
|
collapse into one flat cluster that then matches everything. The background
|
||||||
|
absorbs them cleanly.
|
||||||
|
|
||||||
|
Two populations, to be precise:
|
||||||
|
- **Structured prefix + random tail** (gzip `1f 8b`, PK zip, zstd, most encrypted
|
||||||
|
*containers*): peaked at positions 0–3, flat after. These form **real clusters
|
||||||
|
for free**: genuine discoveries, no special handling.
|
||||||
|
- **Uniform from byte 0** (raw encrypted streams, key material): nothing in the
|
||||||
|
header to cluster on → absorbed by background.
|
||||||
|
|
||||||
|
The background is never promotable. But it is not a silent sink: its
|
||||||
|
size / growth / entropy histogram is surfaced as a first-class signal ("12% of
|
||||||
|
this week's unknowns are structureless"). If sub-clustering the structureless
|
||||||
|
residue ever matters, that needs a *different* feature (byte histogram / entropy),
|
||||||
|
a separate v3 model. Header bytes genuinely cannot do it.
|
||||||
|
|
||||||
|
### 4b. Feature window
|
||||||
|
|
||||||
|
**Front-only, `n = 32`** (config knob; try 64 if under-resolved). Magic lives at
|
||||||
|
offset 0. Tail window deferred to v2: a minority of formats have trailers
|
||||||
|
(ZIP EOCD, ID3v1, PDF `%%EOF`); add as an independent *second block* of positions
|
||||||
|
only if real trailer-formats show up in the residue.
|
||||||
|
|
||||||
|
**Known blind spot: tar.** `ustar` magic is at offset 257, outside the window,
|
||||||
|
so all 100 training tars scatter to background. Accepted for v1, since tar is
|
||||||
|
already a *known* format and discovery doesn't need it. General lesson: a
|
||||||
|
minority of formats put magic at a fixed deeper offset; the fix (if ever needed)
|
||||||
|
is a sparse probe window at that offset (e.g. bytes 257–262 as a third block),
|
||||||
|
rather than densely modeling 257 front bytes, which would 8× every cluster's
|
||||||
|
`n×257` sufficient-stat table to catch one format.
|
||||||
|
|
||||||
|
## 5. Inference: different mode per phase (resolves the Bayesian-vs-catalog tension)
|
||||||
|
|
||||||
|
A sampler yields a *posterior over partitions*; a catalog needs *one partition
|
||||||
|
with durable IDs*. Two MCMC gotchas: label switching (cluster #3 is not a
|
||||||
|
stable identity across iterations/runs) and distribution-not-answer (1000
|
||||||
|
partitions, not one). We sidestep both by using two inference modes:
|
||||||
|
|
||||||
|
- **Phase (A), offline:** full collapsed Gibbs sampler over the
|
||||||
|
Dirichlet-Categorical (conjugacy → ~100 lines, no continuous approximation,
|
||||||
|
unknown *k* native). Used to validate signal, tune `α` + Dirichlet strength,
|
||||||
|
and seed the initial catalog (summarize to a point partition once, via a
|
||||||
|
VI/Binder loss over the posterior similarity matrix, tolerated because it is
|
||||||
|
offline, never in the hot path).
|
||||||
|
- **Phase (B), live:** deterministic sequential CRP-predictive assignment.
|
||||||
|
Each catalog cluster stores per-position 257-count vectors (sufficient stats).
|
||||||
|
A new file's CRP predictive probability of joining each existing cluster vs.
|
||||||
|
the background vs. spawning a new cluster is computed; assign to the argmax.
|
||||||
|
A new cluster is minted only if the new-cluster evidence beats the background
|
||||||
|
by a margin. IDs are frozen at birth → no label switching. This is exactly
|
||||||
|
the Gibbs predictive rule with existing assignments held fixed: the same math,
|
||||||
|
not an ad-hoc hack.
|
||||||
|
- **Periodic compaction, offline:** re-run Gibbs seeded from the current catalog
|
||||||
|
to merge drifted clusters / split bloated ones.
|
||||||
|
|
||||||
|
## 6. Promotion (closing the loop to the classifier)
|
||||||
|
|
||||||
|
**Layered known-check at ingest** becomes:
|
||||||
|
1. Match against promoted signatures (exact, fast), which runs *before* the net.
|
||||||
|
2. Else the Lux `:known` / `:unknown` classifier.
|
||||||
|
3. Else route to `binary/` for this stage.
|
||||||
|
|
||||||
|
**Promotion = append a magic-byte signature to a registry.** A cluster's spiked
|
||||||
|
positions (posterior max-prob `> ~0.9`) become required bytes; flat positions
|
||||||
|
become wildcards, giving a libmagic-style signature. This is a **data change, not a
|
||||||
|
retrain**; interpretable, auditable, reversible. Retraining the Lux net is a
|
||||||
|
separate, *optional periodic* activity using accumulated signature-labeled files,
|
||||||
|
never the promotion mechanism itself.
|
||||||
|
|
||||||
|
**Nominate automatically, activate by hand.** A cluster crossing the thresholds
|
||||||
|
(`≥ N` members, starting at `N ≈ 20–50`, a loose dial since a human is the
|
||||||
|
backstop; `≥ ~3` magic positions; and not the background) is written to a
|
||||||
|
`nominated/` registry with its signature, member count, and example files. A human glance
|
||||||
|
promotes it into the active set. Human gate guards the one hard-to-reverse action
|
||||||
|
(redefining "known"); everything upstream stays automatic.
|
||||||
|
|
||||||
|
## 7. Calibration: recover known formats, then trust on unknowns
|
||||||
|
|
||||||
|
Do not pick priors blind. We have ground truth: `../training_set` (100 each of
|
||||||
|
tgz/tar/pdf/docx, 98 zip, 93 jpg, ELF family) and the `data/done` corpus.
|
||||||
|
|
||||||
|
1. Run labeled known files through the exact clustering pipeline.
|
||||||
|
2. Ground truth = magic-collapsed classes, *not* extensions:
|
||||||
|
`{gzip (tgz), PKzip (docx≡zip), ELF (so/o/elf/out/x86_64), JPEG, PDF, tar}`.
|
||||||
|
Merging docx+zip and the ELF family is the correct answer; scoring
|
||||||
|
against raw extensions would penalize correctness and mistune `α`.
|
||||||
|
3. Measure recovered-vs-truth agreement with Adjusted Rand Index / V-measure.
|
||||||
|
4. **Grid-tune `α` and the Dirichlet pseudocount to maximize agreement**, which
|
||||||
|
finds the settings at which the machine rediscovers formats we already know.
|
||||||
|
5. Freeze, deploy on the `:unknown` pile.
|
||||||
|
|
||||||
|
Splitting docx from zip is a later tier: the discriminating info
|
||||||
|
(central-directory filenames like `word/document.xml`) sits at a *variable
|
||||||
|
offset*, not a fixed position. That is a different feature problem, deferred.
|
||||||
|
|
||||||
|
## 8. Julia package surface
|
||||||
|
|
||||||
|
- **Hand-rolled collapsed Gibbs** over Dirichlet-Categorical. Recommended: the
|
||||||
|
conjugacy makes it short/fast; we own the online + promotion logic; no library
|
||||||
|
impedance. `Distributions.jl` for `Dirichlet`/`Categorical` primitives.
|
||||||
|
- **`CodecZlib`** for the NCD (Normalized Compression Distance) baseline:
|
||||||
|
model-free gzip-similarity clustering. Excellent at format grouping and a great
|
||||||
|
step-(A) sanity check, but O(N²), non-generative (no signature, no online
|
||||||
|
scoring, no promotion) → baseline only, cannot be the catalog.
|
||||||
|
- **`Clustering.jl`** for `randindex` / `vmeasure`, the §7 calibration metric,
|
||||||
|
plus a throwaway k-modes-ish baseline. Not the real model (its k-means is
|
||||||
|
the Euclidean trap of §4).
|
||||||
|
- **`Turing.jl`**, considered and rejected: discrete assignment latents + DP are
|
||||||
|
awkward, and we'd still hand-roll the online path. Overkill.
|
||||||
|
|
||||||
|
## 9. Architecture: single-owner batch stage, NOT inline inference
|
||||||
|
|
||||||
|
The classifier is stateless, immutable, shared read-only across worker threads
|
||||||
|
(see `classify.jl`). The catalog is the opposite: mutable, learned, shared, since
|
||||||
|
every assigned file updates a cluster's counts. It therefore must not copy the
|
||||||
|
classifier's concurrency model (concurrent workers → lock contention, torn reads
|
||||||
|
of sufficient stats, CRP assignment against stale mass).
|
||||||
|
|
||||||
|
Because promotion is human-gated, nothing here is latency-sensitive. So:
|
||||||
|
|
||||||
|
- Workers stay stateless. They deposit `:unknown` files into `binary/` exactly as
|
||||||
|
today. No catalog access on the hot path.
|
||||||
|
- A separate stage-5 process (periodic / cron, single-threaded) owns the
|
||||||
|
catalog exclusively: sweeps newly-arrived `binary/` files, runs sequential
|
||||||
|
CRP-predictive assignment, updates sufficient stats, writes nominations.
|
||||||
|
**One writer, zero locks, no cross-thread shared mutable state.**
|
||||||
|
- The catalog is a durable file mutated by one process, so reuse the stage-2
|
||||||
|
**sidecar-first durable-commit** discipline (`commit_enriched!`: temp write →
|
||||||
|
fsync bytes → rename → fsync dir) so a crash can't corrupt it or lose a rename.
|
||||||
|
|
||||||
|
This slots in as a batch stage, matching how stages 2/3/4 already work. New
|
||||||
|
`Config` knobs follow the existing `FS_*` env-override convention (e.g.
|
||||||
|
`FS_CLUSTER_DIR`, `FS_CLUSTER_N`, `FS_CLUSTER_ALPHA`, `FS_CLUSTER_PSEUDOCOUNT`,
|
||||||
|
`FS_PROMOTE_MIN_MEMBERS`).
|
||||||
|
|
||||||
|
## 10. Concrete test assertions (write these first)
|
||||||
|
|
||||||
|
1. **Discovers nothing from noise.** Current `data/binary` = 20 small random
|
||||||
|
blobs + 1 giant PDF. Correct output: PDF is a singleton that never promotes
|
||||||
|
(N=1), 20 blobs absorbed by background, zero promoted clusters. Any
|
||||||
|
promoted cluster from this pile = broken.
|
||||||
|
2. **Recovers known formats.** On a `../training_set` sample, calibrated settings
|
||||||
|
cluster into the ~6 magic-collapsed classes with high ARI (docx+zip merged,
|
||||||
|
ELF family merged, tar in background as the accepted blind spot).
|
||||||
|
3. **NCD agreement.** Step-(A) Bayesian clusters broadly agree with the NCD
|
||||||
|
baseline on the same input; large disagreement is a red flag to investigate
|
||||||
|
before trusting the generative model.
|
||||||
|
|
||||||
|
## 11. Implementation status & calibration results (v1)
|
||||||
|
|
||||||
|
**Shipped.** `src/cluster.jl` holds feature extraction (`header_symbols`, 257-symbol
|
||||||
|
alphabet), collapsed Gibbs (`gibbs_cluster`, phase A), the sequential
|
||||||
|
CRP-predictive rule (`assign_file`, phase B core), signatures/promotion
|
||||||
|
(`signature`, `is_promotable`), and calibration metrics (`adjusted_rand_index`,
|
||||||
|
`v_measure`). All base-Julia: a base-only Lanczos `loggamma` keeps the
|
||||||
|
Dirichlet-multinomial marginal dependency-free (no Manifest churn). Config knobs
|
||||||
|
`FS_CLUSTER_*` (§9) added. `bin/cluster_calibrate.jl` runs the §7 grid and the §8
|
||||||
|
NCD baseline. Concrete §10 assertions are in the test suite (hermetic synthetic
|
||||||
|
corpora, so they need neither `../training_set` nor gzip).
|
||||||
|
|
||||||
|
**Calibrated defaults** (grid over the 700-file `training_set`, ranked by ARI
|
||||||
|
excluding tar): n=32, α=1.0, β=0.1, bg_mass=5.0 → ARI 0.77 (0.885 excl.
|
||||||
|
tar), V-measure 0.83, homogeneity 0.87. Clusters are clean and promotable:
|
||||||
|
`pkzip:197` (docx+zip correctly merged, §7.2 ✓), `gzip:100`, `jpeg`, and several
|
||||||
|
`pdf` clusters all self-nominate. The NCD baseline agrees (§10.3): on a
|
||||||
|
150-file subsample, NCD 1-NN label purity 0.90 vs. the model's same-cluster
|
||||||
|
purity 0.987. The generative header model separates formats at least as well as
|
||||||
|
model-free gzip similarity.
|
||||||
|
|
||||||
|
### Three assumptions the data corrected
|
||||||
|
|
||||||
|
1. **Tar is not in the background here; it merges into ELF.** §4b/§7 assumed
|
||||||
|
tar's `ustar`-at-257 magic is out of window so tars scatter to background. But
|
||||||
|
98/100 tars in the corpus are Hex/Elixir package tarballs whose *first
|
||||||
|
archived file is named `VERSION`* → a constant, strongly-peaked `VERSION\0`
|
||||||
|
prefix at offset 0. They do form a peaked cluster, but it merges with ELF,
|
||||||
|
because ELF's ident padding and tar's name-field zero-padding give the two a
|
||||||
|
long shared run of `0x00` in bytes 5–31; they differ in only ~3 magic bytes,
|
||||||
|
and 32 equally-weighted positions let ~20 shared zeros outvote 3 real ones. No
|
||||||
|
β both separates ELF/tar and keeps the other formats whole. The honest v1
|
||||||
|
position: this is the *same* "tar is hard" reality §4b flagged, just wearing a
|
||||||
|
different mask. Fix (v2): weight positions by inverse entropy so a
|
||||||
|
low-information shared-zero run stops dominating a few high-information magic
|
||||||
|
bytes. This generalizes beyond tar and is the highest-value next lever.
|
||||||
|
|
||||||
|
2. **You cannot cold-start every point in the background.** A natural reading of
|
||||||
|
§4a/§5 is "everything starts in the junk drawer, real clusters condense out."
|
||||||
|
That deadlocks: at a format's first file a fresh cluster and the background
|
||||||
|
are *both* uniform, so with the `bg_mass ≥ α` that §4a needs for absorption,
|
||||||
|
the background always wins and no cluster is ever seeded. Fix: **initialize
|
||||||
|
every file in its own singleton**; same-format singletons merge and snowball,
|
||||||
|
while a lone random-blob singleton dissolves on resample and is reclaimed by
|
||||||
|
the (stickier) background. Absorption still works, just not as the *initial*
|
||||||
|
state.
|
||||||
|
|
||||||
|
3. **Two pieces of math that look optional but aren't.** (a) Signature peakedness
|
||||||
|
is a Bernoulli question ("is this position fixed to byte v?"), so it uses a
|
||||||
|
2-way posterior `(count+β)/(members+2β)`, not the 257-way mixture
|
||||||
|
predictive, because the alphabet-wide denominator drags even a unanimous position
|
||||||
|
below 0.9 once β<1, which would make promotion *impossible*. (b) Ranking Gibbs
|
||||||
|
restarts needs the collapsed Dirichlet-multinomial marginal (with its
|
||||||
|
`loggamma` normalizer / Occam penalty); a plain product-of-predictives score
|
||||||
|
omits the penalty and actively rewards merging everything into one blob
|
||||||
|
(observed, then fixed).
|
||||||
|
|
||||||
|
### Known v1 limitations (accepted)
|
||||||
|
|
||||||
|
- **β=0.1 over-splits** PDF and JPEG into several *pure* sub-clusters (e.g. PDF by
|
||||||
|
version byte). This costs completeness/ARI but not the mission: each sub-cluster
|
||||||
|
still carries valid magic and promotes independently, and a human dedupes
|
||||||
|
overlapping `%PDF-1.x` nominations at the gate.
|
||||||
|
- The point partition is the best of N Gibbs restarts by marginal likelihood,
|
||||||
|
a MAP-style stand-in for the VI/Binder posterior summary §5 defers. Adequate,
|
||||||
|
because the formats are strongly separated; revisit if compaction (§5) needs it.
|
||||||
|
- Phase B's live single-owner batch process (§9) and the durable catalog file
|
||||||
|
are implemented in `src/catalog.jl` (the `Catalog` durable state, the
|
||||||
|
incremental `catalog_sweep!`, the offline `compact!` seed/recompaction, and
|
||||||
|
`write_nominations!`), driven by `bin/cluster_sweep.jl` (cron/periodic; the
|
||||||
|
first run auto-compacts to seed, subsequent runs sweep incrementally). The
|
||||||
|
catalog is persisted with the stage-2 sidecar-first temp→fsync→rename→fsync-dir
|
||||||
|
discipline. As §5B predicts, under the calibrated `bg_mass > α` the live sweep
|
||||||
|
never mints single-file clusters: new formats are discovered by the offline
|
||||||
|
`compact!` re-clustering the background residue, not by the live path.
|
||||||
|
|
||||||
|
## Open items (deferred, intentionally)
|
||||||
|
|
||||||
|
- **v2, now top priority: inverse-entropy position weighting** (unblocks ELF/tar
|
||||||
|
and any format pair that shares a long constant run; see §11).
|
||||||
|
|
||||||
|
- v2: tail-window block; sparse deep-offset probe (tar-class).
|
||||||
|
- v3: sub-clustering structureless high-entropy residue (needs entropy/histogram
|
||||||
|
feature, not header bytes).
|
||||||
|
- Later tier: docx-vs-zip split via variable-offset central-directory names.
|
||||||
|
- Periodic Lux retrain from accumulated signature-labeled files.
|
||||||
254
src/DarkStruct.jl
Normal file
254
src/DarkStruct.jl
Normal file
@@ -0,0 +1,254 @@
|
|||||||
|
module DarkStruct
|
||||||
|
|
||||||
|
using Logging
|
||||||
|
using Random
|
||||||
|
using UUIDs
|
||||||
|
using HTTP
|
||||||
|
using JSON3
|
||||||
|
using Oxygen
|
||||||
|
using Lux
|
||||||
|
using JLD2
|
||||||
|
using Languages
|
||||||
|
|
||||||
|
include("multipart.jl") # streaming multipart reader (defines UPLOAD_CHUNK_BYTES, used by config.jl)
|
||||||
|
include("config.jl")
|
||||||
|
include("job.jl")
|
||||||
|
include("queue.jl")
|
||||||
|
include("stats.jl") # per-stage counters behind GET /stats (needs Config/Job/JobQueue)
|
||||||
|
include("rabbit.jl") # RabbitMQ-backed JobQueue (needs Config/Job/JobQueue/STAGE_KEYS)
|
||||||
|
include("spool.jl")
|
||||||
|
include("model.jl") # build_model() + read_features(); shared with bin/train.jl
|
||||||
|
include("classify.jl") # Classifier + load_classifier/classify (needs model.jl)
|
||||||
|
include("metadata.jl") # exiftool extraction + sidecar enrichment (stage 2)
|
||||||
|
include("content.jl") # binary-vs-text triage for unknown files (stage 3)
|
||||||
|
include("language.jl") # natural + programming language enrichment for text (stage 4)
|
||||||
|
include("cluster.jl") # unknown-format discovery by header clustering (stage 5, science)
|
||||||
|
include("catalog.jl") # durable single-owner format catalog (stage 5, phase B; needs cluster.jl + metadata.jl fsync)
|
||||||
|
include("worker.jl")
|
||||||
|
|
||||||
|
# Globals the HTTP handlers read at request time. Set once in `run`, before the
|
||||||
|
# server starts accepting connections. Declared after the includes above so the
|
||||||
|
# `Config`/`JobQueue` types exist.
|
||||||
|
#
|
||||||
|
# Typed as the abstract `JobQueue` because which implementation these hold is a
|
||||||
|
# runtime choice (`FS_QUEUE_BACKEND`). That makes `enqueue!` a dynamic dispatch
|
||||||
|
# on the intake path — nanoseconds against a handler that writes a file to disk,
|
||||||
|
# and the price of the backend being swappable at all.
|
||||||
|
const CONFIG = Ref{Config}()
|
||||||
|
const QUEUE = Ref{JobQueue}() # stage-1 (classification) queue; HTTP intake enqueues here
|
||||||
|
const KNOWN_QUEUE = Ref{JobQueue}() # stage-2 (enrichment) queue; stage-1 workers enqueue here
|
||||||
|
const UNKNOWN_QUEUE = Ref{JobQueue}() # stage-3 (content triage) queue; stage-1 workers enqueue here
|
||||||
|
const TEXT_QUEUE = Ref{JobQueue}() # stage-4 (language enrichment) queue; stage-3 workers enqueue here
|
||||||
|
const CLASSIFIER = Ref{Classifier}() # loaded once at startup, shared read-only across workers
|
||||||
|
const DETECTOR = Ref{LanguageDetector}() # natural-language detector; built once at startup, shared read-only
|
||||||
|
|
||||||
|
include("server.jl") # registers routes (references CONFIG/QUEUE at call time)
|
||||||
|
|
||||||
|
export run
|
||||||
|
|
||||||
|
# When stderr is a live terminal Julia flushes each write, but when it's
|
||||||
|
# redirected to a file or pipe (a log file, `tee`, journald, a container log
|
||||||
|
# driver) Julia block-buffers it, so a long-running server's logs sit unseen in
|
||||||
|
# the buffer until it fills or the process exits, making it look like nothing is
|
||||||
|
# happening. This wrapper delegates to a normal logger and flushes after every
|
||||||
|
# message so output appears immediately wherever stderr is pointed.
|
||||||
|
struct FlushLogger{L<:AbstractLogger} <: AbstractLogger
|
||||||
|
inner::L
|
||||||
|
end
|
||||||
|
Logging.min_enabled_level(f::FlushLogger) = Logging.min_enabled_level(f.inner)
|
||||||
|
Logging.shouldlog(f::FlushLogger, args...) = Logging.shouldlog(f.inner, args...)
|
||||||
|
Logging.catch_exceptions(f::FlushLogger) = Logging.catch_exceptions(f.inner)
|
||||||
|
function Logging.handle_message(f::FlushLogger, args...; kwargs...)
|
||||||
|
Logging.handle_message(f.inner, args...; kwargs...)
|
||||||
|
flush(f.inner.stream)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
run(; overrides...)
|
||||||
|
|
||||||
|
Start DarkStruct: build config, ensure directories, recover any leftover
|
||||||
|
spooled files, spawn the worker pool, then serve HTTP until interrupted
|
||||||
|
(Ctrl-C / SIGINT or SIGTERM). On shutdown it stops accepting uploads, drains the
|
||||||
|
queue, waits for workers to finish in-flight files, and exits cleanly.
|
||||||
|
|
||||||
|
Keyword `overrides` (e.g. `port=9000`) take precedence over environment vars.
|
||||||
|
"""
|
||||||
|
function run(; overrides...)
|
||||||
|
# By default a Julia script exits immediately on SIGINT (Ctrl-C). Disable
|
||||||
|
# that so the interrupt surfaces as a catchable InterruptException, letting
|
||||||
|
# us drain the queue gracefully below.
|
||||||
|
Base.exit_on_sigint(false)
|
||||||
|
|
||||||
|
# Flush every log message so output is visible in real time even when stderr
|
||||||
|
# is redirected to a file/pipe (see FlushLogger). Set before anything logs.
|
||||||
|
global_logger(FlushLogger(ConsoleLogger(stderr)))
|
||||||
|
|
||||||
|
cfg = config_from_env(; overrides...)
|
||||||
|
ensure_dirs(cfg)
|
||||||
|
|
||||||
|
# All pools draw from the same OS threads. Warn on the *combined* size (still
|
||||||
|
# allowed): oversubscription just means tasks share threads, not a failure.
|
||||||
|
total_workers = cfg.worker_count + cfg.known_worker_count + cfg.unknown_worker_count + cfg.text_worker_count
|
||||||
|
if total_workers > Threads.nthreads()
|
||||||
|
@warn "combined worker count exceeds available threads; workers will share threads (start Julia with -t N for real parallelism)" classify_workers=cfg.worker_count known_workers=cfg.known_worker_count unknown_workers=cfg.unknown_worker_count text_workers=cfg.text_worker_count total=total_workers nthreads=Threads.nthreads()
|
||||||
|
end
|
||||||
|
|
||||||
|
# exiftool is a hard prerequisite for stage-2 enrichment. Fail fast at
|
||||||
|
# startup rather than discover it missing on the first known file.
|
||||||
|
assert_exiftool()
|
||||||
|
|
||||||
|
# github-linguist powers stage-4 *programming*-language detection, but it's
|
||||||
|
# best-effort (natural-language enrichment stands on its own), so a missing
|
||||||
|
# binary is a warning, not a fatal error; per-file lookups degrade to none.
|
||||||
|
linguist_available() || @warn "github-linguist not found on PATH; stage-4 text files will have no programming language (install it to enable)"
|
||||||
|
|
||||||
|
# The one place the backend choice is made. Everything downstream — handlers,
|
||||||
|
# workers, /stats — sees only `JobQueue` (src/queue.jl).
|
||||||
|
# Set if the broker connection dies under us. This backend does not
|
||||||
|
# reconnect — see `poll_loop!` — so the honest response is to drain and let
|
||||||
|
# the restart policy bring us back to redelivered messages.
|
||||||
|
connection_lost = Threads.Atomic{Bool}(false)
|
||||||
|
backend = cfg.queue_backend === :rabbitmq ? connect_backend(cfg) : nothing
|
||||||
|
queues = backend === nothing ?
|
||||||
|
(classify = ChannelQueue(cfg.queue_capacity),
|
||||||
|
enrich = ChannelQueue(cfg.known_queue_capacity),
|
||||||
|
triage = ChannelQueue(cfg.unknown_queue_capacity),
|
||||||
|
language = ChannelQueue(cfg.text_queue_capacity)) :
|
||||||
|
open_queues(backend, cfg; on_lost = () -> (connection_lost[] = true))
|
||||||
|
queue, known_queue = queues.classify, queues.enrich
|
||||||
|
unknown_queue, text_queue = queues.triage, queues.language
|
||||||
|
CONFIG[] = cfg
|
||||||
|
QUEUE[] = queue
|
||||||
|
KNOWN_QUEUE[] = known_queue
|
||||||
|
UNKNOWN_QUEUE[] = unknown_queue
|
||||||
|
TEXT_QUEUE[] = text_queue
|
||||||
|
|
||||||
|
# Load the classifier before serving. Fail fast: a server that silently
|
||||||
|
# doesn't classify is a worse surprise than a clear startup error.
|
||||||
|
CLASSIFIER[] = load_classifier(cfg.model_path)
|
||||||
|
@info "loaded classifier" path=cfg.model_path
|
||||||
|
|
||||||
|
# Build the natural-language detector once (it loads the whatlang n-gram
|
||||||
|
# model) and share it read-only across the stage-4 pool, like the classifier.
|
||||||
|
DETECTOR[] = LanguageDetector()
|
||||||
|
@info "loaded language detector"
|
||||||
|
|
||||||
|
@info "starting darkstruct" backend=cfg.queue_backend host=cfg.host port=cfg.port workers=cfg.worker_count known_workers=cfg.known_worker_count unknown_workers=cfg.unknown_worker_count text_workers=cfg.text_worker_count capacity=cfg.queue_capacity known_capacity=cfg.known_queue_capacity unknown_capacity=cfg.unknown_queue_capacity text_capacity=cfg.text_queue_capacity
|
||||||
|
|
||||||
|
# Zero the counters here, not at module load: `since` should mean "serving
|
||||||
|
# since", so a scrape's totals cover the run, not the minutes spent loading
|
||||||
|
# the classifier. It must also come *before* the pools start, since recovery
|
||||||
|
# below runs against live workers. The replayed files are real work and
|
||||||
|
# belong in the totals, but only from this point on.
|
||||||
|
reset_metrics!()
|
||||||
|
|
||||||
|
# Each pool gets its stage's counters (src/stats.jl); `worker_loop` records
|
||||||
|
# into them, `GET /stats` reads them out. Stage 1 and 3 also hand theirs to
|
||||||
|
# their handler, which charges time parked on a full downstream queue to
|
||||||
|
# `blocked_ns` so it isn't mistaken for work.
|
||||||
|
st = METRICS.stages
|
||||||
|
workers = [Threads.@spawn worker_loop(i, cfg, queue,
|
||||||
|
(job, c, wid) -> handle_classify_job(job, c, wid, known_queue, unknown_queue, st.classify),
|
||||||
|
st.classify)
|
||||||
|
for i in 1:cfg.worker_count]
|
||||||
|
known_workers = [Threads.@spawn worker_loop(i, cfg, known_queue, handle_known_job, st.enrich)
|
||||||
|
for i in 1:cfg.known_worker_count]
|
||||||
|
unknown_workers = [Threads.@spawn worker_loop(i, cfg, unknown_queue,
|
||||||
|
(job, c, wid) -> handle_unknown_job(job, c, wid, text_queue, st.triage),
|
||||||
|
st.triage)
|
||||||
|
for i in 1:cfg.unknown_worker_count]
|
||||||
|
text_workers = [Threads.@spawn worker_loop(i, cfg, text_queue,
|
||||||
|
(job, c, wid) -> handle_text_job(job, c, wid, DETECTOR[]),
|
||||||
|
st.language)
|
||||||
|
for i in 1:cfg.text_worker_count]
|
||||||
|
|
||||||
|
# Recovery runs *after* the pools are live, and blocks rather than dropping
|
||||||
|
# when the queue is full. `spool/` holds every in-flight file now, so a crash
|
||||||
|
# under load can leave far more leftovers than `queue_capacity`; with the
|
||||||
|
# consumers already running the queue drains as we fill it, so a big backlog
|
||||||
|
# takes longer to re-drive but none of it is abandoned. Serving hasn't
|
||||||
|
# started yet, so intake cannot race us for the directory.
|
||||||
|
#
|
||||||
|
# Everything re-enters at stage 1 and replays: in-process queues die with the
|
||||||
|
# process, so a crash loses the only record of how far each file had got.
|
||||||
|
# Safe, because every stage is a pure function of the file's bytes and every
|
||||||
|
# commit is idempotent; see the header of src/worker.jl.
|
||||||
|
#
|
||||||
|
# Broker-backed, this is exactly the wrong thing to do: the broker still
|
||||||
|
# holds every unacked job at the stage it reached, so re-driving `spool/`
|
||||||
|
# would duplicate the entire in-flight backlog on every restart and drag
|
||||||
|
# finished stages back to the start. `FS_RECOVER_SPOOL` forces it anyway, for
|
||||||
|
# the one case the broker cannot cover: it was purged or recreated, and the
|
||||||
|
# spooled files are the only surviving record of the work.
|
||||||
|
if cfg.queue_backend === :channel || cfg.recover_spool
|
||||||
|
recovered = recover_dir!(cfg.spool_dir, queue)
|
||||||
|
recovered > 0 && @info "recovered leftover files; replaying from stage 1" recovered=recovered
|
||||||
|
else
|
||||||
|
in_flight = count(p -> isfile(p) && !endswith(p, ".meta.json"),
|
||||||
|
readdir(cfg.spool_dir; join = true))
|
||||||
|
@info "broker holds the in-flight record; not re-driving spool/ (set FS_RECOVER_SPOOL=true if the broker was purged)" spooled=in_flight
|
||||||
|
end
|
||||||
|
|
||||||
|
register_routes()
|
||||||
|
# `handler` replaces Oxygen's root stream handler so POST /upload can read its
|
||||||
|
# body incrementally instead of having it buffered into memory first; every
|
||||||
|
# other route still goes through Oxygen (see `root_stream_handler`).
|
||||||
|
serve(; host = cfg.host, port = cfg.port, async = true, show_banner = false,
|
||||||
|
handler = root_stream_handler)
|
||||||
|
|
||||||
|
# Idempotent graceful drain: stop accepting uploads, let workers finish the
|
||||||
|
# buffered jobs, then exit. Called from two places:
|
||||||
|
# * the `finally` below, for SIGINT (Ctrl-C) and normal return, and
|
||||||
|
# * an `atexit` hook, for SIGTERM (systemd/Docker/k8s `stop`).
|
||||||
|
# We can't intercept SIGTERM directly: Julia blocks it on worker threads and
|
||||||
|
# handles it in its own runtime, so a user signal() handler never fires. But
|
||||||
|
# Julia's SIGTERM path runs `atexit` hooks, which gives us a reliable seam.
|
||||||
|
drained = Threads.Atomic{Bool}(false)
|
||||||
|
function drain()
|
||||||
|
Threads.atomic_xchg!(drained, true) && return # run at most once
|
||||||
|
@info "draining queues and stopping workers"
|
||||||
|
terminate() # 1. stop accepting new HTTP requests
|
||||||
|
close!(queue) # 2. no new classify jobs; stage-1 drains buffered
|
||||||
|
foreach(wait, workers) # 3. wait out stage-1, the ONLY producer of BOTH the
|
||||||
|
# known and unknown queues, so nothing else enqueues
|
||||||
|
close!(known_queue) # 4. now safe to close the queues stage-1 fed
|
||||||
|
close!(unknown_queue)
|
||||||
|
foreach(wait, known_workers) # 5. wait out stage-2 (terminal) and stage-3; stage-3 is
|
||||||
|
foreach(wait, unknown_workers) # the ONLY producer of the text queue
|
||||||
|
close!(text_queue) # 6. now safe to close the queue stage-3 fed
|
||||||
|
foreach(wait, text_workers) # 7. wait out stage-4
|
||||||
|
backend === nothing || close_backend!(backend) # 8. release the broker connection
|
||||||
|
@info "shutdown complete"
|
||||||
|
end
|
||||||
|
atexit(drain)
|
||||||
|
|
||||||
|
# Two ways out of here: an interrupt, or the broker connection dying.
|
||||||
|
#
|
||||||
|
# SIGINT is delivered to whichever task happens to be running on thread 1,
|
||||||
|
# which is the main task in the common case but is not guaranteed to be:
|
||||||
|
# AMQPClient runs its own receiver tasks, and an interrupt that lands in one
|
||||||
|
# of those kills the connection instead of reaching this loop. That is not a
|
||||||
|
# hole so much as the same door: the poller sees the dead connection within
|
||||||
|
# a second and sets `connection_lost`, so the drain happens either way. Under
|
||||||
|
# a process manager, prefer SIGTERM for the broker backend — it goes through
|
||||||
|
# Julia's `atexit`, which has no such lottery (see "Shutdown" in the README).
|
||||||
|
try
|
||||||
|
while !connection_lost[]
|
||||||
|
sleep(0.5) # interruptible; SIGINT throws in here
|
||||||
|
end
|
||||||
|
@error "stopping: broker connection closed"
|
||||||
|
catch e
|
||||||
|
e isa InterruptException || rethrow(e)
|
||||||
|
@info "shutdown requested (SIGINT)"
|
||||||
|
finally
|
||||||
|
drain()
|
||||||
|
end
|
||||||
|
|
||||||
|
# A lost broker is a failure, and saying so in the exit status is what lets a
|
||||||
|
# supervisor tell "asked to stop" from "fell over".
|
||||||
|
connection_lost[] && exit(1)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
end # module
|
||||||
@@ -1,124 +0,0 @@
|
|||||||
module FileServer
|
|
||||||
|
|
||||||
using Logging
|
|
||||||
using UUIDs
|
|
||||||
using HTTP
|
|
||||||
using JSON3
|
|
||||||
using Oxygen
|
|
||||||
using Lux
|
|
||||||
using JLD2
|
|
||||||
|
|
||||||
include("config.jl")
|
|
||||||
include("job.jl")
|
|
||||||
include("queue.jl")
|
|
||||||
include("spool.jl")
|
|
||||||
include("model.jl") # build_model() + read_features(); shared with bin/train.jl
|
|
||||||
include("classify.jl") # Classifier + load_classifier/classify (needs model.jl)
|
|
||||||
include("worker.jl")
|
|
||||||
|
|
||||||
# Globals the HTTP handlers read at request time. Set once in `run`, before the
|
|
||||||
# server starts accepting connections. Declared after the includes above so the
|
|
||||||
# `Config`/`ChannelQueue` types exist.
|
|
||||||
const CONFIG = Ref{Config}()
|
|
||||||
const QUEUE = Ref{ChannelQueue}()
|
|
||||||
const CLASSIFIER = Ref{Classifier}() # loaded once at startup, shared read-only across workers
|
|
||||||
|
|
||||||
include("server.jl") # registers routes (references CONFIG/QUEUE at call time)
|
|
||||||
|
|
||||||
export run
|
|
||||||
|
|
||||||
# When stderr is a live terminal Julia flushes each write, but when it's
|
|
||||||
# redirected to a file or pipe (a log file, `tee`, journald, a container log
|
|
||||||
# driver) Julia block-buffers it — so a long-running server's logs sit unseen in
|
|
||||||
# the buffer until it fills or the process exits, making it look like nothing is
|
|
||||||
# happening. This wrapper delegates to a normal logger and flushes after every
|
|
||||||
# message so output appears immediately wherever stderr is pointed.
|
|
||||||
struct FlushLogger{L<:AbstractLogger} <: AbstractLogger
|
|
||||||
inner::L
|
|
||||||
end
|
|
||||||
Logging.min_enabled_level(f::FlushLogger) = Logging.min_enabled_level(f.inner)
|
|
||||||
Logging.shouldlog(f::FlushLogger, args...) = Logging.shouldlog(f.inner, args...)
|
|
||||||
Logging.catch_exceptions(f::FlushLogger) = Logging.catch_exceptions(f.inner)
|
|
||||||
function Logging.handle_message(f::FlushLogger, args...; kwargs...)
|
|
||||||
Logging.handle_message(f.inner, args...; kwargs...)
|
|
||||||
flush(f.inner.stream)
|
|
||||||
return nothing
|
|
||||||
end
|
|
||||||
|
|
||||||
"""
|
|
||||||
run(; overrides...)
|
|
||||||
|
|
||||||
Start the file server: build config, ensure directories, recover any leftover
|
|
||||||
spooled files, spawn the worker pool, then serve HTTP until interrupted
|
|
||||||
(Ctrl-C / SIGINT or SIGTERM). On shutdown it stops accepting uploads, drains the
|
|
||||||
queue, waits for workers to finish in-flight files, and exits cleanly.
|
|
||||||
|
|
||||||
Keyword `overrides` (e.g. `port=9000`) take precedence over environment vars.
|
|
||||||
"""
|
|
||||||
function run(; overrides...)
|
|
||||||
# By default a Julia script exits immediately on SIGINT (Ctrl-C). Disable
|
|
||||||
# that so the interrupt surfaces as a catchable InterruptException, letting
|
|
||||||
# us drain the queue gracefully below.
|
|
||||||
Base.exit_on_sigint(false)
|
|
||||||
|
|
||||||
# Flush every log message so output is visible in real time even when stderr
|
|
||||||
# is redirected to a file/pipe (see FlushLogger). Set before anything logs.
|
|
||||||
global_logger(FlushLogger(ConsoleLogger(stderr)))
|
|
||||||
|
|
||||||
cfg = config_from_env(; overrides...)
|
|
||||||
ensure_dirs(cfg)
|
|
||||||
|
|
||||||
if cfg.worker_count > Threads.nthreads()
|
|
||||||
@warn "worker_count exceeds available threads; workers will share threads (start Julia with -t N for real parallelism)" worker_count=cfg.worker_count nthreads=Threads.nthreads()
|
|
||||||
end
|
|
||||||
|
|
||||||
queue = ChannelQueue(cfg.queue_capacity)
|
|
||||||
CONFIG[] = cfg
|
|
||||||
QUEUE[] = queue
|
|
||||||
|
|
||||||
# Load the classifier before serving. Fail fast: a server that silently
|
|
||||||
# doesn't classify is a worse surprise than a clear startup error.
|
|
||||||
CLASSIFIER[] = load_classifier(cfg.model_path)
|
|
||||||
@info "loaded classifier" path=cfg.model_path
|
|
||||||
|
|
||||||
recovered = recover_spool!(cfg, queue)
|
|
||||||
@info "starting file-server" host=cfg.host port=cfg.port workers=cfg.worker_count capacity=cfg.queue_capacity recovered=recovered
|
|
||||||
|
|
||||||
workers = [Threads.@spawn worker_loop(i, cfg, queue) for i in 1:cfg.worker_count]
|
|
||||||
|
|
||||||
register_routes()
|
|
||||||
serve(; host = cfg.host, port = cfg.port, async = true, show_banner = false)
|
|
||||||
|
|
||||||
# Idempotent graceful drain: stop accepting uploads, let workers finish the
|
|
||||||
# buffered jobs, then exit. Called from two places:
|
|
||||||
# * the `finally` below, for SIGINT (Ctrl-C) and normal return, and
|
|
||||||
# * an `atexit` hook, for SIGTERM (systemd/Docker/k8s `stop`).
|
|
||||||
# We can't intercept SIGTERM directly — Julia blocks it on worker threads and
|
|
||||||
# handles it in its own runtime, so a user signal() handler never fires. But
|
|
||||||
# Julia's SIGTERM path runs `atexit` hooks, which gives us a reliable seam.
|
|
||||||
drained = Threads.Atomic{Bool}(false)
|
|
||||||
function drain()
|
|
||||||
Threads.atomic_xchg!(drained, true) && return # run at most once
|
|
||||||
@info "draining queue and stopping workers"
|
|
||||||
terminate() # stop accepting new HTTP requests
|
|
||||||
close!(queue) # let workers drain buffered jobs, then exit
|
|
||||||
foreach(wait, workers)
|
|
||||||
@info "shutdown complete"
|
|
||||||
end
|
|
||||||
atexit(drain)
|
|
||||||
|
|
||||||
try
|
|
||||||
while true
|
|
||||||
sleep(0.5) # interruptible; SIGINT throws in here
|
|
||||||
end
|
|
||||||
catch e
|
|
||||||
e isa InterruptException || rethrow(e)
|
|
||||||
@info "shutdown requested (SIGINT)"
|
|
||||||
finally
|
|
||||||
drain()
|
|
||||||
end
|
|
||||||
|
|
||||||
return nothing
|
|
||||||
end
|
|
||||||
|
|
||||||
end # module
|
|
||||||
404
src/catalog.jl
Normal file
404
src/catalog.jl
Normal file
@@ -0,0 +1,404 @@
|
|||||||
|
# Stage-5 phase B: the live, single-owner format catalog (DESIGN §5B/§9).
|
||||||
|
#
|
||||||
|
# `cluster.jl` is the *science*: feature extraction, the collapsed Gibbs sampler
|
||||||
|
# (phase A, offline), and `assign_file` (the phase-B scoring core). This file is
|
||||||
|
# the *plumbing* that turns that core into a durable, growing catalog:
|
||||||
|
#
|
||||||
|
# * a `Catalog` = the surviving clusters' sufficient statistics + a record of
|
||||||
|
# which files have already been folded in + a few example filenames each;
|
||||||
|
# * `catalog_sweep!` = the phase-B loop. For every *new* file in `binary/`,
|
||||||
|
# run the deterministic CRP-predictive `assign_file` and fold it into the
|
||||||
|
# chosen cluster's stats (DESIGN §5B);
|
||||||
|
# * `compact!` = the offline Gibbs pass that *seeds* the catalog on first run
|
||||||
|
# and periodically re-clusters the pile (DESIGN §5, "seed the initial
|
||||||
|
# catalog" / "periodic compaction");
|
||||||
|
# * `write_nominations!` = surface promotable clusters to a human (DESIGN §6).
|
||||||
|
#
|
||||||
|
# Concurrency model is the deliberate opposite of the stateless classify workers
|
||||||
|
# (DESIGN §9): exactly ONE process owns the catalog, so there are no locks and no
|
||||||
|
# torn reads of sufficient stats. The catalog is a single durable file mutated by
|
||||||
|
# that one process; it is committed with the same sidecar-first temp→fsync→rename
|
||||||
|
# →fsync-dir discipline as the stage-2 sidecars (`commit_enriched!`), so a crash
|
||||||
|
# mid-write can neither corrupt it nor lose the rename. This file lives in the
|
||||||
|
# module (not dependency-flat like cluster.jl) because it needs JSON3 + the fsync
|
||||||
|
# helpers from metadata.jl.
|
||||||
|
|
||||||
|
"How many example filenames to retain per cluster (for the human promotion gate,
|
||||||
|
DESIGN §6). A handful is plenty to eyeball; the count/signature carry the weight."
|
||||||
|
const CATALOG_EXAMPLE_CAP = 8
|
||||||
|
|
||||||
|
"""
|
||||||
|
Catalog
|
||||||
|
|
||||||
|
The mutable phase-B state owned by the single sweep process:
|
||||||
|
|
||||||
|
* `n` : header window these clusters were built at (must match the
|
||||||
|
feature window used to score new files; frozen once seeded).
|
||||||
|
* `clusters` : id → `ClusterStats` (per-position 257-counts + member count).
|
||||||
|
Ids are frozen at birth and never renumbered, so there is no
|
||||||
|
label switching across sweeps (DESIGN §5).
|
||||||
|
* `examples` : id → up to `CATALOG_EXAMPLE_CAP` member filenames, for the
|
||||||
|
human nomination glance.
|
||||||
|
* `next_id` : the next fresh cluster id to hand out (monotone; retired ids are
|
||||||
|
never reused, keeping ids globally unique over the catalog's life).
|
||||||
|
* `processed` : basenames of every `binary/` file already folded in, so a sweep
|
||||||
|
is incremental: it touches only files it has not seen. This set
|
||||||
|
grows with the `binary/` pile it mirrors, at the same rate and
|
||||||
|
no faster, which is acceptable for v1 (DESIGN §9).
|
||||||
|
"""
|
||||||
|
mutable struct Catalog
|
||||||
|
n::Int
|
||||||
|
clusters::Dict{Int,ClusterStats}
|
||||||
|
examples::Dict{Int,Vector{String}}
|
||||||
|
next_id::Int
|
||||||
|
processed::Set{String}
|
||||||
|
end
|
||||||
|
|
||||||
|
"An empty catalog at header window `n` (no clusters seen yet)."
|
||||||
|
Catalog(n::Integer) = Catalog(Int(n), Dict{Int,ClusterStats}(),
|
||||||
|
Dict{Int,Vector{String}}(), 1, Set{String}())
|
||||||
|
|
||||||
|
"Record `name` as an example of cluster `id`, capped at `CATALOG_EXAMPLE_CAP`."
|
||||||
|
function record_example!(cat::Catalog, id::Integer, name::AbstractString)
|
||||||
|
ex = get!(cat.examples, id, String[])
|
||||||
|
length(ex) < CATALOG_EXAMPLE_CAP && push!(ex, String(name))
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Durable persistence (reuse the stage-2 sidecar-first commit discipline)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
# Count tables are stored sparsely, as only the non-zero
|
||||||
|
# `(position, symbol, count)` triples, because a cluster's n×257 table is
|
||||||
|
# overwhelmingly zero (a peaked
|
||||||
|
# magic byte touches one of 257 symbols per position). Sparse keeps the catalog
|
||||||
|
# file small and its load O(non-zeros), not O(n·257·K).
|
||||||
|
function _sparse_counts(counts::Matrix{Int})
|
||||||
|
n, A = size(counts)
|
||||||
|
triples = Vector{Vector{Int}}()
|
||||||
|
@inbounds for i in 1:n, v in 1:A
|
||||||
|
c = counts[i, v]
|
||||||
|
c != 0 && push!(triples, [i, v, c])
|
||||||
|
end
|
||||||
|
return triples
|
||||||
|
end
|
||||||
|
|
||||||
|
function _dense_counts(triples, n::Integer)
|
||||||
|
counts = zeros(Int, n, ALPHABET)
|
||||||
|
for t in triples
|
||||||
|
counts[Int(t[1]), Int(t[2])] = Int(t[3])
|
||||||
|
end
|
||||||
|
return counts
|
||||||
|
end
|
||||||
|
|
||||||
|
"Serialize a `Catalog` to a plain `NamedTuple` ready for `JSON3.write`."
|
||||||
|
function catalog_payload(cat::Catalog)
|
||||||
|
clusters = [(
|
||||||
|
id = id,
|
||||||
|
members = c.members,
|
||||||
|
counts = _sparse_counts(c.counts),
|
||||||
|
examples = get(cat.examples, id, String[]),
|
||||||
|
) for (id, c) in sort(collect(cat.clusters); by=first)]
|
||||||
|
return (
|
||||||
|
n = cat.n,
|
||||||
|
next_id = cat.next_id,
|
||||||
|
clusters = clusters,
|
||||||
|
processed = sort(collect(cat.processed)),
|
||||||
|
)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
save_catalog!(path, cat)
|
||||||
|
|
||||||
|
Durably write `cat` to `path` with the sidecar-first ordering (DESIGN §9): write
|
||||||
|
to a temp name, fsync the bytes, atomically rename into place, then fsync the
|
||||||
|
containing directory so the rename itself survives power loss. A crash can leave
|
||||||
|
at most a stale `.tmp`, never a torn catalog.
|
||||||
|
"""
|
||||||
|
function save_catalog!(path::AbstractString, cat::Catalog)
|
||||||
|
dir = dirname(path)
|
||||||
|
isempty(dir) || mkpath(dir)
|
||||||
|
tmp = string(path, ".tmp")
|
||||||
|
open(tmp, "w") do io
|
||||||
|
write(io, JSON3.write(catalog_payload(cat)))
|
||||||
|
flush(io)
|
||||||
|
fsync_fd(fd(io)) # persist bytes before the rename
|
||||||
|
end
|
||||||
|
mv(tmp, path; force=true) # atomic replace
|
||||||
|
fsync_dir(isempty(dir) ? "." : dir) # persist the rename itself
|
||||||
|
return path
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
load_catalog(path; n) -> Catalog
|
||||||
|
|
||||||
|
Load the durable catalog from `path`, or return a fresh empty `Catalog(n)` if it
|
||||||
|
does not exist yet (first run). `n` is the configured header window used only for
|
||||||
|
the empty case; a loaded catalog keeps its own frozen `n`.
|
||||||
|
"""
|
||||||
|
function load_catalog(path::AbstractString; n::Integer=HEADER_N)
|
||||||
|
isfile(path) || return Catalog(n)
|
||||||
|
doc = JSON3.read(read(path, String))
|
||||||
|
cn = Int(doc.n)
|
||||||
|
cat = Catalog(cn)
|
||||||
|
cat.next_id = Int(doc.next_id)
|
||||||
|
for entry in doc.clusters
|
||||||
|
id = Int(entry.id)
|
||||||
|
c = ClusterStats(_dense_counts(entry.counts, cn), Int(entry.members))
|
||||||
|
cat.clusters[id] = c
|
||||||
|
cat.examples[id] = String[String(e) for e in entry.examples]
|
||||||
|
end
|
||||||
|
for name in doc.processed
|
||||||
|
push!(cat.processed, String(name))
|
||||||
|
end
|
||||||
|
return cat
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Listing the input pile
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
binary_files(dir) -> Vector{String}
|
||||||
|
|
||||||
|
Sorted absolute paths of the regular files in `dir` to be swept, skipping
|
||||||
|
`.meta.json` sidecars and any `.tmp` scratch. Sorted so a sweep's sequential
|
||||||
|
CRP-predictive assignment (which is order-dependent) is deterministic run to run.
|
||||||
|
"""
|
||||||
|
function binary_files(dir::AbstractString)
|
||||||
|
isdir(dir) || return String[]
|
||||||
|
paths = String[]
|
||||||
|
for name in readdir(dir; join=true)
|
||||||
|
isfile(name) || continue
|
||||||
|
(endswith(name, ".meta.json") || endswith(name, ".tmp")) && continue
|
||||||
|
push!(paths, name)
|
||||||
|
end
|
||||||
|
sort!(paths)
|
||||||
|
return paths
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Phase B: the incremental sweep (the deliverable)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
catalog_sweep!(cat, cfg) -> NamedTuple
|
||||||
|
|
||||||
|
Fold every *new* file in `cfg.cluster_dir` into `cat` using the deterministic
|
||||||
|
CRP-predictive rule (`assign_file`, DESIGN §5B), updating the chosen cluster's
|
||||||
|
sufficient statistics in place. Files already in `cat.processed` are skipped, so
|
||||||
|
repeated sweeps are incremental and idempotent over the pile.
|
||||||
|
|
||||||
|
Per file, `assign_file` returns the argmax component:
|
||||||
|
* an existing cluster id → the file joins it (`add!`);
|
||||||
|
* `0` (background) → counted, not clustered. A novel-but-unmatched file parks
|
||||||
|
here by design; genuinely new formats are discovered by the offline
|
||||||
|
`compact!` re-clustering the residue, not by single-file minting;
|
||||||
|
* `-1` (mint) → a fresh cluster is seeded with a frozen `next_id`. Under the
|
||||||
|
calibrated `bg_mass > α` this never fires on the live path (fresh and
|
||||||
|
background share the one-file likelihood, so background always wins); the
|
||||||
|
branch exists for correctness, not as a routine outcome (DESIGN §5B).
|
||||||
|
|
||||||
|
Mutates `cat` but does NOT persist it; the caller commits once, after the sweep.
|
||||||
|
Returns a summary of what happened this sweep.
|
||||||
|
"""
|
||||||
|
function catalog_sweep!(cat::Catalog, cfg::Config)
|
||||||
|
n_seen = 0; n_bg = 0; n_joined = 0; n_minted = 0
|
||||||
|
for path in binary_files(cfg.cluster_dir)
|
||||||
|
base = basename(path)
|
||||||
|
base in cat.processed && continue
|
||||||
|
x = header_symbols(path; n=cat.n)
|
||||||
|
ids = sort!(collect(keys(cat.clusters)))
|
||||||
|
k = assign_file(x, cat.clusters, ids;
|
||||||
|
α=cfg.cluster_alpha, β=cfg.cluster_pseudocount,
|
||||||
|
bg_mass=cfg.cluster_bg_mass)
|
||||||
|
if k == -1
|
||||||
|
id = cat.next_id
|
||||||
|
cat.next_id += 1
|
||||||
|
c = ClusterStats(cat.n)
|
||||||
|
add!(c, x)
|
||||||
|
cat.clusters[id] = c
|
||||||
|
record_example!(cat, id, base)
|
||||||
|
n_minted += 1
|
||||||
|
elseif k == 0
|
||||||
|
n_bg += 1
|
||||||
|
else
|
||||||
|
add!(cat.clusters[k], x)
|
||||||
|
record_example!(cat, k, base)
|
||||||
|
n_joined += 1
|
||||||
|
end
|
||||||
|
push!(cat.processed, base)
|
||||||
|
n_seen += 1
|
||||||
|
end
|
||||||
|
return (; n_seen, n_joined, n_bg, n_minted,
|
||||||
|
n_clusters=length(cat.clusters), n_processed=length(cat.processed))
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Offline seed / compaction (wraps the phase-A Gibbs sampler)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
compact!(cat, cfg; sweeps, restarts, rng) -> NamedTuple
|
||||||
|
|
||||||
|
Re-cluster the *entire* `binary/` pile with the offline collapsed Gibbs sampler
|
||||||
|
and adopt the winning partition as the catalog's clusters (DESIGN §5). This is
|
||||||
|
both the **seed** on first run (an empty catalog has no clusters, so the live
|
||||||
|
sweep alone would send everything to background) and the **periodic compaction**
|
||||||
|
that merges drifted clusters / splits bloated ones later.
|
||||||
|
|
||||||
|
Ids are taken from the Gibbs partition and frozen; because compaction re-derives
|
||||||
|
the whole partition, this is a wholesale replace of `clusters`/`examples`, and
|
||||||
|
every file in the pile is marked processed. Callers run this on an explicit
|
||||||
|
schedule (e.g. `--compact`), never on the latency path.
|
||||||
|
"""
|
||||||
|
function compact!(cat::Catalog, cfg::Config;
|
||||||
|
sweeps::Integer=150, restarts::Integer=6,
|
||||||
|
rng::AbstractRNG=Random.default_rng())
|
||||||
|
paths = binary_files(cfg.cluster_dir)
|
||||||
|
if isempty(paths)
|
||||||
|
return (; n_files=0, n_clusters=length(cat.clusters), n_bg=0)
|
||||||
|
end
|
||||||
|
X = header_matrix(paths; n=cat.n)
|
||||||
|
result = gibbs_cluster(X; α=cfg.cluster_alpha, β=cfg.cluster_pseudocount,
|
||||||
|
bg_mass=cfg.cluster_bg_mass, sweeps=sweeps,
|
||||||
|
restarts=restarts, rng=rng)
|
||||||
|
empty!(cat.clusters)
|
||||||
|
empty!(cat.examples)
|
||||||
|
empty!(cat.processed)
|
||||||
|
for (id, c) in result.clusters
|
||||||
|
cat.clusters[id] = c
|
||||||
|
end
|
||||||
|
cat.next_id = (isempty(result.clusters) ? 0 : maximum(keys(result.clusters))) + 1
|
||||||
|
for (j, path) in enumerate(paths)
|
||||||
|
base = basename(path)
|
||||||
|
push!(cat.processed, base)
|
||||||
|
z = result.assignments[j]
|
||||||
|
z > 0 && record_example!(cat, z, base)
|
||||||
|
end
|
||||||
|
n_bg = count(==(0), result.assignments)
|
||||||
|
return (; n_files=length(paths), n_clusters=length(cat.clusters), n_bg)
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Nominations (closing the loop to the classifier; DESIGN §6)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"Render a signature (from `signature`) into a human-readable hex template:
|
||||||
|
two hex digits for a required byte, `EOF` for a required past-EOF, `??` for a
|
||||||
|
wildcard. This is what a human eyeballs at the promotion gate."
|
||||||
|
function signature_hex(sig::AbstractVector)
|
||||||
|
parts = map(sig) do s
|
||||||
|
s === nothing ? "??" :
|
||||||
|
s == PAST_EOF ? "EOF" :
|
||||||
|
string(s; base=16, pad=2)
|
||||||
|
end
|
||||||
|
return join(parts, " ")
|
||||||
|
end
|
||||||
|
|
||||||
|
"The required (non-wildcard) positions of a signature as `(position, byte)`
|
||||||
|
records; `byte` is the raw 0–255 value, or the string `\"past_eof\"`."
|
||||||
|
function signature_magic(sig::AbstractVector)
|
||||||
|
magic = Vector{NamedTuple{(:position, :byte),Tuple{Int,Any}}}()
|
||||||
|
for (i, s) in enumerate(sig)
|
||||||
|
s === nothing && continue
|
||||||
|
push!(magic, (position=i, byte=(s == PAST_EOF ? "past_eof" : s)))
|
||||||
|
end
|
||||||
|
return magic
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
write_nominations!(cat, cfg) -> Vector{String}
|
||||||
|
|
||||||
|
Write one JSON nomination per promotable cluster (`is_promotable`, DESIGN §6)
|
||||||
|
into `cfg.nominated_dir`, each carrying the cluster's signature (hex template +
|
||||||
|
required magic positions), member count, and example filenames: everything a
|
||||||
|
human needs to glance and promote. The background (id 0) is never a cluster here,
|
||||||
|
so it can never be nominated, by construction.
|
||||||
|
|
||||||
|
Nominations are rewritten every sweep (membership only grows), so each file is
|
||||||
|
durably replaced via the same temp→fsync→rename→fsync-dir commit as the catalog.
|
||||||
|
Returns the paths written. Non-promotable clusters are left alone; a cluster
|
||||||
|
that *was* nominated and later fell below threshold cannot happen (members only
|
||||||
|
grow), so there is nothing to retract.
|
||||||
|
"""
|
||||||
|
function write_nominations!(cat::Catalog, cfg::Config)
|
||||||
|
mkpath(cfg.nominated_dir)
|
||||||
|
written = String[]
|
||||||
|
for (id, c) in sort(collect(cat.clusters); by=first)
|
||||||
|
sig = signature(c) # default β; signature β is
|
||||||
|
# decoupled from clustering β (DESIGN §11.3a)
|
||||||
|
is_promotable(c, sig; min_members=cfg.promote_min_members,
|
||||||
|
min_magic=cfg.promote_min_magic) || continue
|
||||||
|
payload = (
|
||||||
|
cluster_id = id,
|
||||||
|
members = c.members,
|
||||||
|
magic_length = magic_positions(sig),
|
||||||
|
signature_hex = signature_hex(sig),
|
||||||
|
magic = signature_magic(sig),
|
||||||
|
examples = get(cat.examples, id, String[]),
|
||||||
|
)
|
||||||
|
dest = joinpath(cfg.nominated_dir, "cluster-$(id).json")
|
||||||
|
tmp = string(dest, ".tmp")
|
||||||
|
open(tmp, "w") do io
|
||||||
|
write(io, JSON3.write(payload))
|
||||||
|
flush(io)
|
||||||
|
fsync_fd(fd(io))
|
||||||
|
end
|
||||||
|
mv(tmp, dest; force=true)
|
||||||
|
push!(written, dest)
|
||||||
|
end
|
||||||
|
fsync_dir(cfg.nominated_dir)
|
||||||
|
return written
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Orchestration + CLI (the periodic single-owner process; DESIGN §9)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
run_cluster_sweep(cfg; compact=false, rng) -> NamedTuple
|
||||||
|
|
||||||
|
One end-to-end pass of the single-owner stage-5 process: load the durable
|
||||||
|
catalog, either `compact!` (offline Gibbs, used to seed on first run or to
|
||||||
|
recompact) or `catalog_sweep!` (the incremental live path), then durably persist
|
||||||
|
the catalog and (re)write nominations. This is the whole job the cron/periodic
|
||||||
|
runner performs; `bin/cluster_sweep.jl` is a thin shell around it.
|
||||||
|
|
||||||
|
`compact` is forced when the catalog is empty (no clusters AND nothing processed
|
||||||
|
yet): a first live sweep against no clusters would send every file to background,
|
||||||
|
so the catalog must be seeded by an offline Gibbs pass before it can assign.
|
||||||
|
"""
|
||||||
|
function run_cluster_sweep(cfg::Config; compact::Bool=false,
|
||||||
|
rng::AbstractRNG=Random.default_rng())
|
||||||
|
cat = load_catalog(cfg.cluster_catalog_path; n=cfg.cluster_n)
|
||||||
|
if cat.n != cfg.cluster_n
|
||||||
|
@warn "configured cluster_n differs from the catalog's frozen window; using the catalog's" catalog_n=cat.n configured_n=cfg.cluster_n
|
||||||
|
end
|
||||||
|
is_empty = isempty(cat.clusters) && isempty(cat.processed)
|
||||||
|
mode = (compact || is_empty) ? :compact : :sweep
|
||||||
|
summary = mode === :compact ? compact!(cat, cfg; rng=rng) : catalog_sweep!(cat, cfg)
|
||||||
|
save_catalog!(cfg.cluster_catalog_path, cat)
|
||||||
|
nominated = write_nominations!(cat, cfg)
|
||||||
|
return (; mode, summary, n_nominated=length(nominated), nominated,
|
||||||
|
n_clusters=length(cat.clusters), n_processed=length(cat.processed))
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
cluster_sweep_cli(args)
|
||||||
|
|
||||||
|
Entry point for `bin/cluster_sweep.jl`. Builds a `Config` from the environment,
|
||||||
|
runs one `run_cluster_sweep`, and logs a one-line summary. `--compact` forces the
|
||||||
|
offline Gibbs re-cluster (seed / periodic compaction) instead of the incremental
|
||||||
|
live sweep.
|
||||||
|
"""
|
||||||
|
function cluster_sweep_cli(args::AbstractVector{<:AbstractString}=String[])
|
||||||
|
compact = "--compact" in args
|
||||||
|
cfg = config_from_env()
|
||||||
|
ensure_dirs(cfg)
|
||||||
|
@info "stage-5 sweep starting" catalog=cfg.cluster_catalog_path input=cfg.cluster_dir compact=compact
|
||||||
|
r = run_cluster_sweep(cfg; compact=compact)
|
||||||
|
@info "stage-5 sweep complete" mode=r.mode clusters=r.n_clusters processed=r.n_processed nominated=r.n_nominated summary=r.summary
|
||||||
|
return r
|
||||||
|
end
|
||||||
@@ -8,7 +8,7 @@
|
|||||||
|
|
||||||
An in-memory, ready-to-run classifier: the (stateless) Lux `model` plus the
|
An in-memory, ready-to-run classifier: the (stateless) Lux `model` plus the
|
||||||
learned parameters `ps` and states `st` loaded from the artifact. Immutable and
|
learned parameters `ps` and states `st` loaded from the artifact. Immutable and
|
||||||
shared read-only across all worker threads — Lux inference is a pure function
|
shared read-only across all worker threads, since Lux inference is a pure function
|
||||||
over `ps`/`st`, so no locking is needed.
|
over `ps`/`st`, so no locking is needed.
|
||||||
"""
|
"""
|
||||||
struct Classifier{M,P,S}
|
struct Classifier{M,P,S}
|
||||||
@@ -21,7 +21,7 @@ end
|
|||||||
load_classifier(path) -> Classifier
|
load_classifier(path) -> Classifier
|
||||||
|
|
||||||
Rebuild the fixed architecture and load `ps`/`st` from the JLD2 artifact at
|
Rebuild the fixed architecture and load `ps`/`st` from the JLD2 artifact at
|
||||||
`path`. Throws if the file is missing or unreadable — the server fails fast at
|
`path`. Throws if the file is missing or unreadable, so the server fails fast at
|
||||||
startup rather than run silently without classification.
|
startup rather than run silently without classification.
|
||||||
"""
|
"""
|
||||||
function load_classifier(path::AbstractString)
|
function load_classifier(path::AbstractString)
|
||||||
|
|||||||
518
src/cluster.jl
Normal file
518
src/cluster.jl
Normal file
@@ -0,0 +1,518 @@
|
|||||||
|
# Stage-5: unknown-format discovery by Bayesian header clustering.
|
||||||
|
#
|
||||||
|
# See model/DESIGN_clustering.md for the full rationale. In brief: files that
|
||||||
|
# stage-3 sorted into `binary/` are the `:unknown` sink: genuinely
|
||||||
|
# unrecognized bytes. This stage clusters them by *file format* (not producer)
|
||||||
|
# using the first `HEADER_N` header bytes, modeled as a Dirichlet-process
|
||||||
|
# mixture of per-position categoricals over a 257-symbol alphabet
|
||||||
|
# (byte 0–255, plus symbol 257 = "past EOF"). Each cluster's signature is a
|
||||||
|
# magic-number template that can be promoted into the classifier's fast path.
|
||||||
|
#
|
||||||
|
# This file is deliberately dependency-light: everything below is base Julia
|
||||||
|
# (only `log`, no `SpecialFunctions`), so it drops into the existing module and
|
||||||
|
# the offline calibration script alike without touching the Manifest. The model
|
||||||
|
# is categorical on purpose. Do NOT reuse model.jl's [0,1] byte scaling here
|
||||||
|
# (that metric is meaningful for the Lux net and meaningless for header bytes,
|
||||||
|
# where 0x89 and 0x88 are not "close"; see DESIGN §4).
|
||||||
|
|
||||||
|
"Number of leading header bytes modeled per file (the feature window). DESIGN §4b."
|
||||||
|
const HEADER_N = 32
|
||||||
|
|
||||||
|
"Alphabet size: byte values 0–255 plus one extra symbol for 'past end of file'."
|
||||||
|
const ALPHABET = 257
|
||||||
|
|
||||||
|
"The 'past EOF' symbol (1-based index `ALPHABET`). A file shorter than a given
|
||||||
|
position emits this here: real, discriminative signal for fixed-length formats,
|
||||||
|
and it avoids colliding zero-padding with genuine 0x00 header bytes (DESIGN §4)."
|
||||||
|
const PAST_EOF = ALPHABET
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Feature extraction
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
header_symbols(path; n=HEADER_N) -> Vector{Int}
|
||||||
|
|
||||||
|
Read the first `n` bytes of the file at `path` and map them to a length-`n`
|
||||||
|
vector of 1-based categorical symbols: byte value `b` → `b + 1` (so `1..256`),
|
||||||
|
and every position at or beyond end-of-file → `PAST_EOF` (`257`). Never reads
|
||||||
|
more than `n` bytes, so memory stays flat regardless of file size.
|
||||||
|
"""
|
||||||
|
function header_symbols(path::AbstractString; n::Integer=HEADER_N)
|
||||||
|
syms = fill(PAST_EOF, n)
|
||||||
|
open(path, "r") do io
|
||||||
|
bytes = read(io, n)
|
||||||
|
@inbounds for i in eachindex(bytes)
|
||||||
|
syms[i] = Int(bytes[i]) + 1
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return syms
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
header_matrix(paths; n=HEADER_N) -> Matrix{Int}
|
||||||
|
|
||||||
|
Stack `header_symbols` for every path into an `n × length(paths)` matrix (one
|
||||||
|
column per file), the input layout the Gibbs sampler and predictive scorer both
|
||||||
|
consume.
|
||||||
|
"""
|
||||||
|
function header_matrix(paths::AbstractVector{<:AbstractString}; n::Integer=HEADER_N)
|
||||||
|
X = Matrix{Int}(undef, n, length(paths))
|
||||||
|
for (j, p) in enumerate(paths)
|
||||||
|
X[:, j] = header_symbols(p; n=n)
|
||||||
|
end
|
||||||
|
return X
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Model: DP mixture of per-position categoricals (Dirichlet-Categorical)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
ClusterStats
|
||||||
|
|
||||||
|
Sufficient statistics for one cluster: a per-position count table `counts`
|
||||||
|
(`n × ALPHABET`; `counts[i, v]` = how many member files show symbol `v` at
|
||||||
|
position `i`) and the member count `members`. These are exactly what phase-B
|
||||||
|
persists per catalog entry, and everything the collapsed predictive needs.
|
||||||
|
A slot with `members == 0` is inactive (reusable): the Gibbs sweep prunes
|
||||||
|
emptied clusters without renumbering, so surviving cluster ids stay stable.
|
||||||
|
"""
|
||||||
|
mutable struct ClusterStats
|
||||||
|
counts::Matrix{Int} # n × ALPHABET
|
||||||
|
members::Int
|
||||||
|
end
|
||||||
|
|
||||||
|
ClusterStats(n::Integer) = ClusterStats(zeros(Int, n, ALPHABET), 0)
|
||||||
|
|
||||||
|
"Add file `x` (a length-n symbol vector) into cluster `c`'s sufficient stats."
|
||||||
|
function add!(c::ClusterStats, x::AbstractVector{<:Integer})
|
||||||
|
@inbounds for i in eachindex(x)
|
||||||
|
c.counts[i, x[i]] += 1
|
||||||
|
end
|
||||||
|
c.members += 1
|
||||||
|
return c
|
||||||
|
end
|
||||||
|
|
||||||
|
"Remove file `x` from cluster `c`'s sufficient stats (inverse of `add!`)."
|
||||||
|
function remove!(c::ClusterStats, x::AbstractVector{<:Integer})
|
||||||
|
@inbounds for i in eachindex(x)
|
||||||
|
c.counts[i, x[i]] -= 1
|
||||||
|
end
|
||||||
|
c.members -= 1
|
||||||
|
return c
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
log_predictive(c, x, β) -> Float64
|
||||||
|
|
||||||
|
Log probability that file `x` was produced by cluster `c` under the collapsed
|
||||||
|
Dirichlet-Categorical predictive, given `c`'s current counts: at each position
|
||||||
|
`i`, `p(x_i | c) = (counts[i, x_i] + β) / (members + ALPHABET·β)`, summed in log
|
||||||
|
space over positions. Call with `c` NOT containing `x` (Gibbs excludes the point
|
||||||
|
being resampled), so an emptied cluster reduces to the uniform prior `1/ALPHABET`
|
||||||
|
per position, identical to a brand-new cluster, as it should be.
|
||||||
|
"""
|
||||||
|
function log_predictive(c::ClusterStats, x::AbstractVector{<:Integer}, β::Float64)
|
||||||
|
denom = log(c.members + ALPHABET * β)
|
||||||
|
s = 0.0
|
||||||
|
@inbounds for i in eachindex(x)
|
||||||
|
s += log(c.counts[i, x[i]] + β) - denom
|
||||||
|
end
|
||||||
|
return s
|
||||||
|
end
|
||||||
|
|
||||||
|
"Log likelihood of `x` under the fixed uniform component (each position uniform
|
||||||
|
over the 257 symbols): `n · log(1/ALPHABET)`. Used for both the never-adaptive
|
||||||
|
background 'junk drawer' and the prior predictive of a fresh cluster (DESIGN §4a)."
|
||||||
|
log_uniform(n::Integer) = -n * log(ALPHABET)
|
||||||
|
|
||||||
|
# Lanczos approximation to log Γ(x) for x > 0, so partition scoring (below) needs
|
||||||
|
# the Dirichlet-multinomial marginal's gamma terms without pulling in
|
||||||
|
# SpecialFunctions, keeping this stage dependency-flat (no Manifest churn).
|
||||||
|
# g = 7, standard coefficients; accurate to ~1e-14 over the range we use.
|
||||||
|
const _LANCZOS_G = 7
|
||||||
|
const _LANCZOS_C = (0.99999999999980993, 676.5203681218851, -1259.1392167224028,
|
||||||
|
771.32342877765313, -176.61502916214059, 12.507343278686905,
|
||||||
|
-0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7)
|
||||||
|
|
||||||
|
function loggamma(x::Float64)
|
||||||
|
x < 0.5 && return log(π / sin(π * x)) - loggamma(1.0 - x) # reflection
|
||||||
|
x -= 1.0
|
||||||
|
a = _LANCZOS_C[1]
|
||||||
|
t = x + _LANCZOS_G + 0.5
|
||||||
|
@inbounds for i in 1:_LANCZOS_G + 1
|
||||||
|
a += _LANCZOS_C[i + 1] / (x + i)
|
||||||
|
end
|
||||||
|
return 0.5 * log(2π) + (x + 0.5) * log(t) - t + log(a)
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Phase A: collapsed Gibbs sampler (offline; the science)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
GibbsResult
|
||||||
|
|
||||||
|
Output of `gibbs_cluster`: `assignments` (one per input file; `0` = absorbed by
|
||||||
|
the background junk drawer, positive ints = cluster id), the surviving
|
||||||
|
`clusters` keyed by id, and `score` (the partition's collapsed pseudo-likelihood,
|
||||||
|
used to rank restarts).
|
||||||
|
"""
|
||||||
|
struct GibbsResult
|
||||||
|
assignments::Vector{Int}
|
||||||
|
clusters::Dict{Int,ClusterStats}
|
||||||
|
score::Float64
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
gibbs_cluster(X; α, β, bg_mass, sweeps, restarts, rng) -> GibbsResult
|
||||||
|
|
||||||
|
Cluster the columns of `X` (an `n × N` header-symbol matrix) with a collapsed
|
||||||
|
Gibbs sampler over a CRP/Dirichlet-Categorical mixture plus a fixed uniform
|
||||||
|
background (DESIGN §5). Unknown *k* falls out of the CRP natively.
|
||||||
|
|
||||||
|
Per point, per sweep, the point is removed from its cluster and reassigned by
|
||||||
|
sampling from the CRP-predictive weights:
|
||||||
|
|
||||||
|
* existing cluster `k`: `members_k · exp(log_predictive)`
|
||||||
|
* background: `bg_mass · (1/ALPHABET)^n` (never adapts)
|
||||||
|
* a fresh cluster: `α · (1/ALPHABET)^n`
|
||||||
|
|
||||||
|
A uniform/high-entropy blob matches no structured cluster, and background vs.
|
||||||
|
fresh is then decided by `bg_mass` vs. `α`; with `bg_mass ≥ α` such blobs are
|
||||||
|
absorbed rather than minting singletons. `restarts` independent runs are made
|
||||||
|
from different seeds and the highest-scoring partition is returned (a cheap,
|
||||||
|
base-only stand-in for the offline Binder/VI point-summary the design defers).
|
||||||
|
"""
|
||||||
|
function gibbs_cluster(X::AbstractMatrix{<:Integer};
|
||||||
|
α::Float64=1.0, β::Float64=0.5, bg_mass::Float64=5.0,
|
||||||
|
sweeps::Integer=80, restarts::Integer=4,
|
||||||
|
rng::AbstractRNG=Random.default_rng())
|
||||||
|
best = nothing
|
||||||
|
for _ in 1:restarts
|
||||||
|
r = _gibbs_once(X; α=α, β=β, bg_mass=bg_mass, sweeps=sweeps, rng=rng)
|
||||||
|
if best === nothing || r.score > best.score
|
||||||
|
best = r
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return best
|
||||||
|
end
|
||||||
|
|
||||||
|
function _gibbs_once(X::AbstractMatrix{<:Integer};
|
||||||
|
α::Float64, β::Float64, bg_mass::Float64,
|
||||||
|
sweeps::Integer, rng::AbstractRNG)
|
||||||
|
n, N = size(X)
|
||||||
|
# Seed every file in its own singleton (NOT the background). Cold-starting
|
||||||
|
# from the background deadlocks: at a format's first file, a fresh cluster
|
||||||
|
# and the background are equally uniform, so with bg_mass ≥ α the background
|
||||||
|
# always wins and no real cluster is ever seeded. Singleton init sidesteps
|
||||||
|
# this: same-format singletons merge and snowball, while a lone
|
||||||
|
# random-blob singleton dissolves on resample and is reclaimed by the
|
||||||
|
# (stickier) background. See DESIGN §4a.
|
||||||
|
z = collect(1:N)
|
||||||
|
clusters = Dict{Int,ClusterStats}()
|
||||||
|
for j in 1:N
|
||||||
|
c = ClusterStats(n)
|
||||||
|
add!(c, view(X, :, j))
|
||||||
|
clusters[j] = c
|
||||||
|
end
|
||||||
|
next_id = N + 1
|
||||||
|
log_u = log_uniform(n)
|
||||||
|
log_bg = log(bg_mass) + log_u
|
||||||
|
log_new = log(α) + log_u
|
||||||
|
|
||||||
|
# Reused across every point-visit so the sampler's hot loop allocates nothing
|
||||||
|
# per step (2M+ visits per run): `idbuf[t]` is the cluster id whose weight is
|
||||||
|
# `logw[t+1]` (logw[1] = background, logw[end] = fresh). Rebuilding these with
|
||||||
|
# fresh `Vector`/`collect(keys(...))` each step was both slow and enough GC
|
||||||
|
# churn to trip a Julia GC segfault on long grid runs.
|
||||||
|
idbuf = Int[]
|
||||||
|
logw = Float64[]
|
||||||
|
for _ in 1:sweeps
|
||||||
|
for j in 1:N
|
||||||
|
x = view(X, :, j)
|
||||||
|
|
||||||
|
# Remove point j from its current component.
|
||||||
|
zj = z[j]
|
||||||
|
if zj > 0
|
||||||
|
c = clusters[zj]
|
||||||
|
remove!(c, x)
|
||||||
|
if c.members == 0
|
||||||
|
delete!(clusters, zj) # prune emptied cluster; id retired
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# Candidate log-weights: background, each live cluster, fresh.
|
||||||
|
empty!(idbuf); empty!(logw)
|
||||||
|
push!(logw, log_bg)
|
||||||
|
for (k, c) in clusters
|
||||||
|
push!(idbuf, k)
|
||||||
|
push!(logw, log(c.members) + log_predictive(c, x, β))
|
||||||
|
end
|
||||||
|
push!(logw, log_new)
|
||||||
|
|
||||||
|
# Gumbel-max sample from the categorical over components.
|
||||||
|
pick = _gumbel_argmax(logw, rng)
|
||||||
|
|
||||||
|
if pick == 1
|
||||||
|
z[j] = 0 # background
|
||||||
|
elseif pick == length(logw)
|
||||||
|
id = next_id; next_id += 1 # fresh cluster
|
||||||
|
c = ClusterStats(n)
|
||||||
|
add!(c, x)
|
||||||
|
clusters[id] = c
|
||||||
|
z[j] = id
|
||||||
|
else
|
||||||
|
id = idbuf[pick - 1]
|
||||||
|
add!(clusters[id], x)
|
||||||
|
z[j] = id
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
return GibbsResult(z, clusters, partition_logmarginal(X, z, clusters, α, β))
|
||||||
|
end
|
||||||
|
|
||||||
|
"Argmax of `logw .+ Gumbel noise`: an exact draw from softmax(logw) without
|
||||||
|
normalizing (numerically safe for the tiny header-likelihood magnitudes)."
|
||||||
|
function _gumbel_argmax(logw::AbstractVector{Float64}, rng::AbstractRNG)
|
||||||
|
best_i = 1
|
||||||
|
best_v = -Inf
|
||||||
|
@inbounds for i in eachindex(logw)
|
||||||
|
g = logw[i] - log(-log(rand(rng)))
|
||||||
|
if g > best_v
|
||||||
|
best_v = g
|
||||||
|
best_i = i
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return best_i
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
partition_logmarginal(X, z, clusters, α, β) -> Float64
|
||||||
|
|
||||||
|
The joint log-evidence `log p(z, X)` of a partition under the CRP prior and the
|
||||||
|
collapsed Dirichlet-Categorical likelihood. This is the principled score for ranking
|
||||||
|
Gibbs restarts (higher = better). It is the sum of:
|
||||||
|
|
||||||
|
* the Dirichlet-multinomial **marginal** of each cluster's per-position counts,
|
||||||
|
`lΓ(Aβ) − lΓ(mₖ+Aβ) + Σ_v [lΓ(c_v+β) − lΓ(β)]`, whose normalizer supplies the
|
||||||
|
Occam penalty that a plain product-of-predictives lacks. It is what makes a
|
||||||
|
*merged, heterogeneous* cluster score **worse** than two clean ones (an
|
||||||
|
earlier pseudo-likelihood scorer omitted this and wrongly rewarded merging);
|
||||||
|
* the CRP prior over the clustered points, `K·log α + Σₖ lΓ(mₖ) + lΓ(α) −
|
||||||
|
lΓ(α+N_clustered)`, penalizing gratuitous extra clusters; and
|
||||||
|
* the fixed uniform term for background-assigned files.
|
||||||
|
"""
|
||||||
|
function partition_logmarginal(X::AbstractMatrix{<:Integer}, z::AbstractVector{<:Integer},
|
||||||
|
clusters::Dict{Int,ClusterStats}, α::Float64, β::Float64)
|
||||||
|
n, N = size(X)
|
||||||
|
Aβ = ALPHABET * β
|
||||||
|
lg_Aβ = loggamma(Aβ)
|
||||||
|
lg_β = loggamma(β)
|
||||||
|
s = 0.0
|
||||||
|
# Dirichlet-Categorical marginal likelihood, per cluster × position.
|
||||||
|
for (_, c) in clusters
|
||||||
|
lg_denom = loggamma(c.members + Aβ)
|
||||||
|
@inbounds for i in 1:n, v in 1:ALPHABET
|
||||||
|
cv = c.counts[i, v]
|
||||||
|
cv > 0 && (s += loggamma(cv + β) - lg_β)
|
||||||
|
end
|
||||||
|
s += n * (lg_Aβ - lg_denom)
|
||||||
|
end
|
||||||
|
# CRP prior over the partition of the clustered points.
|
||||||
|
n_bg = count(==(0), z)
|
||||||
|
n_clustered = N - n_bg
|
||||||
|
K = length(clusters)
|
||||||
|
s += K * log(α) + loggamma(α) - loggamma(α + n_clustered)
|
||||||
|
for (_, c) in clusters
|
||||||
|
s += loggamma(float(c.members))
|
||||||
|
end
|
||||||
|
# Background files: independent, uniform.
|
||||||
|
s += n_bg * log_uniform(n)
|
||||||
|
return s
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Phase B: sequential CRP-predictive assignment (online; the catalog)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
assign_file(x, clusters, ids; α, β, bg_mass) -> Int
|
||||||
|
|
||||||
|
Deterministically assign a single file `x` against an existing catalog: the
|
||||||
|
same CRP-predictive rule as Gibbs but at the **argmax** (no sampling) with the
|
||||||
|
current assignments held fixed (DESIGN §5B). Returns the id of the chosen
|
||||||
|
cluster, `0` for the background, or `-1` to signal "mint a new cluster". `ids`
|
||||||
|
is the caller's stable ordering of `keys(clusters)`.
|
||||||
|
|
||||||
|
A new cluster is minted (`-1`) only when the fresh-cluster weight strictly wins.
|
||||||
|
Fresh and background share the same `(1/ALPHABET)^n` likelihood (one file, however
|
||||||
|
structured, is indistinguishable from a uniform blob until a *second* like file
|
||||||
|
appears), so this reduces to `α > bg_mass`. Under the calibrated `bg_mass > α`,
|
||||||
|
minting is therefore effectively off on the live path **by design**: a novel file
|
||||||
|
that matches nothing parks in the background, and genuinely new formats are
|
||||||
|
discovered by the periodic **offline Gibbs compaction** re-clustering that
|
||||||
|
residue (DESIGN §5), not by single-file minting. Because ids are frozen at birth
|
||||||
|
by the caller, there is no label switching.
|
||||||
|
"""
|
||||||
|
function assign_file(x::AbstractVector{<:Integer}, clusters::Dict{Int,ClusterStats},
|
||||||
|
ids::AbstractVector{<:Integer};
|
||||||
|
α::Float64=1.0, β::Float64=0.5, bg_mass::Float64=5.0)
|
||||||
|
n = length(x)
|
||||||
|
log_u = log_uniform(n)
|
||||||
|
best_kind = :bg # :bg, :existing, :new
|
||||||
|
best_id = 0
|
||||||
|
best = log(bg_mass) + log_u
|
||||||
|
new_w = log(α) + log_u
|
||||||
|
if new_w > best
|
||||||
|
best = new_w; best_kind = :new
|
||||||
|
end
|
||||||
|
for k in ids
|
||||||
|
c = clusters[k]
|
||||||
|
w = log(c.members) + log_predictive(c, x, β)
|
||||||
|
if w > best
|
||||||
|
best = w; best_kind = :existing; best_id = k
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return best_kind === :bg ? 0 : best_kind === :new ? -1 : best_id
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Signatures and promotion (closing the loop to the classifier; DESIGN §6)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
signature(c; peak_threshold=0.9, β=0.5) -> Vector{Union{Int,Nothing}}
|
||||||
|
|
||||||
|
Turn a cluster's counts into a libmagic-style template: at each position, if the
|
||||||
|
modal symbol's posterior probability exceeds `peak_threshold`, that byte is
|
||||||
|
*required* (returned as the 0–255 byte value, or `PAST_EOF`); otherwise the
|
||||||
|
position is a wildcard (`nothing`). The vector of required bytes IS the
|
||||||
|
magic-number template, the whole point of the categorical model (DESIGN §4).
|
||||||
|
|
||||||
|
Peakedness is a **Bernoulli** question ("is this position fixed to byte `v`, or
|
||||||
|
not?"), so it uses a 2-way posterior mean `(count + β)/(members + 2β)`, NOT the
|
||||||
|
257-way mixture predictive. The alphabet-wide version would smear the estimate
|
||||||
|
across 257 symbols (`members + 257β` in the denominator), pulling even a
|
||||||
|
unanimous position below any sane threshold once β is small, which would make
|
||||||
|
promotion impossible. This decouples signature detection from the clustering
|
||||||
|
pseudocount and the alphabet size.
|
||||||
|
"""
|
||||||
|
function signature(c::ClusterStats; peak_threshold::Float64=0.9, β::Float64=0.5)
|
||||||
|
n = size(c.counts, 1)
|
||||||
|
sig = Vector{Union{Int,Nothing}}(nothing, n)
|
||||||
|
c.members == 0 && return sig
|
||||||
|
denom = c.members + 2β
|
||||||
|
@inbounds for i in 1:n
|
||||||
|
v = argmax(view(c.counts, i, :))
|
||||||
|
p = (c.counts[i, v] + β) / denom
|
||||||
|
if p > peak_threshold
|
||||||
|
sig[i] = v == PAST_EOF ? PAST_EOF : v - 1 # back to raw byte value
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return sig
|
||||||
|
end
|
||||||
|
|
||||||
|
"Number of fixed (non-wildcard) positions in a signature: its 'magic length'."
|
||||||
|
magic_positions(sig::AbstractVector) = count(!isnothing, sig)
|
||||||
|
|
||||||
|
"""
|
||||||
|
is_promotable(c, sig; min_members=20, min_magic=3) -> Bool
|
||||||
|
|
||||||
|
A cluster qualifies for *nomination* (still human-gated, DESIGN §6) when it has
|
||||||
|
at least `min_members` files AND at least `min_magic` fixed signature positions.
|
||||||
|
The background (id 0) is never passed here; it is never promotable by design.
|
||||||
|
"""
|
||||||
|
function is_promotable(c::ClusterStats, sig::AbstractVector;
|
||||||
|
min_members::Integer=20, min_magic::Integer=3)
|
||||||
|
return c.members >= min_members && magic_positions(sig) >= min_magic
|
||||||
|
end
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Calibration metrics (DESIGN §7): agreement of recovered clusters vs. truth
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
"Map a label vector to consecutive integer ids and a group→indices table."
|
||||||
|
function _groups(labels::AbstractVector)
|
||||||
|
g = Dict{Any,Vector{Int}}()
|
||||||
|
for (i, l) in enumerate(labels)
|
||||||
|
push!(get!(g, l, Int[]), i)
|
||||||
|
end
|
||||||
|
return g
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
adjusted_rand_index(a, b) -> Float64
|
||||||
|
|
||||||
|
Adjusted Rand Index between two labelings of the same items: 1.0 = identical
|
||||||
|
partitions (up to relabeling), ~0.0 = chance agreement, can go negative. This is
|
||||||
|
the §7 calibration objective: grid-tuning maximizes ARI of recovered-vs-truth
|
||||||
|
(magic-collapsed) labels. Hand-rolled to keep the dependency footprint flat;
|
||||||
|
matches `Clustering.randindex`.
|
||||||
|
"""
|
||||||
|
function adjusted_rand_index(a::AbstractVector, b::AbstractVector)
|
||||||
|
length(a) == length(b) || throw(DimensionMismatch("label vectors differ in length"))
|
||||||
|
n = length(a)
|
||||||
|
n < 2 && return 1.0
|
||||||
|
ga = collect(values(_groups(a)))
|
||||||
|
gb = collect(values(_groups(b)))
|
||||||
|
# Contingency-table sum of C(n_ij, 2).
|
||||||
|
comb2(x) = x * (x - 1) / 2
|
||||||
|
sa = Set.(ga)
|
||||||
|
index = 0.0
|
||||||
|
for A in sa, B in gb
|
||||||
|
nij = count(in(A), B)
|
||||||
|
index += comb2(nij)
|
||||||
|
end
|
||||||
|
sum_a = sum(comb2(length(g)) for g in ga)
|
||||||
|
sum_b = sum(comb2(length(g)) for g in gb)
|
||||||
|
total = comb2(n)
|
||||||
|
expected = sum_a * sum_b / total
|
||||||
|
maxi = (sum_a + sum_b) / 2
|
||||||
|
denom = maxi - expected
|
||||||
|
return denom == 0 ? 1.0 : (index - expected) / denom
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
v_measure(truth, pred; β=1.0) -> (v, homogeneity, completeness)
|
||||||
|
|
||||||
|
Entropy-based cluster agreement (Rosenberg & Hirschberg): homogeneity (each
|
||||||
|
predicted cluster holds one true class), completeness (each true class stays in
|
||||||
|
one predicted cluster), and their weighted harmonic mean `v`. Reported alongside
|
||||||
|
ARI in §7 calibration as a second, differently-biased view.
|
||||||
|
"""
|
||||||
|
function v_measure(truth::AbstractVector, pred::AbstractVector; β::Float64=1.0)
|
||||||
|
n = length(truth)
|
||||||
|
n == 0 && return (1.0, 1.0, 1.0)
|
||||||
|
gt = _groups(truth)
|
||||||
|
gp = _groups(pred)
|
||||||
|
entropy(g) = -sum((length(v) / n) * log(length(v) / n) for v in values(g))
|
||||||
|
H_C = entropy(gt)
|
||||||
|
H_K = entropy(gp)
|
||||||
|
# Conditional entropies via the contingency table.
|
||||||
|
H_CK = 0.0 # H(truth | pred)
|
||||||
|
H_KC = 0.0 # H(pred | truth)
|
||||||
|
for (_, P) in gp
|
||||||
|
Ps = Set(P)
|
||||||
|
for (_, C) in gt
|
||||||
|
nij = count(in(Ps), C)
|
||||||
|
nij == 0 && continue
|
||||||
|
H_CK -= (nij / n) * log(nij / length(P))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
for (_, C) in gt
|
||||||
|
Cs = Set(C)
|
||||||
|
for (_, P) in gp
|
||||||
|
nij = count(in(Cs), P)
|
||||||
|
nij == 0 && continue
|
||||||
|
H_KC -= (nij / n) * log(nij / length(C))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
homogeneity = H_C == 0 ? 1.0 : 1 - H_CK / H_C
|
||||||
|
completeness = H_K == 0 ? 1.0 : 1 - H_KC / H_K
|
||||||
|
v = (homogeneity + completeness == 0) ? 0.0 :
|
||||||
|
(1 + β) * homogeneity * completeness / (β * homogeneity + completeness)
|
||||||
|
return (v, homogeneity, completeness)
|
||||||
|
end
|
||||||
156
src/config.jl
156
src/config.jl
@@ -7,10 +7,100 @@ Base.@kwdef struct Config
|
|||||||
port::Int = 8080
|
port::Int = 8080
|
||||||
worker_count::Int = Threads.nthreads()
|
worker_count::Int = Threads.nthreads()
|
||||||
queue_capacity::Int = 1000
|
queue_capacity::Int = 1000
|
||||||
spool_dir::String = "data/spool" # files land here on intake (pending)
|
# Stage 2 (enrichment) has its own pool + queue: exiftool work is process-spawn
|
||||||
done_dir::String = "data/done" # files move here after successful processing
|
# and I/O bound, a different cost profile than the CPU-bound Lux classify, so
|
||||||
|
# the two pools are tuned independently.
|
||||||
|
known_worker_count::Int = Threads.nthreads()
|
||||||
|
known_queue_capacity::Int = 1000
|
||||||
|
# Stage 3 (content triage) also gets its own pool + queue: sorting an
|
||||||
|
# unrecognized file into binary vs text is cheap I/O, tuned independently
|
||||||
|
# of the classify and enrich pools.
|
||||||
|
unknown_worker_count::Int = Threads.nthreads()
|
||||||
|
unknown_queue_capacity::Int = 1000
|
||||||
|
# Stage 4 (language enrichment) has its own pool + queue too: detecting a text
|
||||||
|
# file's natural language (Languages.jl) and programming language (shelling to
|
||||||
|
# github-linguist) is a mix of CPU and process-spawn work, tuned independently.
|
||||||
|
text_worker_count::Int = Threads.nthreads()
|
||||||
|
text_queue_capacity::Int = 1000
|
||||||
|
# Which `JobQueue` implementation backs the four stage queues (src/queue.jl):
|
||||||
|
# `:channel` (in-process, the default, nothing external to run) or
|
||||||
|
# `:rabbitmq` (durable broker queues, src/rabbit.jl). The queue *capacities*
|
||||||
|
# above apply to both, but mean different things: a hard limit in-process, an
|
||||||
|
# advisory one against a polled depth on the broker (see `RabbitQueue`).
|
||||||
|
queue_backend::Symbol = :channel
|
||||||
|
amqp_url::String = "amqp://guest:guest@localhost:5672/"
|
||||||
|
amqp_prefix::String = "darkstruct" # queue names are "<prefix>.<stage>"
|
||||||
|
# Unacked messages the broker will hand one stage at a time. Defaults to that
|
||||||
|
# stage's worker count (0 below means "use it"), so a worker holds at most the
|
||||||
|
# one it is working on and a restart redelivers the minimum.
|
||||||
|
amqp_prefetch::Int = 0
|
||||||
|
# Wait for a publisher confirm before intake reports a file accepted, so a
|
||||||
|
# 202 means the broker has it rather than only that our socket took it. Costs
|
||||||
|
# a round trip per file on the intake path *only*: inter-stage publishes stay
|
||||||
|
# fire-and-forget, because a lost one leaves its source message unacked and
|
||||||
|
# redelivery repairs it. See "Queue backends" in the README.
|
||||||
|
amqp_confirms::Bool = true
|
||||||
|
# Broker-backed startup does not re-drive `spool/`: the broker is the record
|
||||||
|
# of what is in flight, and re-driving would duplicate the whole backlog at
|
||||||
|
# stage 1 on every restart. Set this when the broker was purged or recreated
|
||||||
|
# and the spooled files are the only surviving record. Ignored (always on)
|
||||||
|
# for the in-process backend, whose queues never survive a restart.
|
||||||
|
recover_spool::Bool = false
|
||||||
|
# Every in-flight file lives here, at every stage, from intake until it is
|
||||||
|
# committed to a terminal sink below. Stages route by enqueueing a reference,
|
||||||
|
# not by moving bytes (src/worker.jl header), so there are no per-stage
|
||||||
|
# directories; the queues are what know how far a file has got.
|
||||||
|
spool_dir::String = "data/spool"
|
||||||
|
binary_dir::String = "data/binary" # stage-3 sink: unknown files that look like binary data
|
||||||
|
done_dir::String = "data/done" # fully enriched known files (+ .meta.json sidecars)
|
||||||
|
text_done_dir::String = "data/text_done" # fully enriched text files (+ .meta.json sidecars)
|
||||||
failed_dir::String = "data/failed" # files move here if a worker throws
|
failed_dir::String = "data/failed" # files move here if a worker throws
|
||||||
model_path::String = "model/classifier.jld2" # committed classifier artifact, loaded at startup
|
model_path::String = "model/classifier.jld2" # committed classifier artifact, loaded at startup
|
||||||
|
# Intake reads each upload off the socket in chunks of this size and streams
|
||||||
|
# them straight to the spool file, so this, not the file size, is what
|
||||||
|
# bounds intake memory per in-flight upload (see src/multipart.jl).
|
||||||
|
upload_chunk_bytes::Int = UPLOAD_CHUNK_BYTES
|
||||||
|
exiftool_timeout::Int = 30 # seconds before a stuck exiftool is killed → degraded sidecar
|
||||||
|
linguist_timeout::Int = 30 # seconds before a stuck github-linguist is killed → no programming language
|
||||||
|
# Stage 5 (unknown-format discovery). A separate single-owner batch process
|
||||||
|
# sweeps binary/ and clusters headers; these are its knobs (see
|
||||||
|
# model/DESIGN_clustering.md §9). Values are the calibrated defaults from
|
||||||
|
# bin/cluster_calibrate.jl on the training corpus.
|
||||||
|
cluster_dir::String = "data/binary" # stage-5 input: the :unknown/binary sink to sweep
|
||||||
|
cluster_n::Int = 32 # header bytes modeled per file (HEADER_N)
|
||||||
|
cluster_alpha::Float64 = 1.0 # CRP concentration: propensity to spawn new formats
|
||||||
|
cluster_pseudocount::Float64 = 0.1 # Dirichlet pseudocount β; calibrated on the training corpus
|
||||||
|
cluster_bg_mass::Float64 = 5.0 # fixed mass of the uniform background 'junk drawer'
|
||||||
|
promote_min_members::Int = 20 # cluster size threshold for promotion nomination
|
||||||
|
promote_min_magic::Int = 3 # required fixed signature positions for nomination
|
||||||
|
# The stage-5 catalog is a single durable file mutated by the one sweep
|
||||||
|
# process (never a worker), and nominations are written as one file per
|
||||||
|
# promotable cluster for a human to glance at before promoting (DESIGN §9/§6).
|
||||||
|
cluster_catalog_path::String = "data/catalog.json" # durable phase-B catalog (sufficient stats + processed set)
|
||||||
|
nominated_dir::String = "data/nominated" # one JSON per self-nominated cluster, awaiting a human promote
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
parse_backend(s) -> Symbol
|
||||||
|
|
||||||
|
Turn `FS_QUEUE_BACKEND` into a `Config.queue_backend`. Unknown values are a
|
||||||
|
startup error rather than a silent fallback: a typo that quietly left you on the
|
||||||
|
in-process queue would look exactly like durability working right up until a
|
||||||
|
crash proved it wasn't.
|
||||||
|
"""
|
||||||
|
function parse_backend(s::AbstractString)
|
||||||
|
b = Symbol(lowercase(strip(s)))
|
||||||
|
b in (:channel, :rabbitmq) ||
|
||||||
|
throw(ArgumentError("FS_QUEUE_BACKEND must be \"channel\" or \"rabbitmq\", got \"$s\""))
|
||||||
|
return b
|
||||||
|
end
|
||||||
|
|
||||||
|
"Parse a boolean env var. Accepts the usual spellings; anything else is an error."
|
||||||
|
function parse_bool(s::AbstractString)
|
||||||
|
v = lowercase(strip(s))
|
||||||
|
v in ("1", "true", "yes", "on") && return true
|
||||||
|
v in ("0", "false", "no", "off") && return false
|
||||||
|
throw(ArgumentError("expected a boolean (true/false), got \"$s\""))
|
||||||
end
|
end
|
||||||
|
|
||||||
"""
|
"""
|
||||||
@@ -18,30 +108,80 @@ end
|
|||||||
|
|
||||||
Build a `Config` from environment variables, falling back to the struct
|
Build a `Config` from environment variables, falling back to the struct
|
||||||
defaults. Any keyword `overrides` win over the environment (useful for tests
|
defaults. Any keyword `overrides` win over the environment (useful for tests
|
||||||
and for `FileServer.run(; port=...)`).
|
and for `DarkStruct.run(; port=...)`).
|
||||||
|
|
||||||
Recognised variables:
|
Recognised variables:
|
||||||
FS_HOST, FS_PORT, FS_WORKERS, FS_QUEUE_CAPACITY,
|
FS_HOST, FS_PORT, FS_WORKERS, FS_QUEUE_CAPACITY,
|
||||||
FS_SPOOL_DIR, FS_DONE_DIR, FS_FAILED_DIR, FS_MODEL_PATH
|
FS_KNOWN_WORKERS, FS_KNOWN_QUEUE_CAPACITY,
|
||||||
|
FS_UNKNOWN_WORKERS, FS_UNKNOWN_QUEUE_CAPACITY,
|
||||||
|
FS_TEXT_WORKERS, FS_TEXT_QUEUE_CAPACITY,
|
||||||
|
FS_SPOOL_DIR, FS_BINARY_DIR,
|
||||||
|
FS_DONE_DIR, FS_TEXT_DONE_DIR, FS_FAILED_DIR, FS_MODEL_PATH,
|
||||||
|
FS_UPLOAD_CHUNK_BYTES, FS_EXIFTOOL_TIMEOUT, FS_LINGUIST_TIMEOUT,
|
||||||
|
FS_CLUSTER_DIR, FS_CLUSTER_N, FS_CLUSTER_ALPHA, FS_CLUSTER_PSEUDOCOUNT,
|
||||||
|
FS_CLUSTER_BG_MASS, FS_PROMOTE_MIN_MEMBERS, FS_PROMOTE_MIN_MAGIC,
|
||||||
|
FS_CLUSTER_CATALOG, FS_NOMINATED_DIR,
|
||||||
|
FS_QUEUE_BACKEND, FS_AMQP_URL, FS_AMQP_PREFIX, FS_AMQP_PREFETCH,
|
||||||
|
FS_AMQP_CONFIRMS, FS_RECOVER_SPOOL
|
||||||
"""
|
"""
|
||||||
function config_from_env(; host=nothing, port=nothing, worker_count=nothing,
|
function config_from_env(; host=nothing, port=nothing, worker_count=nothing,
|
||||||
queue_capacity=nothing, spool_dir=nothing,
|
queue_capacity=nothing, known_worker_count=nothing,
|
||||||
done_dir=nothing, failed_dir=nothing, model_path=nothing)
|
known_queue_capacity=nothing, unknown_worker_count=nothing,
|
||||||
|
unknown_queue_capacity=nothing, text_worker_count=nothing,
|
||||||
|
text_queue_capacity=nothing, spool_dir=nothing,
|
||||||
|
binary_dir=nothing, done_dir=nothing, text_done_dir=nothing,
|
||||||
|
failed_dir=nothing, model_path=nothing,
|
||||||
|
upload_chunk_bytes=nothing, exiftool_timeout=nothing,
|
||||||
|
linguist_timeout=nothing, cluster_dir=nothing, cluster_n=nothing,
|
||||||
|
cluster_alpha=nothing, cluster_pseudocount=nothing,
|
||||||
|
cluster_bg_mass=nothing, promote_min_members=nothing,
|
||||||
|
promote_min_magic=nothing, cluster_catalog_path=nothing,
|
||||||
|
nominated_dir=nothing, queue_backend=nothing,
|
||||||
|
amqp_url=nothing, amqp_prefix=nothing,
|
||||||
|
amqp_prefetch=nothing, amqp_confirms=nothing,
|
||||||
|
recover_spool=nothing)
|
||||||
Config(
|
Config(
|
||||||
host = something(host, get(ENV, "FS_HOST", "127.0.0.1")),
|
host = something(host, get(ENV, "FS_HOST", "127.0.0.1")),
|
||||||
port = something(port, parse(Int, get(ENV, "FS_PORT", "8080"))),
|
port = something(port, parse(Int, get(ENV, "FS_PORT", "8080"))),
|
||||||
worker_count = something(worker_count, parse(Int, get(ENV, "FS_WORKERS", string(Threads.nthreads())))),
|
worker_count = something(worker_count, parse(Int, get(ENV, "FS_WORKERS", string(Threads.nthreads())))),
|
||||||
queue_capacity = something(queue_capacity, parse(Int, get(ENV, "FS_QUEUE_CAPACITY", "1000"))),
|
queue_capacity = something(queue_capacity, parse(Int, get(ENV, "FS_QUEUE_CAPACITY", "1000"))),
|
||||||
|
known_worker_count = something(known_worker_count, parse(Int, get(ENV, "FS_KNOWN_WORKERS", string(Threads.nthreads())))),
|
||||||
|
known_queue_capacity = something(known_queue_capacity, parse(Int, get(ENV, "FS_KNOWN_QUEUE_CAPACITY", "1000"))),
|
||||||
|
unknown_worker_count = something(unknown_worker_count, parse(Int, get(ENV, "FS_UNKNOWN_WORKERS", string(Threads.nthreads())))),
|
||||||
|
unknown_queue_capacity = something(unknown_queue_capacity, parse(Int, get(ENV, "FS_UNKNOWN_QUEUE_CAPACITY", "1000"))),
|
||||||
|
text_worker_count = something(text_worker_count, parse(Int, get(ENV, "FS_TEXT_WORKERS", string(Threads.nthreads())))),
|
||||||
|
text_queue_capacity = something(text_queue_capacity, parse(Int, get(ENV, "FS_TEXT_QUEUE_CAPACITY", "1000"))),
|
||||||
spool_dir = something(spool_dir, get(ENV, "FS_SPOOL_DIR", "data/spool")),
|
spool_dir = something(spool_dir, get(ENV, "FS_SPOOL_DIR", "data/spool")),
|
||||||
|
binary_dir = something(binary_dir, get(ENV, "FS_BINARY_DIR", "data/binary")),
|
||||||
done_dir = something(done_dir, get(ENV, "FS_DONE_DIR", "data/done")),
|
done_dir = something(done_dir, get(ENV, "FS_DONE_DIR", "data/done")),
|
||||||
|
text_done_dir = something(text_done_dir, get(ENV, "FS_TEXT_DONE_DIR", "data/text_done")),
|
||||||
failed_dir = something(failed_dir, get(ENV, "FS_FAILED_DIR", "data/failed")),
|
failed_dir = something(failed_dir, get(ENV, "FS_FAILED_DIR", "data/failed")),
|
||||||
model_path = something(model_path, get(ENV, "FS_MODEL_PATH", "model/classifier.jld2")),
|
model_path = something(model_path, get(ENV, "FS_MODEL_PATH", "model/classifier.jld2")),
|
||||||
|
upload_chunk_bytes = something(upload_chunk_bytes, parse(Int, get(ENV, "FS_UPLOAD_CHUNK_BYTES", string(UPLOAD_CHUNK_BYTES)))),
|
||||||
|
exiftool_timeout = something(exiftool_timeout, parse(Int, get(ENV, "FS_EXIFTOOL_TIMEOUT", "30"))),
|
||||||
|
linguist_timeout = something(linguist_timeout, parse(Int, get(ENV, "FS_LINGUIST_TIMEOUT", "30"))),
|
||||||
|
cluster_dir = something(cluster_dir, get(ENV, "FS_CLUSTER_DIR", "data/binary")),
|
||||||
|
cluster_n = something(cluster_n, parse(Int, get(ENV, "FS_CLUSTER_N", "32"))),
|
||||||
|
cluster_alpha = something(cluster_alpha, parse(Float64, get(ENV, "FS_CLUSTER_ALPHA", "1.0"))),
|
||||||
|
cluster_pseudocount = something(cluster_pseudocount, parse(Float64, get(ENV, "FS_CLUSTER_PSEUDOCOUNT", "0.1"))),
|
||||||
|
cluster_bg_mass = something(cluster_bg_mass, parse(Float64, get(ENV, "FS_CLUSTER_BG_MASS", "5.0"))),
|
||||||
|
promote_min_members = something(promote_min_members, parse(Int, get(ENV, "FS_PROMOTE_MIN_MEMBERS", "20"))),
|
||||||
|
promote_min_magic = something(promote_min_magic, parse(Int, get(ENV, "FS_PROMOTE_MIN_MAGIC", "3"))),
|
||||||
|
cluster_catalog_path = something(cluster_catalog_path, get(ENV, "FS_CLUSTER_CATALOG", "data/catalog.json")),
|
||||||
|
nominated_dir = something(nominated_dir, get(ENV, "FS_NOMINATED_DIR", "data/nominated")),
|
||||||
|
queue_backend = something(queue_backend, parse_backend(get(ENV, "FS_QUEUE_BACKEND", "channel"))),
|
||||||
|
amqp_url = something(amqp_url, get(ENV, "FS_AMQP_URL", "amqp://guest:guest@localhost:5672/")),
|
||||||
|
amqp_prefix = something(amqp_prefix, get(ENV, "FS_AMQP_PREFIX", "darkstruct")),
|
||||||
|
amqp_prefetch = something(amqp_prefetch, parse(Int, get(ENV, "FS_AMQP_PREFETCH", "0"))),
|
||||||
|
amqp_confirms = something(amqp_confirms, parse_bool(get(ENV, "FS_AMQP_CONFIRMS", "true"))),
|
||||||
|
recover_spool = something(recover_spool, parse_bool(get(ENV, "FS_RECOVER_SPOOL", "false"))),
|
||||||
)
|
)
|
||||||
end
|
end
|
||||||
|
|
||||||
"Create the spool/done/failed directories if they don't already exist."
|
"Create all the pipeline-stage directories if they don't already exist."
|
||||||
function ensure_dirs(cfg::Config)
|
function ensure_dirs(cfg::Config)
|
||||||
for d in (cfg.spool_dir, cfg.done_dir, cfg.failed_dir)
|
for d in (cfg.spool_dir, cfg.binary_dir, cfg.done_dir, cfg.text_done_dir,
|
||||||
|
cfg.failed_dir, cfg.nominated_dir)
|
||||||
mkpath(d)
|
mkpath(d)
|
||||||
end
|
end
|
||||||
return nothing
|
return nothing
|
||||||
|
|||||||
67
src/content.jl
Normal file
67
src/content.jl
Normal file
@@ -0,0 +1,67 @@
|
|||||||
|
# Stage-3 content triage for unknown files.
|
||||||
|
#
|
||||||
|
# A file that stage-1 couldn't recognize is still sorted into one of two coarse
|
||||||
|
# buckets so downstream tooling can treat them differently: text (human-readable)
|
||||||
|
# and binary (everything else). Only binary is a directory, because it is terminal:
|
||||||
|
# the file is committed to `binary/`; a text file has stage 4 still to come, so it
|
||||||
|
# stays in `spool/` and only its queue reference moves on. We sniff only the first
|
||||||
|
# `CONTENT_SNIFF_BYTES` (no full read) and ask two questions: does the window
|
||||||
|
# decode as valid UTF-8, and are any of its control bytes ones that don't belong
|
||||||
|
# in text? This is the Unicode-aware successor to the classic "NUL byte" test.
|
||||||
|
# It accepts non-ASCII text (accents, CJK, emoji) instead of misfiling it as
|
||||||
|
# binary, while still rejecting binary formats, which almost never form valid
|
||||||
|
# UTF-8 near their start (and a NUL is never a valid UTF-8 scalar, so it still
|
||||||
|
# reads as binary for free).
|
||||||
|
|
||||||
|
const CONTENT_SNIFF_BYTES = 8000
|
||||||
|
|
||||||
|
# Control bytes (< 0x20) that appear legitimately in text: BS, TAB, LF, VT, FF,
|
||||||
|
# CR, and ESC (ANSI-colored logs). Any *other* control byte is a binary signal.
|
||||||
|
const TEXT_CONTROL_BYTES = (0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x1b)
|
||||||
|
|
||||||
|
# Drop a trailing UTF-8 sequence that the sniff window cut in half, so a
|
||||||
|
# multi-byte character straddling the boundary isn't mistaken for invalid bytes.
|
||||||
|
# Continuation bytes are 0x80–0xBF; a lead byte encodes its own sequence length
|
||||||
|
# in its high bits. We walk back over the trailing continuation bytes, and if
|
||||||
|
# the lead byte we land on expects more bytes than the window actually holds,
|
||||||
|
# trim the whole incomplete sequence.
|
||||||
|
function trim_truncated_utf8(chunk::AbstractVector{UInt8})
|
||||||
|
n = length(chunk)
|
||||||
|
n == 0 && return chunk
|
||||||
|
# Find the start of the final byte sequence: skip back over continuations.
|
||||||
|
i = n
|
||||||
|
while i > 0 && (chunk[i] & 0xc0) == 0x80
|
||||||
|
i -= 1
|
||||||
|
end
|
||||||
|
i == 0 && return chunk # all continuations; leave as-is
|
||||||
|
lead = chunk[i]
|
||||||
|
# How many bytes does this lead byte announce?
|
||||||
|
expected = lead < 0x80 ? 1 : # ASCII
|
||||||
|
lead < 0xe0 ? 2 : # 110xxxxx
|
||||||
|
lead < 0xf0 ? 3 : # 1110xxxx
|
||||||
|
4 # 11110xxx
|
||||||
|
have = n - i + 1
|
||||||
|
return have < expected ? view(chunk, 1:i-1) : chunk
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
is_binary(path) -> Bool
|
||||||
|
|
||||||
|
Classify a file as binary (`true`) or text (`false`) by sniffing its first
|
||||||
|
`CONTENT_SNIFF_BYTES` bytes. A file is text when that window (minus any
|
||||||
|
multi-byte character truncated by the window edge) is valid UTF-8 and contains
|
||||||
|
no control bytes outside the text-safe set (`TEXT_CONTROL_BYTES`). An empty file
|
||||||
|
is treated as text.
|
||||||
|
"""
|
||||||
|
function is_binary(path::AbstractString)::Bool
|
||||||
|
open(path, "r") do io
|
||||||
|
chunk = read(io, CONTENT_SNIFF_BYTES)
|
||||||
|
isempty(chunk) && return false
|
||||||
|
window = trim_truncated_utf8(chunk)
|
||||||
|
# Malformed UTF-8 → binary.
|
||||||
|
isvalid(String(copy(window))) || return true
|
||||||
|
# Valid UTF-8, but a stray non-text control byte still means binary.
|
||||||
|
# (NUL is a valid UTF-8 scalar, so it's rejected here, not above.)
|
||||||
|
return any(b -> b < 0x20 && !(b in TEXT_CONTROL_BYTES), window)
|
||||||
|
end
|
||||||
|
end
|
||||||
21
src/job.jl
21
src/job.jl
@@ -1,8 +1,11 @@
|
|||||||
# A unit of work on the queue. Deliberately lightweight: the file *bytes* live
|
# A unit of work on the queue. Deliberately lightweight: the file *bytes* live
|
||||||
# on disk in the spool directory, and only this small reference travels through
|
# on disk in the spool directory, and only this small reference travels through
|
||||||
# the queue. This is what keeps intake fast and memory flat regardless of file
|
# the queue. This is what keeps intake fast and memory flat regardless of file
|
||||||
# size, and it's the shape you'd publish to RabbitMQ later (the "claim check"
|
# size, and it's the shape published to RabbitMQ (the "claim check" pattern:
|
||||||
# pattern — enqueue a reference, not the payload).
|
# enqueue a reference, not the payload).
|
||||||
|
#
|
||||||
|
# `path` is a *local* filesystem path, which is why a broker-backed pipeline is
|
||||||
|
# still single-consumer-process: see "Queue backends" in the README.
|
||||||
|
|
||||||
struct Job
|
struct Job
|
||||||
id::String # server-minted UUID; also the on-disk filename prefix
|
id::String # server-minted UUID; also the on-disk filename prefix
|
||||||
@@ -10,4 +13,18 @@ struct Job
|
|||||||
path::String # absolute-ish path to the spooled file
|
path::String # absolute-ish path to the spooled file
|
||||||
size::Int # bytes
|
size::Int # bytes
|
||||||
received_at::Float64 # time() at intake
|
received_at::Float64 # time() at intake
|
||||||
|
# Broker delivery tag, and the one field that is *not* part of the job: it
|
||||||
|
# identifies this delivery on the channel that delivered it, so it is set on
|
||||||
|
# receipt and never serialized into the message body (a tag from a previous
|
||||||
|
# connection is meaningless, and a redelivery gets a fresh one). 0 means "no
|
||||||
|
# broker" — every job under `ChannelQueue`, and every job before publishing.
|
||||||
|
# Int64 rather than UInt64 to match AMQPClient's `TAMQPDeliveryTag`.
|
||||||
|
delivery_tag::Int64
|
||||||
end
|
end
|
||||||
|
|
||||||
|
Job(id, original_name, path, size, received_at) =
|
||||||
|
Job(id, original_name, path, size, received_at, Int64(0))
|
||||||
|
|
||||||
|
"Copy of `job` carrying the delivery tag it arrived with."
|
||||||
|
with_delivery_tag(job::Job, tag::Integer) =
|
||||||
|
Job(job.id, job.original_name, job.path, job.size, job.received_at, Int64(tag))
|
||||||
|
|||||||
155
src/language.jl
Normal file
155
src/language.jl
Normal file
@@ -0,0 +1,155 @@
|
|||||||
|
# Stage-4 language enrichment for text files.
|
||||||
|
#
|
||||||
|
# A file that stage-3 sorted as text is human-readable, but we don't yet know
|
||||||
|
# *what* it is. This stage answers two questions and records them in a
|
||||||
|
# `.meta.json` sidecar, exactly like the stage-2 known-file enrichment:
|
||||||
|
#
|
||||||
|
# * natural language, via Languages.jl's `LanguageDetector` (a Julia port of
|
||||||
|
# the `whatlang` n-gram model): English vs. French vs. Japanese, plus a
|
||||||
|
# confidence score. Pure Julia, no subprocess.
|
||||||
|
# * programming language, via the `github-linguist` CLI, which recognizes
|
||||||
|
# source and markup by extension + content heuristics. There is no
|
||||||
|
# comparable native Julia library, so we shell out (mirroring stage-2's
|
||||||
|
# exiftool dependency).
|
||||||
|
#
|
||||||
|
# Neither detector failing quarantines the file: a text file is wanted whether
|
||||||
|
# or not we can name its language, so a failure yields a *degraded* sidecar
|
||||||
|
# (what we know plus an `error` note), just like stage 2.
|
||||||
|
#
|
||||||
|
# The github-linguist quirk that shapes this code: run against a path *inside* a
|
||||||
|
# git repository, linguist reads the file's committed git blob, not the bytes on
|
||||||
|
# disk, and an untracked file (which every file under `data/` is) has no blob,
|
||||||
|
# so it crashes. We sidestep this by copying the file to a fresh temp dir outside
|
||||||
|
# any repo (preserving its name so linguist's extension heuristics still fire)
|
||||||
|
# and pointing linguist there.
|
||||||
|
|
||||||
|
# How much of a text file to feed the natural-language detector. The whatlang
|
||||||
|
# model saturates quickly, so a bounded prefix keeps memory flat on huge logs
|
||||||
|
# while still giving the detector plenty of signal.
|
||||||
|
const LANG_SAMPLE_BYTES = 65_536
|
||||||
|
|
||||||
|
"Return true if the `github-linguist` binary is on PATH."
|
||||||
|
function linguist_available()
|
||||||
|
try
|
||||||
|
Base.run(pipeline(`github-linguist --version`; stdout=devnull, stderr=devnull))
|
||||||
|
return true
|
||||||
|
catch
|
||||||
|
return false
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
read_text_sample(path) -> String
|
||||||
|
|
||||||
|
Read up to `LANG_SAMPLE_BYTES` of `path` as UTF-8 text, trimming a multi-byte
|
||||||
|
character the window may have cut in half (reusing stage-3's `trim_truncated_utf8`)
|
||||||
|
so the tail isn't misread as garbage.
|
||||||
|
"""
|
||||||
|
function read_text_sample(path::AbstractString)::String
|
||||||
|
open(path, "r") do io
|
||||||
|
chunk = read(io, LANG_SAMPLE_BYTES)
|
||||||
|
return String(copy(trim_truncated_utf8(chunk)))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
detect_natural_language(detector, text) -> (name, code, confidence)
|
||||||
|
|
||||||
|
Run the `LanguageDetector` on `text`, returning the language's English name
|
||||||
|
(e.g. `"English"`), its ISO 639-3 code (e.g. `"eng"`), and the model's
|
||||||
|
confidence in `[0,1]`. Returns `(nothing, nothing, nothing)` when there is no
|
||||||
|
usable text (empty/whitespace) or the detector errors; the caller records that
|
||||||
|
as a degraded result rather than failing the file.
|
||||||
|
"""
|
||||||
|
function detect_natural_language(detector, text::AbstractString)
|
||||||
|
isempty(strip(text)) && return (nothing, nothing, nothing)
|
||||||
|
try
|
||||||
|
lang, _script, confidence = detector(text)
|
||||||
|
return (Languages.english_name(lang), Languages.isocode(lang), confidence)
|
||||||
|
catch
|
||||||
|
return (nothing, nothing, nothing)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
run_linguist(path, timeout) -> Union{String,Nothing}
|
||||||
|
|
||||||
|
Ask `github-linguist --json` for the programming/markup language of the file at
|
||||||
|
`path`, returning the language name (e.g. `"Python"`, `"Markdown"`) or `nothing`
|
||||||
|
when linguist can't identify one. Plain prose reports as `"Text"` and
|
||||||
|
unrecognized content as JSON `null`; both collapse to `nothing` here, since only
|
||||||
|
a real programming/markup language is worth recording.
|
||||||
|
|
||||||
|
`path` MUST be outside any git repository; see the module header for why.
|
||||||
|
"""
|
||||||
|
function run_linguist(path::AbstractString, timeout::Integer)
|
||||||
|
bytes = run_with_timeout(`github-linguist --json $path`, timeout)
|
||||||
|
bytes === nothing && return nothing
|
||||||
|
|
||||||
|
parsed = try
|
||||||
|
JSON3.read(String(bytes))
|
||||||
|
catch
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
# linguist --json emits a single object keyed by the file path; pull the one
|
||||||
|
# entry rather than depend on the exact key spelling.
|
||||||
|
isempty(parsed) && return nothing
|
||||||
|
entry = first(values(parsed))
|
||||||
|
lang = get(entry, :language, nothing)
|
||||||
|
(lang === nothing || lang == "Text") && return nothing
|
||||||
|
return String(lang)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
detect_programming_language(job, cfg) -> Union{String,Nothing}
|
||||||
|
|
||||||
|
Programming/markup language of a text file, or `nothing`. Copies the file to a
|
||||||
|
throwaway temp dir *outside* the git repo (under its sanitized original name, so
|
||||||
|
linguist's extension heuristics still apply), runs linguist there, and cleans up.
|
||||||
|
"""
|
||||||
|
function detect_programming_language(job::Job, cfg::Config)
|
||||||
|
lang = nothing
|
||||||
|
mktempdir() do dir # tempdir() → /tmp, outside the repo
|
||||||
|
safe = sanitize_filename(job.original_name)
|
||||||
|
tmp = joinpath(dir, safe)
|
||||||
|
cp(job.path, tmp; force=true)
|
||||||
|
lang = run_linguist(tmp, cfg.linguist_timeout)
|
||||||
|
end
|
||||||
|
return lang
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
build_text_metadata(detector, job, cfg) -> NamedTuple
|
||||||
|
|
||||||
|
Build the stage-4 sidecar payload for a text file: its natural language (name +
|
||||||
|
ISO code + confidence) and programming/markup language, plus the Job's
|
||||||
|
authoritative id/name/size. `error` is set only when natural-language detection
|
||||||
|
produced nothing usable (the file is still enriched and committed); programming
|
||||||
|
language is best-effort and its absence is normal, not an error.
|
||||||
|
"""
|
||||||
|
function build_text_metadata(detector, job::Job, cfg::Config)
|
||||||
|
text = read_text_sample(job.path)
|
||||||
|
name, code, confidence = detect_natural_language(detector, text)
|
||||||
|
programming_language = detect_programming_language(job, cfg)
|
||||||
|
|
||||||
|
return (
|
||||||
|
id = job.id,
|
||||||
|
original_name = job.original_name,
|
||||||
|
file_size = job.size, # authoritative, from intake
|
||||||
|
content_type = "text",
|
||||||
|
language = name,
|
||||||
|
language_code = code,
|
||||||
|
language_confidence = confidence,
|
||||||
|
programming_language = programming_language,
|
||||||
|
error = name === nothing ? "language detection produced no result" : nothing,
|
||||||
|
)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
finalize_text!(cfg, job, meta) -> (file_dest, sidecar_dest)
|
||||||
|
|
||||||
|
Commit an enriched text file (stage 4) to `text_done/` via the shared
|
||||||
|
sidecar-first `commit_enriched!`, giving text files the same crash-safe
|
||||||
|
"file implies sidecar" guarantee as stage-2 known files.
|
||||||
|
"""
|
||||||
|
finalize_text!(cfg::Config, job::Job, meta) = commit_enriched!(cfg.text_done_dir, job, meta)
|
||||||
267
src/metadata.jl
Normal file
267
src/metadata.jl
Normal file
@@ -0,0 +1,267 @@
|
|||||||
|
# Stage-2 metadata extraction and enrichment.
|
||||||
|
#
|
||||||
|
# Known files are enriched by shelling out to `exiftool -json -G` (the only tool
|
||||||
|
# with broad, multi-format coverage; there is no comparable native Julia
|
||||||
|
# library), then normalizing its output into a small, stable, documented schema
|
||||||
|
# that downstream consumers can rely on, while preserving the full raw dump.
|
||||||
|
#
|
||||||
|
# exiftool being installed is a hard startup prerequisite (see `assert_exiftool`,
|
||||||
|
# called from `run`). A per-file extraction failure or hang does NOT quarantine
|
||||||
|
# the file. It produces a *degraded* sidecar recording what we know plus the
|
||||||
|
# error, because a file that passed classification is wanted regardless of
|
||||||
|
# whether we could read its metadata.
|
||||||
|
|
||||||
|
"Throw at startup if the `exiftool` binary isn't on PATH: fail fast rather than discover it per file."
|
||||||
|
function assert_exiftool()
|
||||||
|
try
|
||||||
|
Base.run(pipeline(`exiftool -ver`; stdout=devnull, stderr=devnull))
|
||||||
|
catch
|
||||||
|
error("exiftool not found on PATH (install libimage-exiftool-perl / exiftool). It is required for stage-2 metadata enrichment.")
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# Each normalized field is a coalesce over exiftool tag names, tried in order;
|
||||||
|
# the first present, non-empty value wins. exiftool with `-G` prefixes tags by
|
||||||
|
# group (e.g. "EXIF:Software"), so we match on the bare tag name after the last
|
||||||
|
# colon. Extend a field simply by appending tag names here.
|
||||||
|
#
|
||||||
|
# Note the documented `Creator` ambiguity: in PDF it's the authoring app, but
|
||||||
|
# elsewhere it's the person. We take the simple route: `Creator` feeds
|
||||||
|
# `created_by` only, and `author` relies on the person-specific tags.
|
||||||
|
const CREATED_BY_TAGS = ["Producer", "CreatorTool", "Creator", "Software", "Application", "Encoder", "EncodingTool", "HostComputer"]
|
||||||
|
const AUTHOR_TAGS = ["Author", "Artist", "By-line", "Owner", "Artist"]
|
||||||
|
const CREATED_DATE_TAGS = ["DateTimeOriginal", "CreateDate", "MediaCreateDate", "CreationDate"]
|
||||||
|
const MODIFIED_DATE_TAGS = ["ModifyDate", "FileModifyDate"]
|
||||||
|
|
||||||
|
"fsync an open file descriptor, throwing on failure. Makes a write durable before a rename commits it."
|
||||||
|
function fsync_fd(fd)
|
||||||
|
ccall(:fsync, Cint, (Cint,), fd) == 0 || error("fsync failed: $(Base.Libc.strerror())")
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"fsync a directory so a rename into it survives a crash (the rename, not just the file bytes, must be persisted)."
|
||||||
|
function fsync_dir(dir::AbstractString)
|
||||||
|
dfd = ccall(:open, Cint, (Cstring, Cint), dir, 0) # O_RDONLY
|
||||||
|
dfd < 0 && error("cannot open dir for fsync: $dir ($(Base.Libc.strerror()))")
|
||||||
|
try
|
||||||
|
fsync_fd(dfd)
|
||||||
|
finally
|
||||||
|
ccall(:close, Cint, (Cint,), dfd)
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"Strip exiftool's `-G` group prefix (`EXIF:Software` → `Software`) so lookups are group-agnostic."
|
||||||
|
strip_group(tag::AbstractString) = String(last(split(tag, ':')))
|
||||||
|
|
||||||
|
"Return the first present, non-empty value among `tags` in the group-stripped map, or `nothing`."
|
||||||
|
function coalesce_tag(bytag::Dict{String,Any}, tags)
|
||||||
|
for t in tags
|
||||||
|
v = get(bytag, t, nothing)
|
||||||
|
v === nothing && continue
|
||||||
|
s = string(v)
|
||||||
|
isempty(strip(s)) && continue
|
||||||
|
return v
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"How long a child gets to honor SIGTERM before `run_with_timeout` escalates to SIGKILL."
|
||||||
|
const KILL_GRACE_SECONDS = 2.0
|
||||||
|
|
||||||
|
"""
|
||||||
|
Send `signum` to the whole process group `pgid` (a negative pid means "the group"
|
||||||
|
to `kill(2)`). Julia's `kill(::Process, sig)` signals only the child itself,
|
||||||
|
which is not enough to enforce a timeout; see `run_with_timeout`.
|
||||||
|
"""
|
||||||
|
function signal_group(pgid::Integer, signum::Integer)
|
||||||
|
ccall(:kill, Cint, (Cint, Cint), -pgid, signum)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
run_with_timeout(cmd, timeout) -> Union{Vector{UInt8},Nothing}
|
||||||
|
|
||||||
|
Run `cmd`, capturing stdout, and return the captured bytes on clean exit, or
|
||||||
|
`nothing` on non-zero exit or timeout. The subprocess is killed (SIGTERM, then
|
||||||
|
SIGKILL after a grace period) once it overruns `timeout` seconds, so one
|
||||||
|
pathological input can't wedge a worker forever. Shared by the exiftool (stage 2)
|
||||||
|
and github-linguist (stage 4) shells.
|
||||||
|
|
||||||
|
The child runs in its own process group and the timeout signals the *group*, not
|
||||||
|
just the child. This is what makes the timeout enforceable: `wait` below returns
|
||||||
|
only once the captured stdout pipe closes, and any grandchild inherits that pipe,
|
||||||
|
so signalling the child alone leaves a `sh -c "...; sleep 30"`-shaped process
|
||||||
|
tree running to completion with the worker still blocked on it. The cost of the
|
||||||
|
process group is that a hard crash of the server orphans an in-flight child
|
||||||
|
rather than taking it down; these children are short-lived and timeout-bounded,
|
||||||
|
which is the cheaper side of that trade.
|
||||||
|
"""
|
||||||
|
function run_with_timeout(cmd::Cmd, timeout::Integer)
|
||||||
|
out = IOBuffer()
|
||||||
|
# `detach` puts the child in a fresh process group (it becomes the group
|
||||||
|
# leader, so the group id is its pid); see the docstring for why the group,
|
||||||
|
# and not the child, is what the timeout has to signal.
|
||||||
|
proc = Base.run(pipeline(detach(cmd); stdout=out, stderr=devnull); wait=false)
|
||||||
|
pgid = Base.getpid(proc)
|
||||||
|
|
||||||
|
# A one-shot timer, cancelled the moment the child exits, rather than a
|
||||||
|
# polling loop the caller has to join. The polling version charged every
|
||||||
|
# call the remainder of its in-flight `sleep(0.1)` *after* the child had
|
||||||
|
# already exited: ~50 ms on average, and a measured 101 ms on a process
|
||||||
|
# that exits instantly. That is pure latency on the hot path of two stages
|
||||||
|
# (exiftool here, github-linguist in stage 4), and it dwarfed the work on
|
||||||
|
# anything but a slow file. Waiting on the process directly costs nothing
|
||||||
|
# when the child exits normally, which is the overwhelmingly common case.
|
||||||
|
killed = Ref(false)
|
||||||
|
timer = Timer(timeout) do _
|
||||||
|
process_running(proc) || return
|
||||||
|
killed[] = true
|
||||||
|
signal_group(pgid, Base.SIGTERM)
|
||||||
|
# Escalate: a process that ignores/defers SIGTERM would otherwise pin the
|
||||||
|
# worker forever on the wait(proc) below, defeating the timeout. This
|
||||||
|
# runs off the timer's task so the event loop isn't held during the
|
||||||
|
# grace period, and it is not joined: by the time it wakes, `wait(proc)`
|
||||||
|
# has long since returned and `process_running` settles it.
|
||||||
|
Threads.@spawn begin
|
||||||
|
deadline = time() + KILL_GRACE_SECONDS
|
||||||
|
while process_running(proc) && time() < deadline
|
||||||
|
sleep(0.05)
|
||||||
|
end
|
||||||
|
process_running(proc) && signal_group(pgid, Base.SIGKILL)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
try
|
||||||
|
wait(proc)
|
||||||
|
finally
|
||||||
|
close(timer) # cancel the pending kill; a no-op if it already fired
|
||||||
|
end
|
||||||
|
|
||||||
|
(killed[] || !success(proc)) && return nothing
|
||||||
|
return take!(out)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
run_exiftool(path, timeout) -> Union{Dict{String,Any},Nothing}
|
||||||
|
|
||||||
|
Run `exiftool -json -G` on `path`, returning the parsed tag object, or `nothing`
|
||||||
|
on non-zero exit, unparseable output, or timeout.
|
||||||
|
"""
|
||||||
|
function run_exiftool(path::AbstractString, timeout::Integer)
|
||||||
|
# -json: machine output; -G: group-prefixed tags; -n: numeric (unformatted)
|
||||||
|
# values so sizes/durations are numbers, not display strings.
|
||||||
|
bytes = run_with_timeout(`exiftool -json -G -n $path`, timeout)
|
||||||
|
bytes === nothing && return nothing
|
||||||
|
|
||||||
|
parsed = try
|
||||||
|
JSON3.read(String(bytes))
|
||||||
|
catch
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
# exiftool -json emits a one-element array of objects (one per input file).
|
||||||
|
(parsed isa AbstractVector && !isempty(parsed)) || return nothing
|
||||||
|
return Dict{String,Any}(String(strip_group(String(k))) => v for (k, v) in pairs(parsed[1]))
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
build_metadata(job, cfg) -> NamedTuple
|
||||||
|
|
||||||
|
Extract and normalize metadata for a known file. Always returns a sidecar
|
||||||
|
payload: on extraction success, the normalized fields plus the full raw dump; on
|
||||||
|
failure/timeout, a *degraded* payload with what we know from the Job plus an
|
||||||
|
`error` note. `file_size` always comes from the Job (authoritative), never
|
||||||
|
exiftool.
|
||||||
|
"""
|
||||||
|
function build_metadata(job::Job, cfg::Config)
|
||||||
|
bytag = run_exiftool(job.path, cfg.exiftool_timeout)
|
||||||
|
|
||||||
|
if bytag === nothing
|
||||||
|
return (
|
||||||
|
id = job.id,
|
||||||
|
original_name = job.original_name,
|
||||||
|
file_type = nothing,
|
||||||
|
mime_type = nothing,
|
||||||
|
file_size = job.size,
|
||||||
|
created_date = nothing,
|
||||||
|
modified_date = nothing,
|
||||||
|
author = nothing,
|
||||||
|
created_by = nothing,
|
||||||
|
dimensions = nothing,
|
||||||
|
duration = nothing,
|
||||||
|
page_count = nothing,
|
||||||
|
error = "exiftool extraction failed or timed out",
|
||||||
|
raw = nothing,
|
||||||
|
)
|
||||||
|
end
|
||||||
|
|
||||||
|
return normalize_metadata(job, bytag)
|
||||||
|
end
|
||||||
|
|
||||||
|
"Build the normalized sidecar payload from a successful exiftool tag map."
|
||||||
|
function normalize_metadata(job::Job, bytag::Dict{String,Any})
|
||||||
|
w = get(bytag, "ImageWidth", nothing)
|
||||||
|
h = get(bytag, "ImageHeight", nothing)
|
||||||
|
dims = (w !== nothing && h !== nothing) ? (; width = w, height = h) : nothing
|
||||||
|
|
||||||
|
return (
|
||||||
|
id = job.id,
|
||||||
|
original_name = job.original_name,
|
||||||
|
file_type = get(bytag, "FileType", get(bytag, "FileTypeExtension", nothing)),
|
||||||
|
mime_type = get(bytag, "MIMEType", nothing),
|
||||||
|
file_size = job.size, # authoritative, from intake
|
||||||
|
created_date = coalesce_tag(bytag, CREATED_DATE_TAGS),
|
||||||
|
modified_date = coalesce_tag(bytag, MODIFIED_DATE_TAGS),
|
||||||
|
author = coalesce_tag(bytag, AUTHOR_TAGS),
|
||||||
|
created_by = coalesce_tag(bytag, CREATED_BY_TAGS),
|
||||||
|
dimensions = dims,
|
||||||
|
duration = get(bytag, "Duration", get(bytag, "MediaDuration", nothing)),
|
||||||
|
page_count = get(bytag, "PageCount", nothing),
|
||||||
|
error = nothing,
|
||||||
|
raw = bytag,
|
||||||
|
)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
commit_enriched!(dest_dir, job, meta) -> (file_dest, sidecar_dest)
|
||||||
|
|
||||||
|
Commit an enriched file to `dest_dir` with the sidecar-first ordering so the
|
||||||
|
invariant *"a file in dest_dir implies its sidecar is already there"* always
|
||||||
|
holds. Shared by every enrichment stage that emits a `.meta.json` sidecar
|
||||||
|
(stage-2 known files → `done/`, stage-4 text files → `text_done/`).
|
||||||
|
|
||||||
|
Sequence: write `<name>.meta.json` to a temp name, fsync its bytes, rename it
|
||||||
|
into place, fsync `dest_dir` so the rename itself is durable, THEN move the file
|
||||||
|
into `dest_dir`. A crash between the two leaves only a harmless orphan sidecar in
|
||||||
|
`dest_dir` while the file stays in its stage dir, so stage-aware recovery
|
||||||
|
re-drives it and overwrites the sidecar, idempotently. The fsyncs make the
|
||||||
|
ordering hold across power loss, not just process crashes.
|
||||||
|
"""
|
||||||
|
function commit_enriched!(dest_dir::AbstractString, job::Job, meta)
|
||||||
|
base = basename(job.path)
|
||||||
|
sidecar = joinpath(dest_dir, string(base, ".meta.json"))
|
||||||
|
tmp_sidecar = string(sidecar, ".tmp")
|
||||||
|
|
||||||
|
# Write to a temp name then rename, so a reader in dest_dir never sees a
|
||||||
|
# partial sidecar and a crash mid-write can't masquerade as a committed one.
|
||||||
|
open(tmp_sidecar, "w") do io
|
||||||
|
write(io, JSON3.write(meta))
|
||||||
|
flush(io)
|
||||||
|
fsync_fd(fd(io)) # durably persist bytes before the rename
|
||||||
|
end
|
||||||
|
mv(tmp_sidecar, sidecar; force=true) # sidecar committed first
|
||||||
|
fsync_dir(dest_dir) # persist the rename itself, not just the bytes
|
||||||
|
|
||||||
|
file_dest = move_to(dest_dir, job) # file arrival = commit point
|
||||||
|
return (file_dest, sidecar)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
finalize_known!(cfg, job, meta) -> (file_dest, sidecar_dest)
|
||||||
|
|
||||||
|
Commit an enriched known file (stage 2) to `done/` via the shared sidecar-first
|
||||||
|
`commit_enriched!`.
|
||||||
|
"""
|
||||||
|
finalize_known!(cfg::Config, job::Job, meta) = commit_enriched!(cfg.done_dir, job, meta)
|
||||||
|
|
||||||
@@ -4,8 +4,8 @@
|
|||||||
# drift out of sync (same input layout, same scaling, same net shape).
|
# drift out of sync (same input layout, same scaling, same net shape).
|
||||||
#
|
#
|
||||||
# The including scope is expected to have already done `using Lux` and
|
# The including scope is expected to have already done `using Lux` and
|
||||||
# `using JLD2` (FileServer.jl and bin/train.jl both do), so this file adds no
|
# `using JLD2` (DarkStruct.jl and bin/train.jl both do), so this file adds no
|
||||||
# top-level `using` of its own and works whether it's pulled into the FileServer
|
# top-level `using` of its own and works whether it's pulled into the DarkStruct
|
||||||
# module or a bare script.
|
# module or a bare script.
|
||||||
|
|
||||||
"Number of input features: first 16 bytes + last 16 bytes of a file."
|
"Number of input features: first 16 bytes + last 16 bytes of a file."
|
||||||
|
|||||||
296
src/multipart.jl
Normal file
296
src/multipart.jl
Normal file
@@ -0,0 +1,296 @@
|
|||||||
|
# Streaming multipart/form-data reader.
|
||||||
|
#
|
||||||
|
# Why this exists: HTTP.jl's `parse_multipart_form` takes the *complete* request
|
||||||
|
# body as a byte vector, so using it means every file in the request sits in
|
||||||
|
# memory at once, and is then copied again per part. That contradicts the whole
|
||||||
|
# point of this service: file bytes belong on disk, and only a small reference
|
||||||
|
# travels through the queue. So intake needs a parser that never holds a file.
|
||||||
|
#
|
||||||
|
# This reader walks the body incrementally: it pulls fixed-size chunks off the
|
||||||
|
# socket and hands each part's bytes straight to a sink (the spool file). Peak
|
||||||
|
# memory per connection is `chunk_bytes` + the boundary length, regardless of how
|
||||||
|
# large, or how many, the uploaded files are.
|
||||||
|
#
|
||||||
|
# Interface: two calls in a loop, so the caller keeps ordinary control flow
|
||||||
|
# rather than inverting into callbacks.
|
||||||
|
#
|
||||||
|
# r = MultipartReader(io, boundary)
|
||||||
|
# while (part = next_part!(r)) !== nothing
|
||||||
|
# part.filename === nothing ? skip_part_body!(r) : write_part_body!(sink, r)
|
||||||
|
# end
|
||||||
|
#
|
||||||
|
# The grammar it implements (RFC 2046 §5.1, RFC 7578):
|
||||||
|
#
|
||||||
|
# [preamble] "--" boundary CRLF
|
||||||
|
# part-headers CRLF CRLF part-body
|
||||||
|
# CRLF "--" boundary CRLF ... another part ...
|
||||||
|
# CRLF "--" boundary "--" CRLF ... end of form, [epilogue]
|
||||||
|
#
|
||||||
|
# So the delimiter that *closes* a body is CRLF + "--" + boundary, and the two
|
||||||
|
# bytes after it say whether another part follows (CRLF) or the form is over
|
||||||
|
# ("--"). Every read is bounded, and the buffer retains only the last
|
||||||
|
# `length(delimiter)-1` bytes when no delimiter is found; that tail is what
|
||||||
|
# makes a delimiter split across two chunks parse correctly.
|
||||||
|
|
||||||
|
"Default socket read size, and therefore the memory bound per in-flight upload."
|
||||||
|
const UPLOAD_CHUNK_BYTES = 64 * 1024
|
||||||
|
|
||||||
|
"""
|
||||||
|
A part header block bigger than this is abuse, not a filename. Bounding it keeps
|
||||||
|
the one genuinely unbounded-looking read (headers, which must be buffered whole
|
||||||
|
to be parsed) from being a memory hole.
|
||||||
|
"""
|
||||||
|
const MAX_PART_HEADER_BYTES = 16 * 1024
|
||||||
|
|
||||||
|
const CRLF = UInt8[0x0d, 0x0a]
|
||||||
|
const CRLFCRLF = UInt8[0x0d, 0x0a, 0x0d, 0x0a]
|
||||||
|
const DASHDASH = UInt8[0x2d, 0x2d]
|
||||||
|
|
||||||
|
"A malformed (or truncated) multipart body. Callers turn this into a 400."
|
||||||
|
struct MultipartError <: Exception
|
||||||
|
msg::String
|
||||||
|
end
|
||||||
|
|
||||||
|
Base.showerror(io::IO, e::MultipartError) = print(io, "MultipartError: ", e.msg)
|
||||||
|
|
||||||
|
"What a part's headers said about it. `filename === nothing` means a plain form field, not a file."
|
||||||
|
struct MultipartPart
|
||||||
|
name::Union{String,Nothing}
|
||||||
|
filename::Union{String,Nothing}
|
||||||
|
content_type::Union{String,Nothing}
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
MultipartReader(io, boundary; chunk_bytes = UPLOAD_CHUNK_BYTES)
|
||||||
|
|
||||||
|
An incremental reader over the multipart body arriving on `io`. `boundary` is the
|
||||||
|
value from the request's `Content-Type` header (see [`multipart_boundary`](@ref)).
|
||||||
|
"""
|
||||||
|
mutable struct MultipartReader{I<:IO}
|
||||||
|
io::I
|
||||||
|
dash_boundary::Vector{UInt8} # "--" boundary: opens the first part
|
||||||
|
delimiter::Vector{UInt8} # CRLF "--" boundary: closes every part
|
||||||
|
buf::Vector{UInt8} # rolling window; bounded by chunk_bytes + delimiter
|
||||||
|
pos::Int # next unconsumed index in buf
|
||||||
|
scratch::Vector{UInt8} # reused socket read target, so chunks don't churn the GC
|
||||||
|
chunk_bytes::Int
|
||||||
|
state::Symbol # :preamble | :at_delimiter | :body | :done
|
||||||
|
end
|
||||||
|
|
||||||
|
function MultipartReader(io::IO, boundary::AbstractString;
|
||||||
|
chunk_bytes::Int = UPLOAD_CHUNK_BYTES)
|
||||||
|
chunk_bytes > 0 || throw(ArgumentError("chunk_bytes must be positive"))
|
||||||
|
isempty(boundary) && throw(MultipartError("empty multipart boundary"))
|
||||||
|
dash_boundary = Vector{UInt8}(codeunits(string("--", boundary)))
|
||||||
|
return MultipartReader(io, dash_boundary, vcat(CRLF, dash_boundary),
|
||||||
|
UInt8[], 1, Vector{UInt8}(undef, chunk_bytes),
|
||||||
|
chunk_bytes, :preamble)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
multipart_boundary(content_type) -> String | nothing
|
||||||
|
|
||||||
|
Pull the boundary out of a `multipart/form-data` Content-Type header. Returns
|
||||||
|
`nothing` if the header is missing, is some other media type, or has no boundary.
|
||||||
|
All of those are the same 400 to a caller.
|
||||||
|
"""
|
||||||
|
function multipart_boundary(content_type::Union{AbstractString,Nothing})
|
||||||
|
content_type === nothing && return nothing
|
||||||
|
occursin(r"^\s*multipart/form-data"i, content_type) || return nothing
|
||||||
|
m = match(r"(?i:\bboundary)=(?:\"([^\"]+)\"|([^\s;]+))", content_type)
|
||||||
|
m === nothing && return nothing
|
||||||
|
return String(something(m[1], m[2]))
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------ buffer plumbing
|
||||||
|
|
||||||
|
"Unconsumed bytes currently buffered."
|
||||||
|
navail(r::MultipartReader) = length(r.buf) - r.pos + 1
|
||||||
|
|
||||||
|
"Drop already-consumed bytes so the buffer stays bounded across a long body."
|
||||||
|
function compact!(r::MultipartReader)
|
||||||
|
r.pos == 1 && return nothing
|
||||||
|
n = navail(r)
|
||||||
|
n > 0 && copyto!(r.buf, 1, r.buf, r.pos, n)
|
||||||
|
resize!(r.buf, max(n, 0))
|
||||||
|
r.pos = 1
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Pull one more chunk off the wire, returning `false` at end of body.
|
||||||
|
|
||||||
|
`readbytes!` on an `HTTP.Stream` returns at most what remains of the current
|
||||||
|
content-length or chunk, so this is bounded by `chunk_bytes`; `eof` is what
|
||||||
|
advances a chunked-encoded body to its next chunk, hence the guard.
|
||||||
|
"""
|
||||||
|
function fill_more!(r::MultipartReader)
|
||||||
|
eof(r.io) && return false
|
||||||
|
n = readbytes!(r.io, r.scratch, r.chunk_bytes)
|
||||||
|
n == 0 && return false
|
||||||
|
append!(r.buf, view(r.scratch, 1:n))
|
||||||
|
return true
|
||||||
|
end
|
||||||
|
|
||||||
|
"Buffer until `needle` is found, returning its range, or `nothing` at end of body."
|
||||||
|
function seek_needle!(r::MultipartReader, needle::Vector{UInt8}; limit::Int = 0)
|
||||||
|
while true
|
||||||
|
idx = findnext(needle, r.buf, r.pos)
|
||||||
|
idx === nothing || return idx
|
||||||
|
# Only the last length(needle)-1 bytes can still be part of a match, but
|
||||||
|
# the caller may need the skipped bytes (a part body), so trimming is the
|
||||||
|
# caller's job; we only enforce the optional limit.
|
||||||
|
limit > 0 && navail(r) > limit &&
|
||||||
|
throw(MultipartError("no delimiter within $limit bytes"))
|
||||||
|
compact!(r)
|
||||||
|
fill_more!(r) || return nothing
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"Ensure at least `n` bytes are buffered; `false` if the body ended first."
|
||||||
|
function ensure!(r::MultipartReader, n::Int)
|
||||||
|
while navail(r) < n
|
||||||
|
compact!(r)
|
||||||
|
fill_more!(r) || return false
|
||||||
|
end
|
||||||
|
return true
|
||||||
|
end
|
||||||
|
|
||||||
|
# Write `r.buf[range]` to `sink`. Goes through `unsafe_write` because
|
||||||
|
# `write(io, ::SubArray{UInt8})` falls back to a byte-at-a-time loop in Base,
|
||||||
|
# which would dominate the cost of a large upload.
|
||||||
|
function emit!(sink::IO, r::MultipartReader, from::Int, to::Int)
|
||||||
|
n = to - from + 1
|
||||||
|
n <= 0 && return 0
|
||||||
|
buf = r.buf # GC.@preserve needs a plain symbol, not a field access
|
||||||
|
GC.@preserve buf unsafe_write(sink, pointer(buf, from), UInt(n))
|
||||||
|
return n
|
||||||
|
end
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------- parts
|
||||||
|
|
||||||
|
"""
|
||||||
|
next_part!(r) -> MultipartPart | nothing
|
||||||
|
|
||||||
|
Advance to the next part and return its headers, or `nothing` at the end of the
|
||||||
|
form. The previous part's body must have been consumed first (with
|
||||||
|
[`write_part_body!`](@ref) or [`skip_part_body!`](@ref)). The reader cannot skip
|
||||||
|
a body it hasn't been told to, because the body is only bounded by finding the
|
||||||
|
next delimiter.
|
||||||
|
"""
|
||||||
|
function next_part!(r::MultipartReader)
|
||||||
|
r.state === :done && return nothing
|
||||||
|
r.state === :body &&
|
||||||
|
throw(MultipartError("the current part's body must be consumed before the next part"))
|
||||||
|
|
||||||
|
if r.state === :preamble
|
||||||
|
# Discard the preamble (RFC says ignore it) and consume the opening
|
||||||
|
# delimiter. Bounded: real clients send no preamble at all, and an
|
||||||
|
# unbounded scan here would be a way to make us buffer a whole body.
|
||||||
|
idx = seek_needle!(r, r.dash_boundary; limit = MAX_PART_HEADER_BYTES)
|
||||||
|
idx === nothing && throw(MultipartError("no multipart boundary found in body"))
|
||||||
|
r.pos = last(idx) + 1
|
||||||
|
r.state = :at_delimiter
|
||||||
|
end
|
||||||
|
|
||||||
|
# Just after a delimiter: "--" ends the form, CRLF introduces another part.
|
||||||
|
ensure!(r, 2) || throw(MultipartError("truncated body after a boundary delimiter"))
|
||||||
|
if view(r.buf, r.pos:r.pos+1) == DASHDASH
|
||||||
|
r.pos += 2
|
||||||
|
r.state = :done
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
skip_linear_whitespace!(r)
|
||||||
|
ensure!(r, 2) || throw(MultipartError("truncated body after a boundary delimiter"))
|
||||||
|
view(r.buf, r.pos:r.pos+1) == CRLF ||
|
||||||
|
throw(MultipartError("boundary delimiter is not followed by a line ending"))
|
||||||
|
r.pos += 2
|
||||||
|
|
||||||
|
part = read_part_headers!(r)
|
||||||
|
r.state = :body
|
||||||
|
return part
|
||||||
|
end
|
||||||
|
|
||||||
|
"RFC 2046 allows spaces/tabs between the delimiter and its line ending."
|
||||||
|
function skip_linear_whitespace!(r::MultipartReader)
|
||||||
|
while ensure!(r, 1) && (r.buf[r.pos] == 0x20 || r.buf[r.pos] == 0x09)
|
||||||
|
r.pos += 1
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
function read_part_headers!(r::MultipartReader)
|
||||||
|
# A part with no headers at all is `CRLF CRLF body`: the empty line comes
|
||||||
|
# immediately, so searching for CRLFCRLF would run past it into the body.
|
||||||
|
if ensure!(r, 2) && view(r.buf, r.pos:r.pos+1) == CRLF
|
||||||
|
r.pos += 2
|
||||||
|
return MultipartPart(nothing, nothing, nothing)
|
||||||
|
end
|
||||||
|
# The `limit` here bounds *buffering*, and only fires when the headers span
|
||||||
|
# chunks. The explicit length check below is the actual policy, so the rule
|
||||||
|
# doesn't depend on how the body happened to be chunked on the wire.
|
||||||
|
idx = seek_needle!(r, CRLFCRLF; limit = MAX_PART_HEADER_BYTES)
|
||||||
|
idx === nothing && throw(MultipartError("truncated body inside a part's headers"))
|
||||||
|
first(idx) - r.pos > MAX_PART_HEADER_BYTES &&
|
||||||
|
throw(MultipartError("part headers exceed $MAX_PART_HEADER_BYTES bytes"))
|
||||||
|
# Copying is fine: the check above bounds this block.
|
||||||
|
block = String(r.buf[r.pos:first(idx)-1])
|
||||||
|
r.pos = last(idx) + 1
|
||||||
|
return parse_part_headers(block)
|
||||||
|
end
|
||||||
|
|
||||||
|
"Unescape the backslash escapes RFC 2045 allows inside a quoted-string."
|
||||||
|
unquote(s::AbstractString) = replace(s, r"\\(.)" => s"\1")
|
||||||
|
|
||||||
|
function parse_part_headers(block::AbstractString)
|
||||||
|
name = filename = content_type = nothing
|
||||||
|
for line in eachsplit(block, "\r\n")
|
||||||
|
colon = findfirst(':', line)
|
||||||
|
colon === nothing && continue
|
||||||
|
key = lowercase(strip(line[1:colon-1]))
|
||||||
|
value = strip(line[colon+1:end])
|
||||||
|
if key == "content-disposition"
|
||||||
|
# `\b` matters: it keeps the `name=` pattern from matching inside `filename=`.
|
||||||
|
m = match(r"(?i:\bname)=(?:\"((?:[^\"\\]|\\.)*)\"|([^\s;]+))", value)
|
||||||
|
m === nothing || (name = unquote(String(something(m[1], m[2]))))
|
||||||
|
m = match(r"(?i:\bfilename)=(?:\"((?:[^\"\\]|\\.)*)\"|([^\s;]+))", value)
|
||||||
|
m === nothing || (filename = unquote(String(something(m[1], m[2]))))
|
||||||
|
elseif key == "content-type"
|
||||||
|
content_type = String(value)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return MultipartPart(name, filename, content_type)
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
write_part_body!(sink, r) -> Int
|
||||||
|
|
||||||
|
Stream the current part's body into `sink`, returning the number of bytes
|
||||||
|
written. Nothing larger than a chunk is ever held in memory.
|
||||||
|
"""
|
||||||
|
function write_part_body!(sink::IO, r::MultipartReader)
|
||||||
|
r.state === :body || throw(MultipartError("no part body is open"))
|
||||||
|
total = 0
|
||||||
|
keep = length(r.delimiter) - 1 # a delimiter may straddle two chunks
|
||||||
|
while true
|
||||||
|
idx = findnext(r.delimiter, r.buf, r.pos)
|
||||||
|
if idx !== nothing
|
||||||
|
total += emit!(sink, r, r.pos, first(idx) - 1)
|
||||||
|
r.pos = last(idx) + 1
|
||||||
|
r.state = :at_delimiter
|
||||||
|
return total
|
||||||
|
end
|
||||||
|
# Emit only what cannot be the start of a straddling delimiter, then
|
||||||
|
# keep that tail and read more.
|
||||||
|
emit_to = length(r.buf) - keep
|
||||||
|
if emit_to >= r.pos
|
||||||
|
total += emit!(sink, r, r.pos, emit_to)
|
||||||
|
r.pos = emit_to + 1
|
||||||
|
end
|
||||||
|
compact!(r)
|
||||||
|
fill_more!(r) || throw(MultipartError("truncated body inside a part"))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"Consume and discard the current part's body (a form field, or a file we can't take)."
|
||||||
|
skip_part_body!(r::MultipartReader) = write_part_body!(devnull, r)
|
||||||
74
src/queue.jl
74
src/queue.jl
@@ -1,13 +1,36 @@
|
|||||||
# The queue seam.
|
# The queue seam.
|
||||||
#
|
#
|
||||||
# The rest of the app only ever calls `enqueue!`, `dequeue!`, and `close!`.
|
# The rest of the app only ever calls `enqueue!`, `dequeue!`, `ack!`, `nack!`
|
||||||
# Today those are backed by an in-process, bounded, thread-safe buffer
|
# and `close!`. Two implementations sit behind those five methods:
|
||||||
# (the Go-channel / Julia-`Channel` model). To move to RabbitMQ (or any broker)
|
#
|
||||||
# later, implement a new `JobQueue` subtype with these three methods and swap
|
# ChannelQueue in-process, bounded, thread-safe (the Go-channel / Julia-
|
||||||
# the construction in `run` — no HTTP handler or worker code needs to change.
|
# `Channel` model). Lost on crash: everything still in `spool/`
|
||||||
|
# is re-driven from stage 1 by `recover_dir!`.
|
||||||
|
# RabbitQueue durable RabbitMQ queues (src/rabbit.jl). Survives a crash: a
|
||||||
|
# job stays unacked until its handler commits, so a restart
|
||||||
|
# resumes each file at the stage it had actually reached.
|
||||||
|
#
|
||||||
|
# `FS_QUEUE_BACKEND` picks one; `run` constructs accordingly. No HTTP handler or
|
||||||
|
# worker code knows which it got.
|
||||||
|
#
|
||||||
|
# The ack pair is what makes the broker worth having. Without it a delivery is
|
||||||
|
# settled the moment it is handed to a worker, so a crash mid-handler loses the
|
||||||
|
# job exactly as the in-process queue does, and the network hop buys nothing.
|
||||||
|
# `ChannelQueue` implements both as no-ops, which is honest rather than lazy:
|
||||||
|
# an in-process queue has no delivery to settle, and its recovery story is
|
||||||
|
# `recover_dir!`.
|
||||||
|
|
||||||
abstract type JobQueue end
|
abstract type JobQueue end
|
||||||
|
|
||||||
|
# How long a producer backs off before retrying an enqueue onto a full queue.
|
||||||
|
# Blocking backpressure: a file that is already on disk is never dropped, so the
|
||||||
|
# producer parks until the consumer makes room. Used by the stage-1 and stage-3
|
||||||
|
# routing handoffs and by startup recovery (see `enqueue_blocking!`). Intake is
|
||||||
|
# deliberately *not* a user: the HTTP path's `enqueue!` stays non-blocking and
|
||||||
|
# turns a full queue into a 503, so a saturated pipeline slows uploads down
|
||||||
|
# rather than holding request threads hostage.
|
||||||
|
const ROUTE_ENQUEUE_RETRY_SECONDS = 0.05
|
||||||
|
|
||||||
"""
|
"""
|
||||||
enqueue!(q, job) -> Bool
|
enqueue!(q, job) -> Bool
|
||||||
|
|
||||||
@@ -21,10 +44,35 @@ function enqueue! end
|
|||||||
dequeue!(q) -> Union{Job,Nothing}
|
dequeue!(q) -> Union{Job,Nothing}
|
||||||
|
|
||||||
Blocks until a job is available and returns it. Returns `nothing` only when the
|
Blocks until a job is available and returns it. Returns `nothing` only when the
|
||||||
queue has been closed *and* fully drained — the signal for a worker to exit.
|
queue has been closed *and* fully drained, which is a worker's signal to exit.
|
||||||
"""
|
"""
|
||||||
function dequeue! end
|
function dequeue! end
|
||||||
|
|
||||||
|
"""
|
||||||
|
ack!(q, job)
|
||||||
|
|
||||||
|
Settle `job` as done: the broker may forget it. Called by `worker_loop` after
|
||||||
|
the handler returns, and after a quarantine, so the message outlives the process
|
||||||
|
for exactly as long as the work is unfinished.
|
||||||
|
|
||||||
|
No-op for `ChannelQueue`.
|
||||||
|
"""
|
||||||
|
function ack! end
|
||||||
|
|
||||||
|
"""
|
||||||
|
nack!(q, job)
|
||||||
|
|
||||||
|
Settle `job` as rejected, without requeueing. Requeueing is deliberately not
|
||||||
|
offered: a job that failed on its bytes will fail again, and redelivering it is
|
||||||
|
a poison-message loop. With no dead-letter exchange configured the broker simply
|
||||||
|
discards it, so this differs from `ack!` only in what it says, not in what
|
||||||
|
happens — which is the point at the one call site that uses it (see
|
||||||
|
`worker_loop`).
|
||||||
|
|
||||||
|
No-op for `ChannelQueue`.
|
||||||
|
"""
|
||||||
|
function nack! end
|
||||||
|
|
||||||
"""
|
"""
|
||||||
close!(q)
|
close!(q)
|
||||||
|
|
||||||
@@ -70,6 +118,10 @@ function dequeue!(q::ChannelQueue)::Union{Job,Nothing}
|
|||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
|
||||||
|
# Nothing to settle: an in-process job was never a delivery.
|
||||||
|
ack!(::ChannelQueue, ::Job) = nothing
|
||||||
|
nack!(::ChannelQueue, ::Job) = nothing
|
||||||
|
|
||||||
function close!(q::ChannelQueue)
|
function close!(q::ChannelQueue)
|
||||||
lock(q.cond)
|
lock(q.cond)
|
||||||
try
|
try
|
||||||
@@ -81,6 +133,16 @@ function close!(q::ChannelQueue)
|
|||||||
return nothing
|
return nothing
|
||||||
end
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
capacity(q) -> Int
|
||||||
|
|
||||||
|
How many jobs the queue can hold before `enqueue!` starts refusing. Part of the
|
||||||
|
introspection seam alongside `length`: `/stats` reports depth against capacity,
|
||||||
|
because a depth of 900 means nothing without knowing whether the limit is 1000
|
||||||
|
or 1_000_000.
|
||||||
|
"""
|
||||||
|
capacity(q::ChannelQueue) = q.capacity
|
||||||
|
|
||||||
"Number of jobs currently buffered (for logging/introspection)."
|
"Number of jobs currently buffered (for logging/introspection)."
|
||||||
function Base.length(q::ChannelQueue)
|
function Base.length(q::ChannelQueue)
|
||||||
lock(q.cond)
|
lock(q.cond)
|
||||||
|
|||||||
512
src/rabbit.jl
Normal file
512
src/rabbit.jl
Normal file
@@ -0,0 +1,512 @@
|
|||||||
|
# RabbitMQ-backed `JobQueue` (src/queue.jl defines the seam this implements).
|
||||||
|
#
|
||||||
|
# What it buys, and what it doesn't
|
||||||
|
# ---------------------------------
|
||||||
|
# The in-process `ChannelQueue` dies with the process, so a crash leaves no
|
||||||
|
# record of how far each file had got and `recover_dir!` can only re-drive
|
||||||
|
# everything in `spool/` from stage 1. Durable queues plus per-job acks replace
|
||||||
|
# that guess with a fact: a job stays unacked until its handler commits, so a
|
||||||
|
# restart resumes each file at the stage it had actually reached, and no file is
|
||||||
|
# processed from scratch just because the process died.
|
||||||
|
#
|
||||||
|
# It is *at-least-once*, not exactly-once. Stages 1 and 3 publish downstream and
|
||||||
|
# then ack upstream (`worker_loop`), so a crash in that window redelivers a job
|
||||||
|
# that was already routed, and the file gets processed twice. That is safe here
|
||||||
|
# for the same reason replay was: classification and the UTF-8 sniff are pure
|
||||||
|
# functions of the file's bytes, and every terminal commit renames with
|
||||||
|
# `force=true`. The one new wrinkle is that a duplicate can now run *concurrently*
|
||||||
|
# with the original, and the loser finds the file already moved; `worker_loop`
|
||||||
|
# recognises that (a vanished `job.path`) and settles it quietly instead of
|
||||||
|
# quarantining a file that in fact succeeded.
|
||||||
|
#
|
||||||
|
# It also does not buy horizontal scale. `Job.path` is a local filesystem path,
|
||||||
|
# so a second consumer process on another host would be handed jobs whose files
|
||||||
|
# it cannot see. See "Queue backends" in the README.
|
||||||
|
#
|
||||||
|
# Shape
|
||||||
|
# -----
|
||||||
|
# One connection for the whole process, and per stage queue:
|
||||||
|
#
|
||||||
|
# con_chan consumes, with `basic_qos` prefetch bounding unacked messages,
|
||||||
|
# and carries this stage's acks (a delivery tag is only meaningful
|
||||||
|
# on the channel that delivered it)
|
||||||
|
# pub_chan publishes *into* this queue, from whichever stage feeds it
|
||||||
|
# buffer a bounded `Channel{Job}` bridging AMQPClient's push consumer to
|
||||||
|
# the blocking `dequeue!` that `worker_loop` expects
|
||||||
|
#
|
||||||
|
# plus one shared admin channel for the depth polls. Publishing and acking from
|
||||||
|
# many worker tasks at once is safe: AMQPClient serializes multi-frame publishes
|
||||||
|
# under a per-channel send lock and queues frames onto a thread-safe channel.
|
||||||
|
|
||||||
|
using AMQPClient
|
||||||
|
|
||||||
|
# The stage queues, in `STAGE_KEYS` order. The suffix is what appears after
|
||||||
|
# `FS_AMQP_PREFIX`, so a broker's queue list reads like the pipeline.
|
||||||
|
const AMQP_QUEUE_SUFFIXES = (classify = "classify", enrich = "enrich",
|
||||||
|
triage = "triage", language = "language")
|
||||||
|
|
||||||
|
# Startup connect: retry a broker that isn't up yet (compose starts the app long
|
||||||
|
# before RabbitMQ is ready) but give up rather than hang, so a genuinely
|
||||||
|
# misconfigured URL is a startup failure and not a silent hang.
|
||||||
|
const AMQP_CONNECT_RETRY_SECONDS = 30.0
|
||||||
|
const AMQP_CONNECT_RETRY_INTERVAL = 2.0
|
||||||
|
|
||||||
|
# How often the depth poller asks the broker for each queue's ready-message
|
||||||
|
# count. This is what `/stats` reports and what the advisory capacity check
|
||||||
|
# compares against, so it trades staleness for one round trip per queue per
|
||||||
|
# second instead of one per file.
|
||||||
|
const AMQP_DEPTH_POLL_SECONDS = 1.0
|
||||||
|
|
||||||
|
# A publisher confirm that never arrives must not park an HTTP request thread
|
||||||
|
# forever. The poller doubles as the watchdog (see `poll_loop!`).
|
||||||
|
const AMQP_CONFIRM_TIMEOUT_SECONDS = 10.0
|
||||||
|
|
||||||
|
"Connection parameters parsed out of an `amqp://` URL."
|
||||||
|
struct AMQPTarget
|
||||||
|
host::String
|
||||||
|
port::Int
|
||||||
|
virtualhost::String
|
||||||
|
login::String
|
||||||
|
password::String
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
parse_amqp_url(url) -> AMQPTarget
|
||||||
|
|
||||||
|
Parse `amqp://user:pass@host:port/vhost`. Every component is optional and falls
|
||||||
|
back to the AMQP defaults, except that we do *not* default the credentials to
|
||||||
|
`guest`/`guest`: those only authenticate over loopback in a stock RabbitMQ, so
|
||||||
|
inheriting them across a container network would fail at connect time with an
|
||||||
|
authentication error that says nothing about the real mistake.
|
||||||
|
|
||||||
|
The vhost is the path with its leading `/` removed and percent-escapes decoded,
|
||||||
|
so the conventional default vhost `/` can be written either as a bare trailing
|
||||||
|
slash or as the `%2F` form.
|
||||||
|
"""
|
||||||
|
function parse_amqp_url(url::AbstractString)
|
||||||
|
uri = HTTP.URI(url)
|
||||||
|
scheme = lowercase(uri.scheme)
|
||||||
|
scheme in ("amqp", "amqps") ||
|
||||||
|
throw(ArgumentError("FS_AMQP_URL must be an amqp:// or amqps:// URL, got \"$url\""))
|
||||||
|
scheme == "amqps" &&
|
||||||
|
throw(ArgumentError("amqps:// (TLS) is not supported yet; use amqp:// on a trusted network"))
|
||||||
|
|
||||||
|
userinfo = split(uri.userinfo, ':'; limit = 2)
|
||||||
|
login = isempty(userinfo[1]) ? "guest" : HTTP.unescapeuri(userinfo[1])
|
||||||
|
password = length(userinfo) < 2 ? "guest" : HTTP.unescapeuri(userinfo[2])
|
||||||
|
|
||||||
|
path = lstrip(uri.path, '/')
|
||||||
|
vhost = isempty(path) ? "/" : HTTP.unescapeuri(path)
|
||||||
|
|
||||||
|
return AMQPTarget(isempty(uri.host) ? "localhost" : uri.host,
|
||||||
|
isempty(uri.port) ? AMQPClient.AMQP_DEFAULT_PORT : parse(Int, uri.port),
|
||||||
|
vhost, login, password)
|
||||||
|
end
|
||||||
|
|
||||||
|
"The process's single broker connection, plus the shared channel depth polls use."
|
||||||
|
mutable struct RabbitBackend
|
||||||
|
conn::AMQPClient.Connection
|
||||||
|
admin_chan::AMQPClient.MessageChannel
|
||||||
|
const admin_lock::ReentrantLock # one synchronous RPC on admin_chan at a time
|
||||||
|
const prefix::String
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
RabbitQueue(backend, name, capacity, prefetch; confirms = false)
|
||||||
|
|
||||||
|
One durable queue. `capacity` is *advisory*: `enqueue!` compares it against a
|
||||||
|
depth that is polled once a second and adjusted locally in between, so a burst
|
||||||
|
can overshoot by up to a poll interval before the queue starts refusing. The
|
||||||
|
alternative — `x-max-length` with `overflow: reject-publish` — needs a publisher
|
||||||
|
confirm per message to learn about the rejection, which would put a round trip
|
||||||
|
on an inter-stage handoff that currently costs 0.12 µs.
|
||||||
|
|
||||||
|
`confirms` puts this queue's *publish* channel into confirm mode, so `enqueue!`
|
||||||
|
returns only once the broker has accepted the message. It is set on the stage-1
|
||||||
|
queue alone, because intake is the only publisher that answers to someone: a
|
||||||
|
`202` tells a client its file is safe, and that client may delete its copy.
|
||||||
|
Inter-stage publishes stay fire-and-forget, where a lost message leaves its
|
||||||
|
source job unacked and redelivery repairs it for free.
|
||||||
|
"""
|
||||||
|
mutable struct RabbitQueue <: JobQueue
|
||||||
|
const backend::RabbitBackend
|
||||||
|
const name::String
|
||||||
|
const capacity::Int
|
||||||
|
const buffer::Channel{Job} # consumer callback → dequeue!
|
||||||
|
const pub_chan::AMQPClient.MessageChannel
|
||||||
|
const con_chan::AMQPClient.MessageChannel
|
||||||
|
const depth::Threads.Atomic{Int} # last polled ready count, locally adjusted
|
||||||
|
const closed::Threads.Atomic{Bool}
|
||||||
|
const confirms::Bool
|
||||||
|
const pub_lock::ReentrantLock # serializes publish+confirm on pub_chan
|
||||||
|
const confirm_ch::Channel{Bool} # broker's verdict on the outstanding publish
|
||||||
|
const pending_since::Threads.Atomic{Float64} # when it was published (0.0 = nothing pending)
|
||||||
|
const on_lost::Function # called once if the connection dies
|
||||||
|
consumer_tag::String
|
||||||
|
poller::Union{Task,Nothing}
|
||||||
|
end
|
||||||
|
|
||||||
|
capacity(q::RabbitQueue) = q.capacity
|
||||||
|
|
||||||
|
"""
|
||||||
|
Depth as of the last poll, adjusted by our own publishes and takes since.
|
||||||
|
|
||||||
|
The broker's number counts *ready* messages, so it excludes the ones already
|
||||||
|
handed to workers and not yet acked. That is the right number for `/stats`:
|
||||||
|
depth is meant to say how much work is piled up in front of a stage, and the
|
||||||
|
messages in its workers' hands are not piled up, they are being worked.
|
||||||
|
"""
|
||||||
|
Base.length(q::RabbitQueue) = max(0, q.depth[])
|
||||||
|
|
||||||
|
# --- wire format ------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
job_json(job) -> String
|
||||||
|
|
||||||
|
Serialize a `Job` for the wire. `delivery_tag` is deliberately absent: it names a
|
||||||
|
delivery on one channel of one connection, not the job, and a redelivery gets a
|
||||||
|
fresh one.
|
||||||
|
"""
|
||||||
|
job_json(job::Job) = JSON3.write((; job.id, job.original_name, job.path,
|
||||||
|
job.size, job.received_at))
|
||||||
|
|
||||||
|
"Rebuild a `Job` from a message body, carrying the tag it was delivered under."
|
||||||
|
function job_from_message(msg::AMQPClient.Message)
|
||||||
|
o = JSON3.read(String(copy(msg.data)))
|
||||||
|
return Job(String(o.id), String(o.original_name), String(o.path),
|
||||||
|
Int(o.size), Float64(o.received_at), Int64(msg.delivery_tag))
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- connect / declare ------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
connect_backend(cfg) -> RabbitBackend
|
||||||
|
|
||||||
|
Open the process's connection, retrying a broker that is not up yet for up to
|
||||||
|
`AMQP_CONNECT_RETRY_SECONDS`. Compose starts this container long before RabbitMQ
|
||||||
|
finishes booting even with a healthcheck gate, so a first refusal is routine; a
|
||||||
|
persistent one is a startup error, because falling back to the in-process queue
|
||||||
|
would silently deliver none of the durability that was asked for.
|
||||||
|
"""
|
||||||
|
function connect_backend(cfg::Config)
|
||||||
|
t = parse_amqp_url(cfg.amqp_url)
|
||||||
|
auth = Dict{String,Any}("MECHANISM" => "AMQPLAIN",
|
||||||
|
"LOGIN" => t.login, "PASSWORD" => t.password)
|
||||||
|
deadline = time() + AMQP_CONNECT_RETRY_SECONDS
|
||||||
|
attempt = 0
|
||||||
|
while true
|
||||||
|
attempt += 1
|
||||||
|
try
|
||||||
|
conn = AMQPClient.connection(; host = t.host, port = t.port,
|
||||||
|
virtualhost = t.virtualhost,
|
||||||
|
auth_params = auth)
|
||||||
|
admin = AMQPClient.channel(conn, AMQPClient.UNUSED_CHANNEL, true)
|
||||||
|
@info "connected to broker" host=t.host port=t.port vhost=t.virtualhost attempts=attempt
|
||||||
|
return RabbitBackend(conn, admin, ReentrantLock(), cfg.amqp_prefix)
|
||||||
|
catch e
|
||||||
|
time() < deadline || rethrow()
|
||||||
|
@info "broker not ready, retrying" host=t.host port=t.port attempt=attempt
|
||||||
|
sleep(AMQP_CONNECT_RETRY_INTERVAL)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
open_queue(backend, cfg, key, capacity, workers) -> RabbitQueue
|
||||||
|
|
||||||
|
Declare one stage's durable queue and start consuming from it.
|
||||||
|
|
||||||
|
Prefetch defaults to the stage's worker count: enough that every worker can hold
|
||||||
|
a job, and no more, so a crash redelivers the smallest possible set. The bridge
|
||||||
|
`Channel` is sized to match, since it is exactly the set of messages the broker
|
||||||
|
has handed us and we have not yet acked.
|
||||||
|
"""
|
||||||
|
function open_queue(backend::RabbitBackend, cfg::Config, key::Symbol,
|
||||||
|
capacity::Int, workers::Int; on_lost::Function = () -> nothing)
|
||||||
|
name = string(backend.prefix, ".", AMQP_QUEUE_SUFFIXES[key])
|
||||||
|
prefetch = cfg.amqp_prefetch > 0 ? cfg.amqp_prefetch : max(1, workers)
|
||||||
|
|
||||||
|
con_chan = AMQPClient.channel(backend.conn, AMQPClient.UNUSED_CHANNEL, true)
|
||||||
|
pub_chan = AMQPClient.channel(backend.conn, AMQPClient.UNUSED_CHANNEL, true)
|
||||||
|
|
||||||
|
# durable: the queue definition survives a broker restart. Combined with
|
||||||
|
# PERSISTENT messages below, so do the jobs in it.
|
||||||
|
AMQPClient.queue_declare(con_chan, name; durable = true)
|
||||||
|
AMQPClient.basic_qos(con_chan, 0, prefetch, false)
|
||||||
|
|
||||||
|
confirms = cfg.amqp_confirms && key === :classify
|
||||||
|
q = RabbitQueue(backend, name, capacity, Channel{Job}(prefetch),
|
||||||
|
pub_chan, con_chan, Threads.Atomic{Int}(0),
|
||||||
|
Threads.Atomic{Bool}(false), confirms, ReentrantLock(),
|
||||||
|
Channel{Bool}(1), Threads.Atomic{Float64}(0.0), on_lost, "", nothing)
|
||||||
|
|
||||||
|
if confirms
|
||||||
|
confirm_select!(pub_chan)
|
||||||
|
# The broker reports each publish's fate asynchronously by sequence
|
||||||
|
# number. We hold `pub_lock` across publish-and-wait, so at most one
|
||||||
|
# publish is ever outstanding and any verdict that arrives is this one's.
|
||||||
|
AMQPClient.handle(pub_chan, :Basic, :Ack, (_c, _m, _x) -> settle_confirm!(q, true))
|
||||||
|
AMQPClient.handle(pub_chan, :Basic, :Nack, (_c, _m, _x) -> settle_confirm!(q, false))
|
||||||
|
end
|
||||||
|
|
||||||
|
ok, tag = AMQPClient.basic_consume(con_chan, name, msg -> deliver!(q, msg))
|
||||||
|
ok || error("could not start consuming from queue $name")
|
||||||
|
q.consumer_tag = tag
|
||||||
|
q.poller = Threads.@spawn poll_loop!(q)
|
||||||
|
|
||||||
|
@info "queue ready" queue=name prefetch=prefetch capacity=capacity confirms=confirms
|
||||||
|
return q
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
open_queues(backend, cfg) -> NamedTuple
|
||||||
|
|
||||||
|
The four stage queues, keyed by `STAGE_KEYS` so they line up with the per-stage
|
||||||
|
worker counts, capacities and metrics without a second mapping to keep in sync.
|
||||||
|
"""
|
||||||
|
function open_queues(backend::RabbitBackend, cfg::Config; on_lost::Function = () -> nothing)
|
||||||
|
caps = (classify = cfg.queue_capacity, enrich = cfg.known_queue_capacity,
|
||||||
|
triage = cfg.unknown_queue_capacity, language = cfg.text_queue_capacity)
|
||||||
|
workers = (classify = cfg.worker_count, enrich = cfg.known_worker_count,
|
||||||
|
triage = cfg.unknown_worker_count, language = cfg.text_worker_count)
|
||||||
|
return NamedTuple{STAGE_KEYS}(map(STAGE_KEYS) do key
|
||||||
|
open_queue(backend, cfg, key, caps[key], workers[key]; on_lost)
|
||||||
|
end)
|
||||||
|
end
|
||||||
|
|
||||||
|
"Close the shared connection. Safe to call on a connection that is already gone."
|
||||||
|
function close_backend!(backend::RabbitBackend)
|
||||||
|
try
|
||||||
|
close(backend.conn)
|
||||||
|
catch e
|
||||||
|
@debug "error closing broker connection" exception=e
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
confirm_select!(chan)
|
||||||
|
|
||||||
|
Put `chan` into publisher-confirm mode.
|
||||||
|
|
||||||
|
Not `AMQPClient.confirm_select`, which is broken in AMQPClient 0.5.1: it builds
|
||||||
|
the `Confirm.Select` payload with no fields, while the spec (and the library's
|
||||||
|
own `CLASS_MAP`) declares a `Nowait` bit, so it throws a `BoundsError` before
|
||||||
|
anything reaches the wire. This is the same call with that one field supplied.
|
||||||
|
Drop it for the library's version once the fix is released.
|
||||||
|
"""
|
||||||
|
function confirm_select!(chan::AMQPClient.MessageChannel)
|
||||||
|
AMQPClient._wait_resp(chan, true, false, AMQPClient.on_confirm_select_ok,
|
||||||
|
:Confirm, :SelectOk, false, AMQPClient.DEFAULT_TIMEOUT) do
|
||||||
|
AMQPClient.send(chan, AMQPClient.TAMQPMethodPayload(:Confirm, :Select, (false,)))
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- consume ----------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
Consumer callback: hand a delivery to the bridge channel.
|
||||||
|
|
||||||
|
Blocking here is the mechanism, not a problem. The bridge holds `prefetch` jobs,
|
||||||
|
so when workers fall behind this parks, the broker's prefetch window fills, and
|
||||||
|
it stops sending — backpressure all the way to the queue, with no unbounded
|
||||||
|
local buffer.
|
||||||
|
"""
|
||||||
|
function deliver!(q::RabbitQueue, msg::AMQPClient.Message)
|
||||||
|
try
|
||||||
|
put!(q.buffer, job_from_message(msg))
|
||||||
|
Threads.atomic_sub!(q.depth, 1)
|
||||||
|
catch e
|
||||||
|
# Either we are shutting down (buffer closed after `basic_cancel`, and a
|
||||||
|
# delivery raced it) or the body was unparseable. Neither may kill the
|
||||||
|
# consumer task. The message stays unacked, so the broker redelivers it.
|
||||||
|
q.closed[] || @error "could not accept delivery" queue=q.name exception=(e, catch_backtrace())
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
function dequeue!(q::RabbitQueue)::Union{Job,Nothing}
|
||||||
|
try
|
||||||
|
return take!(q.buffer)
|
||||||
|
catch e
|
||||||
|
# Closed and drained: the shutdown signal `worker_loop` waits for.
|
||||||
|
e isa InvalidStateException && return nothing
|
||||||
|
rethrow()
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Acks and rejects go out on the channel that made the delivery, because a
|
||||||
|
delivery tag is scoped to its channel.
|
||||||
|
|
||||||
|
A failure here is logged, not thrown: an unacked job is redelivered, which is
|
||||||
|
the correct outcome, and killing the worker over it would shrink the pool.
|
||||||
|
"""
|
||||||
|
function ack!(q::RabbitQueue, job::Job)
|
||||||
|
job.delivery_tag == 0 && return nothing
|
||||||
|
try
|
||||||
|
AMQPClient.basic_ack(q.con_chan, job.delivery_tag)
|
||||||
|
catch e
|
||||||
|
@error "ack failed; job will be redelivered" queue=q.name id=job.id exception=e
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
function nack!(q::RabbitQueue, job::Job)
|
||||||
|
job.delivery_tag == 0 && return nothing
|
||||||
|
try
|
||||||
|
AMQPClient.basic_reject(q.con_chan, job.delivery_tag; requeue = false)
|
||||||
|
catch e
|
||||||
|
@error "reject failed; job will be redelivered" queue=q.name id=job.id exception=e
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- publish ----------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
enqueue!(q, job) -> Bool
|
||||||
|
|
||||||
|
Publish `job` as a persistent message. Returns `false` — which intake turns into
|
||||||
|
a `503` and the routing stages turn into park-and-retry — when the queue is
|
||||||
|
closed, when the polled depth says it is at capacity, or when a confirm-mode
|
||||||
|
publish is not confirmed.
|
||||||
|
|
||||||
|
`delivery_mode = PERSISTENT` against a durable queue is what makes a broker
|
||||||
|
restart survivable; without it the queue would come back empty and the whole
|
||||||
|
point of the backend would be gone.
|
||||||
|
"""
|
||||||
|
function enqueue!(q::RabbitQueue, job::Job)::Bool
|
||||||
|
q.closed[] && return false
|
||||||
|
q.depth[] >= q.capacity && return false # advisory: see the `RabbitQueue` docstring
|
||||||
|
|
||||||
|
msg = AMQPClient.Message(Vector{UInt8}(job_json(job));
|
||||||
|
content_type = "application/json",
|
||||||
|
delivery_mode = AMQPClient.PERSISTENT)
|
||||||
|
ok = q.confirms ? publish_confirmed!(q, msg) : publish!(q, msg)
|
||||||
|
ok && Threads.atomic_add!(q.depth, 1)
|
||||||
|
return ok
|
||||||
|
end
|
||||||
|
|
||||||
|
function publish!(q::RabbitQueue, msg::AMQPClient.Message)
|
||||||
|
try
|
||||||
|
AMQPClient.basic_publish(q.pub_chan, msg; exchange = "", routing_key = q.name)
|
||||||
|
return true
|
||||||
|
catch e
|
||||||
|
@error "publish failed" queue=q.name exception=(e, catch_backtrace())
|
||||||
|
return false
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Publish and wait for the broker's confirm.
|
||||||
|
|
||||||
|
Serialized on `pub_lock`, so exactly one publish is outstanding per queue and the
|
||||||
|
next verdict to arrive belongs to it. That costs intake a round trip per file and
|
||||||
|
caps its rate at roughly one file per broker round trip; it is why confirms are
|
||||||
|
scoped to intake and configurable (`FS_AMQP_CONFIRMS`).
|
||||||
|
"""
|
||||||
|
function publish_confirmed!(q::RabbitQueue, msg::AMQPClient.Message)
|
||||||
|
lock(q.pub_lock) do
|
||||||
|
while isready(q.confirm_ch) # discard a verdict left by a timed-out publish
|
||||||
|
take!(q.confirm_ch)
|
||||||
|
end
|
||||||
|
q.pending_since[] = time()
|
||||||
|
try
|
||||||
|
publish!(q, msg) || return false
|
||||||
|
return take!(q.confirm_ch) # poll_loop! guarantees this is answered
|
||||||
|
finally
|
||||||
|
q.pending_since[] = 0.0
|
||||||
|
end
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"Record the broker's verdict for the one outstanding publish."
|
||||||
|
function settle_confirm!(q::RabbitQueue, ok::Bool)
|
||||||
|
q.pending_since[] == 0.0 && return nothing # nothing waiting; a late/duplicate frame
|
||||||
|
isready(q.confirm_ch) || put!(q.confirm_ch, ok)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- poll -------------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
poll_loop!(q)
|
||||||
|
|
||||||
|
One task per queue, doing the two things that must happen on a timer:
|
||||||
|
|
||||||
|
* refresh `depth` from the broker (a passive `queue.declare` returns the ready
|
||||||
|
count), which feeds `/stats` and the advisory capacity check, and
|
||||||
|
* answer a publisher confirm that has gone missing, so a broker that accepted
|
||||||
|
a publish and then went quiet can't park an HTTP request thread forever, and
|
||||||
|
* notice that the connection has died and call `on_lost`, which asks `run` to
|
||||||
|
shut down. This backend does not reconnect (every in-flight delivery tag
|
||||||
|
would be invalid, and the jobs holding them would be processed twice), so a
|
||||||
|
lost connection means the process is finished: it drains, exits, and the
|
||||||
|
restart policy brings it back to redelivered messages and an intact
|
||||||
|
`spool/`. Without this the pipeline would sit there looking healthy with
|
||||||
|
four queues that can never receive anything again.
|
||||||
|
|
||||||
|
Folding the watchdog in here keeps it free: it is a comparison on a task that was
|
||||||
|
going to wake up anyway.
|
||||||
|
"""
|
||||||
|
function poll_loop!(q::RabbitQueue)
|
||||||
|
while !q.closed[]
|
||||||
|
sleep(AMQP_DEPTH_POLL_SECONDS)
|
||||||
|
q.closed[] && break
|
||||||
|
|
||||||
|
pending = q.pending_since[]
|
||||||
|
if pending != 0.0 && time() - pending > AMQP_CONFIRM_TIMEOUT_SECONDS
|
||||||
|
@warn "publisher confirm timed out; reporting the file as not accepted" queue=q.name
|
||||||
|
settle_confirm!(q, false)
|
||||||
|
end
|
||||||
|
|
||||||
|
try
|
||||||
|
lock(q.backend.admin_lock) do
|
||||||
|
ok, _, count, _ = AMQPClient.queue_declare(q.backend.admin_chan, q.name;
|
||||||
|
passive = true)
|
||||||
|
ok && (q.depth[] = Int(count))
|
||||||
|
end
|
||||||
|
catch e
|
||||||
|
q.closed[] && break
|
||||||
|
# A dead connection is not something this backend tries to ride out:
|
||||||
|
# every in-flight delivery tag becomes invalid on reconnect, so the
|
||||||
|
# honest recovery is a restart, where unacked messages redeliver and
|
||||||
|
# `spool/` is intact. Report it and let the supervisor act.
|
||||||
|
if !isopen(q.backend.conn)
|
||||||
|
# Two causes, indistinguishable from here: the broker or the
|
||||||
|
# network went away, or a SIGINT was delivered to one of
|
||||||
|
# AMQPClient's tasks instead of the main loop (see `run`).
|
||||||
|
# The response is the same either way — drain and exit.
|
||||||
|
@error "broker connection closed (lost broker, or an interrupt delivered to a broker task); shutting down (unacked jobs will be redelivered)" queue=q.name
|
||||||
|
q.on_lost()
|
||||||
|
break
|
||||||
|
end
|
||||||
|
@warn "queue depth poll failed" queue=q.name exception=e
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- shutdown ---------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
close!(q)
|
||||||
|
|
||||||
|
Stop consuming, then close the bridge so workers drain what is already in hand
|
||||||
|
and exit — the same drain-then-exit contract `ChannelQueue` offers.
|
||||||
|
|
||||||
|
Jobs the broker has delivered but that we never got to stay unacked and are
|
||||||
|
redelivered on the next start. That is the backend working as intended, and the
|
||||||
|
reason shutdown does not need to be careful here.
|
||||||
|
"""
|
||||||
|
function close!(q::RabbitQueue)
|
||||||
|
q.closed[] || try
|
||||||
|
AMQPClient.basic_cancel(q.con_chan, q.consumer_tag)
|
||||||
|
catch e
|
||||||
|
@debug "error cancelling consumer" queue=q.name exception=e
|
||||||
|
end
|
||||||
|
q.closed[] = true
|
||||||
|
isopen(q.buffer) && close(q.buffer)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
257
src/server.jl
257
src/server.jl
@@ -1,13 +1,21 @@
|
|||||||
# HTTP layer: a single multipart upload endpoint plus a health check.
|
# HTTP layer: a single multipart upload endpoint plus a health check.
|
||||||
#
|
#
|
||||||
# The handler's whole job is to get files onto the queue fast and get out of the
|
# The handler's whole job is to get files onto the queue fast and get out of the
|
||||||
# way: spool each uploaded file to disk, enqueue a reference, respond 202. It
|
# way: stream each uploaded file to disk, enqueue a reference, respond 202. It
|
||||||
# never does real processing — that's the workers' job.
|
# never does real processing; that's the workers' job.
|
||||||
|
#
|
||||||
|
# Intake is *streamed*, not buffered (see src/multipart.jl for the reader). This
|
||||||
|
# is what makes the service's memory story hold end to end: bytes go from the
|
||||||
|
# socket to the spool file a chunk at a time, so a 4 GB upload costs the same
|
||||||
|
# resident memory as a 4 KB one. It is also why /upload is served by its own
|
||||||
|
# stream handler rather than an Oxygen route; see `root_stream_handler`.
|
||||||
#
|
#
|
||||||
# NOTE: routes are registered at runtime via `register_routes()` (called from
|
# NOTE: routes are registered at runtime via `register_routes()` (called from
|
||||||
# `run`), NOT with top-level macros. In a precompiled package, top-level
|
# `run`), NOT with top-level macros. In a precompiled package, top-level
|
||||||
# `@get`/`@post` would execute during precompilation and be lost before serving.
|
# `@get`/`@post` would execute during precompilation and be lost before serving.
|
||||||
|
|
||||||
|
const UPLOAD_PATH = "/upload"
|
||||||
|
|
||||||
jsonresp(status::Int, data) =
|
jsonresp(status::Int, data) =
|
||||||
HTTP.Response(status, ["Content-Type" => "application/json"], JSON3.write(data))
|
HTTP.Response(status, ["Content-Type" => "application/json"], JSON3.write(data))
|
||||||
|
|
||||||
@@ -15,48 +23,215 @@ function health_handler(_::HTTP.Request)
|
|||||||
return jsonresp(200, (; status = "ok"))
|
return jsonresp(200, (; status = "ok"))
|
||||||
end
|
end
|
||||||
|
|
||||||
function upload_handler(req::HTTP.Request)
|
"""
|
||||||
cfg = CONFIG[]
|
`GET /stats`: the pipeline's own counters (src/stats.jl), as JSON.
|
||||||
queue = QUEUE[]
|
|
||||||
|
|
||||||
parts = try
|
Read-only and cheap: a few atomic loads and one `length` per queue, no pipeline
|
||||||
HTTP.parse_multipart_form(req)
|
state touched. Two scrapes Δt apart give per-stage throughput and utilization;
|
||||||
catch
|
see bin/bench.jl, which is the intended consumer.
|
||||||
nothing
|
|
||||||
end
|
|
||||||
parts === nothing &&
|
|
||||||
return jsonresp(400, (; error = "expected multipart/form-data"))
|
|
||||||
|
|
||||||
files = filter(p -> p.filename !== nothing && !isempty(p.filename), parts)
|
Unlike `/upload` this is an ordinary Oxygen route: it has no body to stream, and
|
||||||
isempty(files) &&
|
being in Oxygen's middleware chain is a feature here.
|
||||||
return jsonresp(400, (; error = "no files found in request"))
|
"""
|
||||||
|
function stats_handler(_::HTTP.Request)
|
||||||
accepted = NamedTuple{(:id, :name),Tuple{String,String}}[]
|
queues = (classify = QUEUE[], enrich = KNOWN_QUEUE[],
|
||||||
for p in files
|
triage = UNKNOWN_QUEUE[], language = TEXT_QUEUE[])
|
||||||
bytes = read(p.data)
|
return jsonresp(200, stats_snapshot(CONFIG[], queues))
|
||||||
|
|
||||||
job = try
|
|
||||||
spool_file(cfg, p.filename, bytes)
|
|
||||||
catch e
|
|
||||||
@error "spool failed" name=p.filename exception=(e, catch_backtrace())
|
|
||||||
return jsonresp(500, (; error = "failed to store file", accepted))
|
|
||||||
end
|
|
||||||
|
|
||||||
if !enqueue!(queue, job)
|
|
||||||
rm(job.path; force = true) # never queued → don't leave it in spool
|
|
||||||
return jsonresp(503, (; error = "queue full, retry later", accepted))
|
|
||||||
end
|
|
||||||
|
|
||||||
@info "accepted" id=job.id name=job.original_name size=job.size
|
|
||||||
push!(accepted, (; id = job.id, name = job.original_name))
|
|
||||||
end
|
|
||||||
|
|
||||||
return jsonresp(202, (; accepted))
|
|
||||||
end
|
end
|
||||||
|
|
||||||
"Register HTTP routes on the Oxygen instance. Must run at runtime, before serve."
|
"Write a JSON response onto a raw stream (the streaming handler's `jsonresp`)."
|
||||||
function register_routes()
|
function stream_jsonresp(stream::HTTP.Stream, status::Int, data)
|
||||||
@get("/health", health_handler)
|
body = JSON3.write(data)
|
||||||
@post("/upload", upload_handler)
|
HTTP.setstatus(stream, status)
|
||||||
|
HTTP.setheader(stream, "Content-Type" => "application/json")
|
||||||
|
HTTP.setheader(stream, "Content-Length" => string(sizeof(body)))
|
||||||
|
HTTP.startwrite(stream)
|
||||||
|
write(stream, body)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Read and discard whatever is left of the request body.
|
||||||
|
|
||||||
|
HTTP.jl's server calls `closeread` after the handler and *errors* if the body was
|
||||||
|
only partly consumed, because a half-read body can't be followed by another request on a
|
||||||
|
keep-alive connection. So every exit path drains first. Discarding is bounded in
|
||||||
|
memory (one chunk) and costs nothing in the normal case, where the body is
|
||||||
|
already fully consumed and this returns immediately.
|
||||||
|
"""
|
||||||
|
function discard_body!(stream::HTTP.Stream, chunk_bytes::Int)
|
||||||
|
scratch = Vector{UInt8}(undef, chunk_bytes)
|
||||||
|
while !eof(stream)
|
||||||
|
readbytes!(stream, scratch, chunk_bytes) == 0 && break
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Did this exception mean the client went away, rather than something being wrong
|
||||||
|
on our side?
|
||||||
|
|
||||||
|
A client that hangs up mid-upload (user cancels, network drops) is routine, and
|
||||||
|
must not be logged as a server error or reported as a failed write. But it looks
|
||||||
|
like an I/O failure from inside the parser, so the distinction has to be made
|
||||||
|
explicitly. `EOFError` is what `HTTP.Stream` raises when a connection dies with
|
||||||
|
bytes still promised by `Content-Length`.
|
||||||
|
"""
|
||||||
|
is_client_gone(e) =
|
||||||
|
e isa EOFError ||
|
||||||
|
(e isa Base.IOError && e.code in (Base.UV_EPIPE, Base.UV_ECONNRESET, Base.UV_ECONNABORTED))
|
||||||
|
|
||||||
|
"""
|
||||||
|
Stream a multipart upload to disk, one part at a time.
|
||||||
|
|
||||||
|
Each file part is spooled straight from the socket and its reference enqueued.
|
||||||
|
The response reports what was actually queued:
|
||||||
|
|
||||||
|
* `202`: every file in the request was spooled and queued
|
||||||
|
* `400`: not multipart/form-data, no files present, or a malformed body
|
||||||
|
* `503`: the intake queue filled up; `accepted` lists what got in first
|
||||||
|
* `500`: a file could not be written to disk
|
||||||
|
|
||||||
|
Unlike a buffered handler, this one cannot know up front how many files a request
|
||||||
|
holds or whether they will fit. So when the queue fills mid-request it does not
|
||||||
|
abandon the connection: it stops spooling (discarding the remaining parts rather
|
||||||
|
than writing files it can't queue), drains the body, and answers `503` with the
|
||||||
|
`accepted` list. Files already queued stay queued, so a client can retry the rest.
|
||||||
|
"""
|
||||||
|
function upload_stream_handler(stream::HTTP.Stream)
|
||||||
|
try
|
||||||
|
return serve_upload!(stream)
|
||||||
|
catch e
|
||||||
|
# The client vanished while we were reading its body, draining it, or
|
||||||
|
# answering. Nothing is wrong on our side, so log it as the routine event
|
||||||
|
# it is instead of a server error.
|
||||||
|
is_client_gone(e) || rethrow()
|
||||||
|
@info "upload aborted by client"
|
||||||
|
# HTTP.jl insists a handler write *some* response before returning. The
|
||||||
|
# client is probably already gone, so this write is best-effort: attempt
|
||||||
|
# it only if we haven't started a response, and let it fail silently.
|
||||||
|
if isopen(stream) && !iswritable(stream)
|
||||||
|
try
|
||||||
|
stream_jsonresp(stream, 400, (; error = "upload truncated"))
|
||||||
|
catch e2
|
||||||
|
is_client_gone(e2) || rethrow()
|
||||||
|
end
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
function serve_upload!(stream::HTTP.Stream)
|
||||||
|
Threads.atomic_add!(METRICS.intake.requests, 1)
|
||||||
|
cfg = CONFIG[]
|
||||||
|
queue = QUEUE[]
|
||||||
|
chunk = cfg.upload_chunk_bytes
|
||||||
|
|
||||||
|
boundary = multipart_boundary(HTTP.header(stream.message, "Content-Type", nothing))
|
||||||
|
if boundary === nothing
|
||||||
|
discard_body!(stream, chunk)
|
||||||
|
return stream_jsonresp(stream, 400, (; error = "expected multipart/form-data"))
|
||||||
|
end
|
||||||
|
|
||||||
|
reader = MultipartReader(stream, boundary; chunk_bytes = chunk)
|
||||||
|
accepted = NamedTuple{(:id, :name),Tuple{String,String}}[]
|
||||||
|
n_files = 0
|
||||||
|
queue_full = false
|
||||||
|
failure = nothing # (status, message) from a fatal error mid-body
|
||||||
|
|
||||||
|
try
|
||||||
|
while (part = next_part!(reader)) !== nothing
|
||||||
|
# A part with no filename is an ordinary form field, not a file.
|
||||||
|
if part.filename === nothing || isempty(part.filename)
|
||||||
|
skip_part_body!(reader)
|
||||||
|
continue
|
||||||
|
end
|
||||||
|
n_files += 1
|
||||||
|
|
||||||
|
# Already backpressured: consume the part, but don't write a file we
|
||||||
|
# know we cannot enqueue.
|
||||||
|
if queue_full
|
||||||
|
Threads.atomic_add!(METRICS.intake.rejected, 1)
|
||||||
|
skip_part_body!(reader)
|
||||||
|
continue
|
||||||
|
end
|
||||||
|
|
||||||
|
job = try
|
||||||
|
spool_stream(cfg, part.filename) do io
|
||||||
|
write_part_body!(io, reader)
|
||||||
|
end
|
||||||
|
catch e
|
||||||
|
(e isa MultipartError || is_client_gone(e)) && rethrow()
|
||||||
|
# A disk error leaves the reader mid-part, so the rest of the body
|
||||||
|
# can no longer be parsed, so stop and report.
|
||||||
|
@error "spool failed" name=part.filename exception=(e, catch_backtrace())
|
||||||
|
failure = (500, "failed to store file")
|
||||||
|
break
|
||||||
|
end
|
||||||
|
|
||||||
|
if enqueue!(queue, job)
|
||||||
|
@info "accepted" id=job.id name=job.original_name size=job.size
|
||||||
|
push!(accepted, (; id = job.id, name = job.original_name))
|
||||||
|
# Counted at the point the job becomes stage 1's problem, so
|
||||||
|
# intake totals and stage-1 arrivals refer to the same files.
|
||||||
|
Threads.atomic_add!(METRICS.intake.files, 1)
|
||||||
|
Threads.atomic_add!(METRICS.intake.bytes, job.size)
|
||||||
|
else
|
||||||
|
rm(job.path; force = true) # never queued → don't leave it in spool
|
||||||
|
Threads.atomic_add!(METRICS.intake.rejected, 1)
|
||||||
|
queue_full = true
|
||||||
|
end
|
||||||
|
end
|
||||||
|
catch e
|
||||||
|
e isa MultipartError || rethrow()
|
||||||
|
@warn "malformed multipart upload" reason=e.msg accepted=length(accepted)
|
||||||
|
failure = (400, "malformed multipart body")
|
||||||
|
end
|
||||||
|
|
||||||
|
discard_body!(stream, chunk)
|
||||||
|
|
||||||
|
failure !== nothing &&
|
||||||
|
return stream_jsonresp(stream, failure[1], (; error = failure[2], accepted))
|
||||||
|
n_files == 0 &&
|
||||||
|
return stream_jsonresp(stream, 400, (; error = "no files found in request"))
|
||||||
|
queue_full &&
|
||||||
|
return stream_jsonresp(stream, 503, (; error = "queue full, retry later", accepted))
|
||||||
|
return stream_jsonresp(stream, 202, (; accepted))
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
root_stream_handler(middleware) -> (stream -> nothing)
|
||||||
|
|
||||||
|
Oxygen's root handler wraps `HTTP.streamhandler`, which does
|
||||||
|
`request.body = read(stream)`, the entire body into memory, *before* any route
|
||||||
|
is dispatched. That happens even for an Oxygen `@stream` route, so no route can
|
||||||
|
stream an upload. We therefore intercept `POST /upload` at the stream level and
|
||||||
|
hand everything else to Oxygen unchanged.
|
||||||
|
|
||||||
|
Trade-off: `/upload` bypasses Oxygen's middleware chain, so it is absent from
|
||||||
|
Oxygen's built-in metrics and docs. Deliberate: flat intake memory is the point
|
||||||
|
of this service, and intake is covered by our own counters (`/stats`) and `@info`
|
||||||
|
records anyway.
|
||||||
|
"""
|
||||||
|
function root_stream_handler(middleware::Function)
|
||||||
|
oxygen_handler = Oxygen.Core.stream_handler(middleware)
|
||||||
|
return function (stream::HTTP.Stream)
|
||||||
|
req = stream.message
|
||||||
|
if req.method == "POST" && HTTP.URI(req.target).path == UPLOAD_PATH
|
||||||
|
return upload_stream_handler(stream)
|
||||||
|
end
|
||||||
|
return oxygen_handler(stream)
|
||||||
|
end
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
Register HTTP routes on the Oxygen instance. Must run at runtime, before serve.
|
||||||
|
|
||||||
|
`POST /upload` is deliberately absent: it is served by `upload_stream_handler`
|
||||||
|
via `root_stream_handler`, ahead of Oxygen's router.
|
||||||
|
"""
|
||||||
|
function register_routes()
|
||||||
|
@get("/health", health_handler)
|
||||||
|
@get("/stats", stats_handler)
|
||||||
return nothing
|
return nothing
|
||||||
end
|
end
|
||||||
|
|||||||
80
src/spool.jl
80
src/spool.jl
@@ -19,18 +19,44 @@ function sanitize_filename(name::AbstractString)::String
|
|||||||
return first(base, MAX_NAME_LEN)
|
return first(base, MAX_NAME_LEN)
|
||||||
end
|
end
|
||||||
|
|
||||||
"Write `bytes` to the spool dir under `<uuid>-<sanitized>` and return the Job."
|
"Build the spool path for a client-supplied name: `<uuid>-<sanitized>`."
|
||||||
function spool_file(cfg::Config, original_name::AbstractString, bytes::Vector{UInt8})::Job
|
function spool_path(cfg::Config, original_name::AbstractString)
|
||||||
id = string(uuid4())
|
id = string(uuid4())
|
||||||
safe = sanitize_filename(original_name)
|
return id, joinpath(cfg.spool_dir, string(id, "-", sanitize_filename(original_name)))
|
||||||
path = joinpath(cfg.spool_dir, string(id, "-", safe))
|
|
||||||
open(path, "w") do io
|
|
||||||
write(io, bytes)
|
|
||||||
end
|
|
||||||
return Job(id, String(original_name), path, length(bytes), time())
|
|
||||||
end
|
end
|
||||||
|
|
||||||
"Move a spooled file into `dir` (done/ or failed/), returning the destination."
|
"""
|
||||||
|
spool_stream(write_body!, cfg, original_name) -> Job
|
||||||
|
|
||||||
|
Create the spool file for `original_name`, hand the open `IO` to `write_body!`,
|
||||||
|
and build the `Job` from however many bytes it reports writing.
|
||||||
|
|
||||||
|
This is the streaming counterpart to `spool_file`: the caller pumps bytes in from
|
||||||
|
the network as they arrive, so a file never exists in memory in one piece. A
|
||||||
|
partial file left by a failed or abandoned write is removed: intake either
|
||||||
|
produces a complete spooled file or nothing at all, so recovery on restart never
|
||||||
|
picks up a truncated upload.
|
||||||
|
"""
|
||||||
|
function spool_stream(write_body!, cfg::Config, original_name::AbstractString)::Job
|
||||||
|
id, path = spool_path(cfg, original_name)
|
||||||
|
nbytes = try
|
||||||
|
open(write_body!, path, "w")
|
||||||
|
catch
|
||||||
|
rm(path; force = true)
|
||||||
|
rethrow()
|
||||||
|
end
|
||||||
|
return Job(id, String(original_name), path, nbytes, time())
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
move_to(dir, job) -> String
|
||||||
|
|
||||||
|
Move a spooled file into `dir`, returning the destination.
|
||||||
|
|
||||||
|
This runs once per file, at the end: into a terminal sink (`done/`, `text_done/`,
|
||||||
|
`binary/`) or into `failed/` when a handler throws. Files are *not* moved between
|
||||||
|
stages; see the header of src/worker.jl.
|
||||||
|
"""
|
||||||
function move_to(dir::AbstractString, job::Job)::String
|
function move_to(dir::AbstractString, job::Job)::String
|
||||||
dest = joinpath(dir, basename(job.path))
|
dest = joinpath(dir, basename(job.path))
|
||||||
mv(job.path, dest; force=true)
|
mv(job.path, dest; force=true)
|
||||||
@@ -42,17 +68,34 @@ end
|
|||||||
const UUID_LEN = 36
|
const UUID_LEN = 36
|
||||||
|
|
||||||
"""
|
"""
|
||||||
recover_spool!(cfg, queue) -> Int
|
recover_dir!(dir, queue) -> Int
|
||||||
|
|
||||||
Re-enqueue any files already sitting in the spool directory (left by a crash,
|
Re-enqueue any files sitting in `dir` (left by a crash, hard shutdown, or an
|
||||||
a hard shutdown, or an intake that never got processed). This is the payoff of
|
intake that never finished) onto `queue`. This is the payoff of spooling to
|
||||||
spooling to disk: a restart resumes work instead of stranding it. Returns the
|
disk: a restart resumes work instead of stranding it. Returns the number of
|
||||||
number of files recovered.
|
files recovered.
|
||||||
|
|
||||||
|
A file's stage is not recorded on disk. It lives in whichever in-process queue
|
||||||
|
holds its reference, and a crash loses that (src/worker.jl header). So everything
|
||||||
|
still in `spool/` re-enters at stage 1 and replays. Replay is safe rather than
|
||||||
|
merely tolerable: classification and the binary/text sniff are pure functions of
|
||||||
|
the file's bytes, and the terminal commits rename with `force=true`, so a
|
||||||
|
replayed file lands where it would have landed and overwrites its own sidecar.
|
||||||
|
|
||||||
|
**This blocks on a full queue**, and that is the whole point: `spool/` now holds
|
||||||
|
every in-flight file, so a crash under load can easily leave more leftovers than
|
||||||
|
one queue's capacity. Dropping the excess would mean recovery quietly losing the
|
||||||
|
work it exists to save. Because the caller spawns the worker pools *before*
|
||||||
|
recovering, the consumers are live and the queue drains as we fill it, so parking
|
||||||
|
here costs latency, never progress.
|
||||||
|
|
||||||
|
Skips `.meta.json` sidecars: those are enrichment output, not work to redo.
|
||||||
"""
|
"""
|
||||||
function recover_spool!(cfg::Config, queue::JobQueue)::Int
|
function recover_dir!(dir::AbstractString, queue::JobQueue)::Int
|
||||||
n = 0
|
n = 0
|
||||||
for path in sort(readdir(cfg.spool_dir; join=true))
|
for path in sort(readdir(dir; join=true))
|
||||||
isfile(path) || continue
|
isfile(path) || continue
|
||||||
|
endswith(path, ".meta.json") && continue # sidecar, not a work item
|
||||||
fname = basename(path)
|
fname = basename(path)
|
||||||
if length(fname) > UUID_LEN + 1
|
if length(fname) > UUID_LEN + 1
|
||||||
id = fname[1:UUID_LEN]
|
id = fname[1:UUID_LEN]
|
||||||
@@ -62,7 +105,10 @@ function recover_spool!(cfg::Config, queue::JobQueue)::Int
|
|||||||
name = fname
|
name = fname
|
||||||
end
|
end
|
||||||
job = Job(id, name, path, filesize(path), time())
|
job = Job(id, name, path, filesize(path), time())
|
||||||
enqueue!(queue, job) || @warn "queue full during recovery; leaving file for next start" path
|
# `nothing` stats: recovery is not a stage, so its wait must not be
|
||||||
|
# charged to any stage's blocked_ns (see `enqueue_blocking!`).
|
||||||
|
enqueue_blocking!(queue, job, nothing;
|
||||||
|
retry_seconds = ROUTE_ENQUEUE_RETRY_SECONDS)
|
||||||
n += 1
|
n += 1
|
||||||
end
|
end
|
||||||
return n
|
return n
|
||||||
|
|||||||
220
src/stats.jl
Normal file
220
src/stats.jl
Normal file
@@ -0,0 +1,220 @@
|
|||||||
|
# Per-stage pipeline metrics.
|
||||||
|
#
|
||||||
|
# End-to-end throughput says how fast the pipeline is; it does not say which
|
||||||
|
# stage is the reason. The external harness can't answer that either, and now
|
||||||
|
# cannot even approximate it: every in-flight file sits in `spool/` regardless of
|
||||||
|
# how far it has got, because stages route by enqueueing a reference rather than
|
||||||
|
# by moving bytes (src/worker.jl header). A file's stage is a property of the
|
||||||
|
# queue holding it, so only the pipeline itself can see it. There is no
|
||||||
|
# directory to poll.
|
||||||
|
#
|
||||||
|
# So each stage keeps four counters, all incremented in `worker_loop` (the one
|
||||||
|
# place every stage's work passes through) and read by `GET /stats`:
|
||||||
|
#
|
||||||
|
# completed / failed jobs finished, jobs quarantined → throughput
|
||||||
|
# bytes job bytes processed → per-stage MiB/s
|
||||||
|
# busy_ns summed handler wall time → service time
|
||||||
|
# blocked_ns of that, time parked on a full downstream queue
|
||||||
|
# in_flight handlers running right now → saturation
|
||||||
|
#
|
||||||
|
# Rates fall out of a pair of scrapes taken Δt apart:
|
||||||
|
#
|
||||||
|
# throughput = Δcompleted / Δt
|
||||||
|
# utilization = (Δbusy_ns - Δblocked_ns) / (Δt * workers) ~1.0 ⇒ bottleneck
|
||||||
|
#
|
||||||
|
# Utilization is the useful one. Throughput alone can't distinguish a stage that
|
||||||
|
# is saturated from one that is merely starved by the stage ahead of it: both
|
||||||
|
# show the same files/s. Utilization separates them: the saturated stage
|
||||||
|
# sits near 1.0 with its queue backing up, the starved stage sits near 0.
|
||||||
|
#
|
||||||
|
# `blocked_ns` is what keeps that true downstream. Stages 1 and 3 apply blocking
|
||||||
|
# backpressure: when the next queue is full the handler parks and retries rather
|
||||||
|
# than dropping the file (see ROUTE_ENQUEUE_RETRY_SECONDS). That parked time is
|
||||||
|
# inside the handler, so counting it as busy would show stage 1 pinned at 1.0
|
||||||
|
# whenever stage 2 is the real bottleneck, and every stage upstream of the jam would
|
||||||
|
# look like the jam. Subtracting it leaves utilization meaning "doing its own
|
||||||
|
# work", and the blocked share becomes its own signal: a stage blocked 90% of the
|
||||||
|
# time is naming its successor as the bottleneck.
|
||||||
|
#
|
||||||
|
# Everything here is monotonic since process start (or since `reset_metrics!`),
|
||||||
|
# in the Prometheus style: counters, never rates. Rates are the reader's job, so
|
||||||
|
# a scrape is stateless and two readers can't disturb each other.
|
||||||
|
#
|
||||||
|
# Cost is a handful of atomic adds per file, against handlers that spawn
|
||||||
|
# exiftool: unmeasurable in practice, and the counters are never read on the
|
||||||
|
# hot path.
|
||||||
|
|
||||||
|
# Stage identity lives here, in declaration order, so the report, the JSON and
|
||||||
|
# the worker wiring can't drift apart: adding stage 5 to the live path means
|
||||||
|
# adding it here and passing the new `StageStats` to its `worker_loop`.
|
||||||
|
const STAGE_KEYS = (:classify, :enrich, :triage, :language)
|
||||||
|
const STAGE_TITLES = (classify = "classify", enrich = "enrich",
|
||||||
|
triage = "triage", language = "language")
|
||||||
|
|
||||||
|
"""
|
||||||
|
Counters for one pipeline stage. All fields are atomic and monotonic: workers
|
||||||
|
only ever add, readers only ever read, so no lock is needed between them.
|
||||||
|
|
||||||
|
`in_flight` is the exception to monotonic. It goes up and down, and is the one
|
||||||
|
counter that is a *level* rather than a total.
|
||||||
|
"""
|
||||||
|
struct StageStats
|
||||||
|
completed::Threads.Atomic{Int}
|
||||||
|
failed::Threads.Atomic{Int}
|
||||||
|
bytes::Threads.Atomic{Int}
|
||||||
|
busy_ns::Threads.Atomic{Int}
|
||||||
|
blocked_ns::Threads.Atomic{Int}
|
||||||
|
in_flight::Threads.Atomic{Int}
|
||||||
|
end
|
||||||
|
|
||||||
|
StageStats() = StageStats(Threads.Atomic{Int}(0), Threads.Atomic{Int}(0),
|
||||||
|
Threads.Atomic{Int}(0), Threads.Atomic{Int}(0),
|
||||||
|
Threads.Atomic{Int}(0), Threads.Atomic{Int}(0))
|
||||||
|
|
||||||
|
"""
|
||||||
|
Counters for HTTP intake, which has no worker loop to hang them off.
|
||||||
|
|
||||||
|
`files` counts what was actually spooled *and* queued, the same thing the 202
|
||||||
|
body reports as `accepted`, so it lines up with stage 1's arrivals. Files
|
||||||
|
dropped because the queue was full are counted separately in `rejected`, since a
|
||||||
|
run where intake outruns the pipeline should be visible as backpressure rather
|
||||||
|
than as slow intake.
|
||||||
|
"""
|
||||||
|
struct IntakeStats
|
||||||
|
requests::Threads.Atomic{Int}
|
||||||
|
files::Threads.Atomic{Int}
|
||||||
|
bytes::Threads.Atomic{Int}
|
||||||
|
rejected::Threads.Atomic{Int}
|
||||||
|
end
|
||||||
|
|
||||||
|
IntakeStats() = IntakeStats(Threads.Atomic{Int}(0), Threads.Atomic{Int}(0),
|
||||||
|
Threads.Atomic{Int}(0), Threads.Atomic{Int}(0))
|
||||||
|
|
||||||
|
"""
|
||||||
|
Every counter in the process, plus the wall-clock origin the totals are measured
|
||||||
|
from.
|
||||||
|
|
||||||
|
`since` matters to a reader computing rates from a single scrape: without it,
|
||||||
|
"1000 files completed" has no denominator. Readers that scrape twice should use
|
||||||
|
their own Δt instead, which excludes the time before they started watching.
|
||||||
|
"""
|
||||||
|
struct Metrics
|
||||||
|
intake::IntakeStats
|
||||||
|
stages::NamedTuple{STAGE_KEYS,NTuple{4,StageStats}}
|
||||||
|
since::Base.RefValue{Float64}
|
||||||
|
end
|
||||||
|
|
||||||
|
Metrics() = Metrics(IntakeStats(),
|
||||||
|
NamedTuple{STAGE_KEYS}(ntuple(_ -> StageStats(), 4)),
|
||||||
|
Ref(time()))
|
||||||
|
|
||||||
|
# Process-global, like the logger: workers on every thread add to it and the
|
||||||
|
# HTTP layer reads it, with no way to thread a handle through both paths that
|
||||||
|
# wouldn't just be this by another name.
|
||||||
|
const METRICS = Metrics()
|
||||||
|
|
||||||
|
"""
|
||||||
|
reset_metrics!()
|
||||||
|
|
||||||
|
Zero every counter and restart the measurement window. For tests, and for a
|
||||||
|
benchmark that wants totals covering only its own run, though a harness that
|
||||||
|
scrapes before and after can subtract instead, which is safer against a server
|
||||||
|
that is also serving someone else.
|
||||||
|
"""
|
||||||
|
function reset_metrics!(m::Metrics = METRICS)
|
||||||
|
for a in (m.intake.requests, m.intake.files, m.intake.bytes, m.intake.rejected)
|
||||||
|
a[] = 0
|
||||||
|
end
|
||||||
|
for s in m.stages
|
||||||
|
s.completed[] = 0
|
||||||
|
s.failed[] = 0
|
||||||
|
s.bytes[] = 0
|
||||||
|
s.busy_ns[] = 0
|
||||||
|
s.blocked_ns[] = 0
|
||||||
|
s.in_flight[] = 0
|
||||||
|
end
|
||||||
|
m.since[] = time()
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"Record one finished job: `ok` distinguishes a completion from a quarantine."
|
||||||
|
function record_job!(s::StageStats, ok::Bool, bytes::Int, elapsed_ns::Int)
|
||||||
|
Threads.atomic_add!(ok ? s.completed : s.failed, 1)
|
||||||
|
Threads.atomic_add!(s.bytes, bytes)
|
||||||
|
Threads.atomic_add!(s.busy_ns, elapsed_ns)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
enqueue_blocking!(queue, job, stats; retry_seconds) -> nothing
|
||||||
|
|
||||||
|
Hand `job` to a downstream queue, parking and retrying until it fits, and charge
|
||||||
|
the parked time to `stats.blocked_ns`.
|
||||||
|
|
||||||
|
This is the routing half of stages 1 and 3: a classified file is never dropped,
|
||||||
|
so a full downstream queue means waiting, not failing. Wrapping the retry loop
|
||||||
|
here, rather than repeating `while !enqueue! sleep end` at each call site, is
|
||||||
|
what makes that wait measurable at all, and keeps the three routing paths from
|
||||||
|
drifting into three different backoff behaviours.
|
||||||
|
|
||||||
|
`stats` is `nothing` for callers that are not a pipeline stage. Startup
|
||||||
|
recovery (`recover_dir!`) uses the same never-drop retry, but it runs outside any
|
||||||
|
handler, so it has no `busy_ns` to charge against. Billing its wait to a stage's
|
||||||
|
`blocked_ns` would make `/stats` utilization, `(Δbusy_ns - Δblocked_ns) / …`,
|
||||||
|
go *negative* after a long recovery. Skipping the charge keeps `blocked_ns`
|
||||||
|
meaning exactly what its docs say: time a handler spent parked on its successor.
|
||||||
|
"""
|
||||||
|
function enqueue_blocking!(queue::JobQueue, job::Job, stats::Union{StageStats,Nothing};
|
||||||
|
retry_seconds::Real)
|
||||||
|
enqueue!(queue, job) && return nothing
|
||||||
|
t0 = time_ns()
|
||||||
|
while !enqueue!(queue, job)
|
||||||
|
sleep(retry_seconds)
|
||||||
|
end
|
||||||
|
stats === nothing || Threads.atomic_add!(stats.blocked_ns, Int(time_ns() - t0))
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
# --- snapshot ---------------------------------------------------------------
|
||||||
|
|
||||||
|
"""
|
||||||
|
stats_snapshot(cfg, queues, m = METRICS) -> NamedTuple
|
||||||
|
|
||||||
|
Read every counter into a plain value tree for `GET /stats`.
|
||||||
|
|
||||||
|
`queues` is a NamedTuple of `JobQueue`s keyed by `STAGE_KEYS`, supplying each
|
||||||
|
stage's live depth and capacity. The counters above are about work *done*, and
|
||||||
|
a stage's depth is what says whether work is piling up in front of it.
|
||||||
|
|
||||||
|
The read is not atomic across stages: counters keep moving while we walk them,
|
||||||
|
so a snapshot can show a job counted as complete by stage 1 but not yet arrived
|
||||||
|
at stage 2. Over a benchmark window that skew is a job or two and does not move
|
||||||
|
a rate; a consistent snapshot would mean stopping the pipeline to read it.
|
||||||
|
"""
|
||||||
|
function stats_snapshot(cfg::Config, queues, m::Metrics = METRICS)
|
||||||
|
workers = (classify = cfg.worker_count, enrich = cfg.known_worker_count,
|
||||||
|
triage = cfg.unknown_worker_count, language = cfg.text_worker_count)
|
||||||
|
now = time()
|
||||||
|
stages = map(STAGE_KEYS, ntuple(identity, 4)) do key, i
|
||||||
|
s, q = m.stages[key], queues[key]
|
||||||
|
(; stage = i,
|
||||||
|
name = STAGE_TITLES[key],
|
||||||
|
workers = workers[key],
|
||||||
|
queue_depth = length(q),
|
||||||
|
queue_capacity = capacity(q),
|
||||||
|
completed = s.completed[],
|
||||||
|
failed = s.failed[],
|
||||||
|
bytes = s.bytes[],
|
||||||
|
busy_seconds = s.busy_ns[] / 1e9,
|
||||||
|
blocked_seconds = s.blocked_ns[] / 1e9,
|
||||||
|
in_flight = s.in_flight[])
|
||||||
|
end
|
||||||
|
return (; now,
|
||||||
|
since = m.since[],
|
||||||
|
uptime_seconds = now - m.since[],
|
||||||
|
intake = (; requests = m.intake.requests[],
|
||||||
|
files = m.intake.files[],
|
||||||
|
bytes = m.intake.bytes[],
|
||||||
|
rejected = m.intake.rejected[]),
|
||||||
|
stages = collect(stages))
|
||||||
|
end
|
||||||
213
src/worker.jl
213
src/worker.jl
@@ -1,47 +1,222 @@
|
|||||||
# Worker task: pull jobs off the queue and process them. One of these runs per
|
# Worker tasks: pull jobs off a queue and process them. The loop scaffolding
|
||||||
# configured worker, each as its own `Threads.@spawn`'d task.
|
# (dequeue-until-drained, try/catch, quarantine-on-throw) is identical for every
|
||||||
|
# stage, so `worker_loop` is parametrized with a `handler` and reused. Today
|
||||||
|
# there are four stages:
|
||||||
|
#
|
||||||
|
# stage 1 handle_classify_job classify → known queue | unknown queue
|
||||||
|
# stage 2 handle_known_job exiftool enrich → done/ (+ .meta.json)
|
||||||
|
# stage 3 handle_unknown_job binary-vs-text sniff → binary/ | text queue
|
||||||
|
# stage 4 handle_text_job language enrich → text_done/ (+ .meta.json)
|
||||||
|
#
|
||||||
|
# Adding a stage later is just another queue + pool + handler; the loop below
|
||||||
|
# doesn't change.
|
||||||
|
#
|
||||||
|
# A file does not move between stages. It is written once into `spool/` at
|
||||||
|
# intake and stays there for its whole in-flight life; only the small `Job`
|
||||||
|
# reference travels, and `job.path` is therefore constant from intake until the
|
||||||
|
# file is committed. The stage a file has reached lives in the queue holding its
|
||||||
|
# reference, not in which directory the bytes sit, so the intermediate hops
|
||||||
|
# (spool→known, spool→unknown, unknown→text) are three renames per file that buy
|
||||||
|
# nothing on the live path. What remains is one move at the end: into a terminal
|
||||||
|
# sink (`done/`, `text_done/`, `binary/`) or into `failed/` on a throw.
|
||||||
|
#
|
||||||
|
# The cost is on restart, and which cost depends on the backend. Under the
|
||||||
|
# in-process `ChannelQueue` a crash loses the queues, and recovery can only
|
||||||
|
# re-drive everything in `spool/` from stage 1. That is safe, since
|
||||||
|
# classification and the UTF-8 sniff are pure functions of the file's bytes and
|
||||||
|
# every commit is idempotent, but it redoes work the old directory-per-stage
|
||||||
|
# layout could skip. Under `RabbitQueue` (src/rabbit.jl) it doesn't arise: a job
|
||||||
|
# stays unacked until its handler commits, so a restart resumes each file at the
|
||||||
|
# stage it had reached. The fix was the queue seam, not the directories.
|
||||||
|
#
|
||||||
|
# (`ROUTE_ENQUEUE_RETRY_SECONDS`, the backoff the routing handoffs below use,
|
||||||
|
# now lives in src/queue.jl, since startup recovery shares it.)
|
||||||
|
|
||||||
|
# Per-file logging in stage 1 is `@debug`, not `@info`, because it is the
|
||||||
|
# stage's dominant cost. Measured by bin/bench_stage1.jl (2000 x 64 KiB files,
|
||||||
|
# min of 5 trials): a formatted pair of log lines costs 71.2 µs per file, against
|
||||||
|
# the 11.7 µs the whole handler takes with them switched off: six times the rest
|
||||||
|
# of the stage put together. Nearly all of it is `ConsoleLogger` formatting
|
||||||
|
# (64.2 µs); the FlushLogger's per-message flush is only ~7 µs on top. Switching
|
||||||
|
# them back on with `JULIA_DEBUG=DarkStruct` takes the handler to 106.3 µs, i.e.
|
||||||
|
# from 85.4k files/s down to 9.4k on one worker.
|
||||||
|
#
|
||||||
|
# With logging off the stage is the classifier and nothing else: classify 10.30 µs
|
||||||
|
# (of which read_features is 7.82 µs), the queue handoff 0.12 µs, the disabled
|
||||||
|
# `@debug` lines 0.29 µs, for 10.7 µs of the 11.7 µs total. Removing the inter-stage
|
||||||
|
# rename is what left it that way: that rename was 11.6 µs per file, so it was
|
||||||
|
# the equal of the classifier, and dropping it took one worker from ~35k to
|
||||||
|
# ~85k files/s (16 workers: ~333k files/s).
|
||||||
|
#
|
||||||
|
# `@debug` is compiled to a min-level check that doesn't evaluate its arguments,
|
||||||
|
# so a disabled line costs ~0.15 µs rather than ~36 µs. The messages are still
|
||||||
|
# there when wanted: run with `JULIA_DEBUG=DarkStruct` to get them back. Errors,
|
||||||
|
# quarantines and lifecycle events stay at `@error`/`@info`, since they are rare and
|
||||||
|
# their cost doesn't scale with throughput. `GET /stats` (src/stats.jl) is the
|
||||||
|
# per-file observability that survives, and it is counted, not formatted.
|
||||||
|
|
||||||
"""
|
"""
|
||||||
handle_job(job, cfg, worker_id)
|
handle_classify_job(job, cfg, worker_id, known_queue, unknown_queue)
|
||||||
|
|
||||||
Do the work for a single job, then move the file to `done/`.
|
Stage 1. Classify the spooled file and route it to the next stage's queue,
|
||||||
|
retrying on a full queue rather than dropping the file:
|
||||||
|
* `:known` → enqueue onto the known queue for stage 2.
|
||||||
|
* `:unknown` → enqueue onto the unknown queue for stage 3.
|
||||||
|
|
||||||
For now the "work" is just logging the received filename to prove the flow —
|
Routing is the enqueue and nothing else: the file stays where intake wrote it and
|
||||||
this is the seam where real heavy-lifting will go later.
|
the *same* `Job` is handed on, so `job.path` still points at it. Reaching stage 2
|
||||||
|
is a fact about which queue holds the reference, not about which directory holds
|
||||||
|
the bytes.
|
||||||
|
|
||||||
|
Sub-`MIN_FILE_BYTES` files short-circuit to `:unknown` inside `classify`.
|
||||||
"""
|
"""
|
||||||
function handle_job(job::Job, cfg::Config, worker_id::Int)
|
function handle_classify_job(job::Job, cfg::Config, worker_id::Int,
|
||||||
# Classify the spooled file (annotate-only for now: the result is logged but
|
known_queue::JobQueue, unknown_queue::JobQueue,
|
||||||
# every file still moves to done/ regardless of known/unknown). Sub-32-byte
|
stats::StageStats)
|
||||||
# files short-circuit to :unknown inside classify without touching the model.
|
|
||||||
classification = classify(CLASSIFIER[], job.path)
|
classification = classify(CLASSIFIER[], job.path)
|
||||||
@info "received file" worker=worker_id id=job.id name=job.original_name size=job.size classification=classification
|
@debug "classified file" worker=worker_id id=job.id name=job.original_name size=job.size classification=classification
|
||||||
dest = move_to(cfg.done_dir, job)
|
|
||||||
@info "completed" worker=worker_id id=job.id dest=dest
|
if classification === :known
|
||||||
|
# known queue full → park and retry, don't drop (time charged to blocked_ns)
|
||||||
|
enqueue_blocking!(known_queue, job, stats;
|
||||||
|
retry_seconds = ROUTE_ENQUEUE_RETRY_SECONDS)
|
||||||
|
@debug "routed to enrichment" worker=worker_id id=job.id path=job.path
|
||||||
|
else
|
||||||
|
enqueue_blocking!(unknown_queue, job, stats;
|
||||||
|
retry_seconds = ROUTE_ENQUEUE_RETRY_SECONDS)
|
||||||
|
@debug "routed to content triage" worker=worker_id id=job.id path=job.path
|
||||||
|
end
|
||||||
return nothing
|
return nothing
|
||||||
end
|
end
|
||||||
|
|
||||||
"""
|
"""
|
||||||
worker_loop(worker_id, cfg, queue)
|
handle_known_job(job, cfg, worker_id)
|
||||||
|
|
||||||
Consume jobs until the queue is closed and drained. A failure on one job is
|
Stage 2. Extract metadata (exiftool, with timeout) and enrich: build the
|
||||||
logged and the file is quarantined in `failed/` — it must never kill the
|
normalized sidecar and commit both to `done/` sidecar-first: the file's one and
|
||||||
worker, or the pool would silently shrink.
|
only move, straight out of `spool/`. Extraction
|
||||||
|
failure/timeout yields a *degraded* sidecar (the file is still a wanted known
|
||||||
|
file), so the only way to land in `failed/` is a genuine I/O error writing the
|
||||||
|
sidecar or moving the file, handled by `worker_loop`'s quarantine.
|
||||||
"""
|
"""
|
||||||
function worker_loop(worker_id::Int, cfg::Config, queue::JobQueue)
|
function handle_known_job(job::Job, cfg::Config, worker_id::Int)
|
||||||
|
meta = build_metadata(job, cfg)
|
||||||
|
file_dest, sidecar = finalize_known!(cfg, job, meta)
|
||||||
|
@info "enriched" worker=worker_id id=job.id dest=file_dest sidecar=basename(sidecar) file_type=meta.file_type created_by=meta.created_by degraded=(meta.error !== nothing)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
handle_unknown_job(job, cfg, worker_id, text_queue)
|
||||||
|
|
||||||
|
Stage 3. Sort an unrecognized file into a coarse content bucket by sniffing its
|
||||||
|
first bytes.
|
||||||
|
|
||||||
|
The two outcomes are asymmetric, because one is terminal and one is not. Binary
|
||||||
|
is the end of the live path, so the file is committed to `binary/`, which is
|
||||||
|
also where the offline stage-5 discovery sweep reads its corpus, so the move is
|
||||||
|
load-bearing, not bookkeeping. Text has a stage 4 still to come, so nothing moves:
|
||||||
|
the same `Job` goes onto the language-enrichment queue, retrying on a full queue
|
||||||
|
rather than dropping the file (the blocking backpressure stage 1 also uses).
|
||||||
|
"""
|
||||||
|
function handle_unknown_job(job::Job, cfg::Config, worker_id::Int,
|
||||||
|
text_queue::JobQueue, stats::StageStats)
|
||||||
|
if is_binary(job.path)
|
||||||
|
dest = move_to(cfg.binary_dir, job)
|
||||||
|
@info "sorted unknown" worker=worker_id id=job.id name=job.original_name kind=:binary dest=dest
|
||||||
|
else
|
||||||
|
# text queue full → park and retry, don't drop (time charged to blocked_ns)
|
||||||
|
enqueue_blocking!(text_queue, job, stats;
|
||||||
|
retry_seconds = ROUTE_ENQUEUE_RETRY_SECONDS)
|
||||||
|
@info "routed to language enrichment" worker=worker_id id=job.id path=job.path
|
||||||
|
end
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
handle_text_job(job, cfg, worker_id, detector)
|
||||||
|
|
||||||
|
Stage 4. Enrich a text file with its natural language (via `detector`) and
|
||||||
|
programming language (via github-linguist): build the sidecar and commit both to
|
||||||
|
`text_done/` sidecar-first: the file's one and only move, straight out of
|
||||||
|
`spool/`. Detection failure yields a *degraded* sidecar (the
|
||||||
|
file is still wanted text), so the only way to land in `failed/` is a genuine I/O
|
||||||
|
error committing, handled by `worker_loop`'s quarantine.
|
||||||
|
"""
|
||||||
|
function handle_text_job(job::Job, cfg::Config, worker_id::Int, detector)
|
||||||
|
meta = build_text_metadata(detector, job, cfg)
|
||||||
|
file_dest, sidecar = finalize_text!(cfg, job, meta)
|
||||||
|
@info "enriched text" worker=worker_id id=job.id dest=file_dest sidecar=basename(sidecar) language=meta.language confidence=meta.language_confidence programming_language=meta.programming_language degraded=(meta.error !== nothing)
|
||||||
|
return nothing
|
||||||
|
end
|
||||||
|
|
||||||
|
"""
|
||||||
|
worker_loop(worker_id, cfg, queue, handler, stats)
|
||||||
|
|
||||||
|
Consume jobs from `queue` until it is closed and drained, running `handler` on
|
||||||
|
each. A failure on one job is logged and the file quarantined in `failed/`. It
|
||||||
|
must never kill the worker, or the pool would silently shrink.
|
||||||
|
|
||||||
|
This loop is also where per-stage metrics are recorded (`stats`, see
|
||||||
|
src/stats.jl). Instrumenting here rather than in each handler means every stage
|
||||||
|
is measured the same way, by construction, and a new stage is measured the
|
||||||
|
moment it is wired up, with no per-handler bookkeeping to forget.
|
||||||
|
|
||||||
|
The timed region is the handler alone, excluding the `dequeue!` above it: time
|
||||||
|
parked waiting for work is idleness, and counting it as service time would make
|
||||||
|
an idle stage look as busy as a saturated one. A quarantined job still counts
|
||||||
|
its time: the work was done, it just ended in `failed/`.
|
||||||
|
"""
|
||||||
|
function worker_loop(worker_id::Int, cfg::Config, queue::JobQueue, handler,
|
||||||
|
stats::StageStats)
|
||||||
@info "worker started" worker=worker_id
|
@info "worker started" worker=worker_id
|
||||||
while true
|
while true
|
||||||
job = dequeue!(queue)
|
job = dequeue!(queue)
|
||||||
job === nothing && break # queue closed and drained → exit
|
job === nothing && break # queue closed and drained → exit
|
||||||
|
Threads.atomic_add!(stats.in_flight, 1)
|
||||||
|
t0 = time_ns()
|
||||||
|
outcome = :ok
|
||||||
try
|
try
|
||||||
handle_job(job, cfg, worker_id)
|
handler(job, cfg, worker_id)
|
||||||
|
ack!(queue, job)
|
||||||
catch e
|
catch e
|
||||||
|
if !isfile(job.path)
|
||||||
|
# The file is gone and we did not move it: another worker
|
||||||
|
# committed it, i.e. this is the duplicate half of an
|
||||||
|
# at-least-once redelivery (src/rabbit.jl). The work is done, so
|
||||||
|
# settle the message and say nothing louder than @debug. It is
|
||||||
|
# emphatically not a failure: quarantining here would file a
|
||||||
|
# `failed/` entry against a file that succeeded, and counting it
|
||||||
|
# either way would double-count one file's work.
|
||||||
|
outcome = :duplicate
|
||||||
|
@debug "duplicate delivery; already committed by another worker" worker=worker_id id=job.id name=job.original_name
|
||||||
|
ack!(queue, job)
|
||||||
|
else
|
||||||
|
outcome = :failed
|
||||||
@error "processing failed" worker=worker_id id=job.id name=job.original_name exception=(e, catch_backtrace())
|
@error "processing failed" worker=worker_id id=job.id name=job.original_name exception=(e, catch_backtrace())
|
||||||
try
|
try
|
||||||
move_to(cfg.failed_dir, job)
|
move_to(cfg.failed_dir, job)
|
||||||
|
ack!(queue, job) # quarantined: the file's story ends in failed/
|
||||||
catch e2
|
catch e2
|
||||||
@error "could not quarantine failed file" worker=worker_id id=job.id path=job.path exception=(e2, catch_backtrace())
|
@error "could not quarantine failed file" worker=worker_id id=job.id path=job.path exception=(e2, catch_backtrace())
|
||||||
|
# Failed *and* unquarantinable: the file is still in spool/.
|
||||||
|
# Reject rather than ack, because we are not claiming this
|
||||||
|
# one is done. With no dead-letter exchange the broker
|
||||||
|
# discards it either way — `nack!` never requeues, so this
|
||||||
|
# cannot become a poison-message loop — but the distinction
|
||||||
|
# is free and the message is the honest one.
|
||||||
|
nack!(queue, job)
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
finally
|
||||||
|
# In a `finally` so an InterruptException during shutdown can't leave
|
||||||
|
# in_flight permanently above zero, which would read as a stuck job.
|
||||||
|
# A duplicate records nothing: its work was already counted by the
|
||||||
|
# delivery that actually did it.
|
||||||
|
outcome === :duplicate ||
|
||||||
|
record_job!(stats, outcome === :ok, job.size, Int(time_ns() - t0))
|
||||||
|
Threads.atomic_sub!(stats.in_flight, 1)
|
||||||
|
end
|
||||||
end
|
end
|
||||||
@info "worker stopped" worker=worker_id
|
@info "worker stopped" worker=worker_id
|
||||||
return nothing
|
return nothing
|
||||||
|
|||||||
1314
test/runtests.jl
Normal file
1314
test/runtests.jl
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user