Semi-lazy imports + auto-mode forkserver preload

[dodamih]

Summary

Replaces the all-or-nothing forkserver preload with a CUE-driven auto mode that loads only the modules a given spec actually needs. Adds a static @register index, a lookup-miss fallback for late-bound names, dynamic resolvers for np.* / torch.*, and a shared make_lazy_module helper used to lazify ~25 subpackage __init__.py files. Wires the computed preload list to remote workers via env var. Flips the CLI default to auto.

Motivation

zetta run and setup_environment("all") used to import every preload bundle eagerly (~1.4 GB RSS, ~12-15 s on a typical laptop) regardless of what the spec needed. Tri's debug iteration in particular hits this on every CUE typo. After this PR:

Scenario	Before	After
import zetta_utils (bare)	~1.4 s	~1.2 s (no numpy/torch)
Fast-fail on CUE typo	~12-15 s wasted	<2 s
Trivial spec (no torch/mazepa)	~12 s / 1.4 GB	~2-3 s / 105 MB
Production spec (test3.cue: internal alignment + mazepa)	~12 s / 1.4 GB	~5-7 s / 105 MB parent (-89% RSS)
Worker daemon RSS (workers fork from this)	~1.4 GB	~770-830 MB (~600 MB CoW saving per pod)

Everything else stays opt-out: zetta run -l all still works for any case where lazy is undesirable.

Architecture

Phase 0: Static @register index (scan.py)

AST-walk the package looking for @register('name') decorators and register('name')(target) direct calls. Build a name → module index without executing any imports. Disk-cached at ~/.cache/zetta_utils/builder_index.json. Cold scan ~220 ms; warm <50 ms.

Phase 1: Lookup-miss fallback (registry.py)

get_matching_entry consults the index when REGISTRY misses, imports the candidate module(s), retries. Idempotent, thread-safe. Behavior is unchanged when REGISTRY is fully populated; the fallback only fires on lookups that would have raised.

Phase 2: none preload mode

Skip the eager load_*_modules() step in the parent. Daemon still spawns eagerly with preload/none.py (always-eager: parallel, log, builder). Lookup-miss fallback handles everything else on demand.

Phase 3: auto preload mode

Parse the CUE, walk it for @type literals (parsing/spec_scan.py), resolve through the static index to a minimal module set (builder/preload/init.py: compute_preload_set), feed that list directly to multiprocessing.set_forkserver_preload(...) — no preload shim file needed.

Phase 4: Lazy subpackage __init__.py (common/lazy.py)

A package that re-exports many submodules pays a transitive import cost on first import package even when the caller wants one symbol. make_lazy_module(__name__, globals(), subpackages, reexports_by_module) returns __getattr__ + __dir__ that resolve names on first access. Applied to:

• mazepa, mazepa_layer_processing, layer, tensor_ops, augmentations, db_annotations, message_queues, task_management, training
• cloud_management, cloud_management/resource_allocation/{__init__, gcloud, k8s}
• convnet, convnet/architecture, convnet/architecture/deprecated

The helper falls back to importlib.import_module(\".{name}\", pkg) for unknown names, mirroring Python's natural submodule attribute behavior — so pkg.X.Y works whether or not X is explicitly declared.

Dynamic resolvers (built_in_registrations.py)

built_in_registrations.py and convnet/architecture/primitives.py previously enumerated dir(numpy) and dir(torch.nn) etc. at import time, registering ~1400 names + pulling numpy + torch. Now _resolve_numpy(\"np.foo\") and _resolve_torch(\"torch.nn.Bar\") do getattr(numpy, \"foo\") / getattr(torch.nn, \"Bar\") on first lookup. Net: numpy and torch are no longer in the always-eager set; they get loaded only when first needed.

Phase 5: Remote workers (cloud_management/resource_allocation/k8s)

Workers spawned by execute_on_gcp_with_sqs previously used --load_mode try and let each worker load all preloaded subpackages itself. With auto mode, the master computes the minimal preload set from the original CUE and ships it via a ZETTA_PRELOAD_MODULES env var on the worker pod. The CLI reads the env var on startup and passes the explicit list to setup_environment(load_mode='auto', preload_modules=[...]). Master's _compute_worker_preload_modules re-scans ZETTA_RUN_SPEC_PATH to derive the list.

Default flip

--load_mode default is now auto. -l all still available as the safety hatch.

Files changed

New:

• zetta_utils/builder/scan.py — AST-driven static index
• zetta_utils/builder/preload/none.py — minimal preload bundle
• zetta_utils/parsing/spec_scan.py — extract_types for CUE specs
• zetta_utils/common/lazy.py — make_lazy_module helper
• tests/unit/builder/test_scan.py (17 tests)
• tests/unit/builder/test_lazy_fallback.py (11 tests)
• tests/unit/builder/test_index_parity.py (3 tests)
• tests/unit/builder/test_none_mode.py (4 tests)
• tests/unit/builder/test_auto_mode.py (2 tests)
• tests/unit/builder/test_compute_preload.py (5 tests)
• tests/unit/builder/test_resolvers.py (15 tests)
• tests/unit/parsing/test_spec_scan.py (6 tests)
• tests/unit/common/test_lazy.py (7 tests)

Modified:

• zetta_utils/__init__.py — setup_environment(load_mode, cue_path, preload_modules); new auto and none modes
• zetta_utils/builder/{registry, __init__, built_in_registrations}.py — fallback path, dynamic resolvers
• zetta_utils/builder/preload/__init__.py — compute_preload_set
• zetta_utils/cli/main.py — --load_mode auto default; reads ZETTA_PRELOAD_MODULES env var
• zetta_utils/cloud_management/resource_allocation/k8s/{pod, deployment}.py — thread preload_modules to worker pod env
• zetta_utils/mazepa_addons/configurations/execute_on_gcp_with_sqs.py — compute and ship preload list
• zetta_utils/convnet/architecture/primitives.py — drop torch.* import-time loops
• 16 __init__.py files lazified via the helper

Behavior changes worth knowing about

1. setup_environment(\"all\") no longer eagerly cascades into every submodule. The lazified __init__.py files mean from zetta_utils import mazepa only loads the package object, not its submodules. The lookup-miss fallback fills the gap on first use. Existing tests that assumed REGISTRY was "fully populated" after setup_environment(\"all\") may need updating to use a name from the always-eager set (the parity test was updated to use lambda instead of brightness_aug).

2. CLI default flip. zetta run path.cue defaults to auto mode. -l all is the rollback path.

3. internal/__init__.py chunkedgraph try/except. Was eager (logged a warning at import time if missing); now lazy (raises ModuleNotFoundError only on first access of build_chunkedgraph_segmentation_flow).

4. Worker --load_mode try still hardcoded in get_mazepa_worker_command for back-compat. When ZETTA_PRELOAD_MODULES is set by the master, the CLI overrides to auto. Workers without the env var fall back to try (current behavior).

Test plan

• Unit: 100% line coverage on the 12 files I authored/modified for this work (scan.py, registry.py, built_in_registrations.py, preload/{__init__, all, core, inference, none, training}.py, parsing/spec_scan.py, common/lazy.py, cli/main.py). 2 lines in scan.py carry # pragma: no cover for filesystem-race OSError paths that aren't reproducible in unit tests.
• Local full suite: 2008/2009 tests pass; the 1 pre-existing failure is a docker-fixture environmental issue.
• End-to-end smoke:
  • setup_environment(\"none\") then build() of a non-eager spec → succeeds via lookup-miss fallback
  • setup_environment(\"auto\", cue_path=...) → daemon preloads exactly the computed set
  • Subprocess with ZETTA_PRELOAD_MODULES env var → daemon preload list matches the env var
• CI green
• Production validation: zetta run an internal CUE in a dev k8s cluster, confirm worker pods get the env var and resource usage drops

Pairs with

internal #198 — must merge after this lands so the lazy __init__.py updates in internal/ align with the helper this PR introduces.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (f1a2436) to head (615f56e).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #1253      +/-   ##
===========================================
+ Coverage   99.98%   100.00%   +0.01%     
===========================================
  Files         206       211       +5     
  Lines       10920     11292     +372     
===========================================
+ Hits        10918     11292     +374     
+ Misses          2         0       -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[nkemnitz]

Pending tests and coverage