fix(app): cold-replica pickle-by-value + reused-process module purge (0.12.0 transport hardening)

[nilsmechtel]

Two replica-side follow-ups to 0.12.0's move of app source off Ray-GCS onto per-file Hypha sync (#148). Both are gated on that transport and were surfaced by it. Ships as 0.12.1.

---

Fix 1 — cold replica must not import app source during cloudpickle.loads (the KTH cellpose outage)

Symptom. A cold Serve replica of a real app (cellpose; composed model-runner) dies in ServeReplica.__init__:

replica.py:2671  deployment_def = cloudpickle.loads(serialized_deployment_def)
ModuleNotFoundError: No module named 'main'

Cause. 0.12.0 stopped shipping app source on Ray py_modules, so <app_dir>/source is not on the replica's sys.path when Ray runs cloudpickle.loads. Ray Serve pickles a deployment's class by name (ray.cloudpickle) whenever that class's methods reference a module-level symbol of their own source module — a helper function/class, of which any monolith like cellpose has many. Unpickling such a reference does import main/import entry before any bioengine code runs, so before the sys.meta_path finder that would materialise the source is installed → ModuleNotFoundError. The worker_process_setup_hook (the intended pre-import bridge) is not honoured before a Serve replica's unpickle, confirmed on KTH: the runtime_env carried the hook, yet it produced no output before the crash. Both prod apps only survived 0.12.0 as warm-adopted replicas (never re-pickle-loaded); the first genuine cold pickle-load failed. A trivial @bioengine.app class with no module-level self-references pickles fully by value and loads cold fine — which is why simple apps and the pre-merge risk-A test never surfaced it.

Fix (bioengine/_app/bootstrap.py, build_and_run_application). After binding, walk sys.modules and ray.cloudpickle.register_pickle_by_value every module whose __file__ is under <app_dir>/source (covers subpackages like model-runner's model_cache/). The deployment then serializes by value and the replica reconstructs it without importing user modules. Lazy in-method imports still resolve via the finder — installed when the reconstructed class's @bioengine wrappers import bioengine on load.

Fix 2 — reused warm worker process serves the previous deploy's code

Symptom. A redeploy that lands a fresh replica on a reused Ray worker process serves the previous deploy's code even though <app_dir>/source on disk is fresh (model-runner's tagging fix looked deployed but ran old logic on KTH).

Cause. Per-file Hypha sync refreshes the source on disk, but the reused interpreter still has the old app modules in sys.modules, so a lazy import entry returns stale code. Ray reuses a warm process when the runtime_env is byte-identical across deploys — true for public artifacts (no per-deploy download token → identical env → warm GPU process reused). Private apps mint a fresh token each deploy, so they dodge it. Env-var changes (e.g. BIOENGINE_ARTIFACT_VERSION) do not force Ray to recycle a warm process, so this must be fixed replica-side. The old content-hashed GCS py_module incidentally forced a fresh process on any source change, masking this until 0.12.0 removed it.

Fix (bioengine/_app/mixin.py, _purge_stale_app_modules). Run per replica instance in _setup_replica, before the user's __init__: drop every sys.modules entry whose file lives under <app_dir>/source, so the next import re-reads the freshly-synced source. No-op on a genuinely fresh process; framework/stdlib untouched. Pairs with the sync's fresh mtimes (forces a real recompile).

---

Validation

• Unit tests (both green; full tests/_app/ = 87):
  • tests/_app/test_pickle_by_value_cold_replica.py — a source module whose entry class references a module-level symbol, loaded in a clean subprocess (cold replica): by-reference → No module named 'entry'; after register_pickle_by_value → loads + constructs.
  • tests/_app/test_replica_module_purge.py — stale app module + fresh source on disk → purge drops it → re-import reads new code; framework modules untouched.
• cloudpickle-layer proof — clean-interpreter load of a real @bioengine.app deployment reproduces the exact ModuleNotFoundError: No module named 'main' by-reference, and the by-value pickle loads.
• Isolated KTH cold gate (0.12.1-dev2 on a separate test worker in a non-prod workspace, adoption-isolated from prod; prod stayed on 0.11.29 throughout):
  • Monolith cold deploy (module-level ref + lazy in-method import) → RUNNING; marker() returns the computed + lazily-imported values (layers 1 & 2).
  • Composed model-runner cold import entry → RUNNING, lazy model_cache subpackage resolved at call time (layers 1 & 2 on the composed shape).
  • In-process cutover (Fix 2): same application_id redeploy v1.15.24 → v1.15.26, no stop / no restart, on the warm-public-process condition → serves the new tagging code. The exact prod failure mode, proven fixed without a restart.

Rollout

bioengine/** → PR + dev-image workflow. Version bumped to 0.12.1 on this branch. After merge + canonical publish, KTH is upgraded 0.11.29 → 0.12.1 (restoring the Hypha transport) and the model-runner 1.15.26 in-process cutover proceeds. Never merge — left to the maintainer.

🤖 Generated with Claude Code