Skip to main content

Selective module loading

Status: design proposal. Not yet implemented.

Motivating idea: when a user runs a graph that touches only a handful of node families, the worker subprocess shouldn't pay the import cost for every other family in the registry.

The problem​

EdgeWeave runs each graph execution in a fresh subprocess worker. That gives us crash isolation, graph-level cancellation, and a clean namespace per run, but the price is a fresh import sweep every time. Today the worker imports every node module at startup — including heavyweights like torch, sklearn, simweave[viz], the (eventual) larger sklearn/onnx/jax integrations — even when the user's graph only contains a handful of numpy_* nodes.

Concrete cost on a typical desktop today:

ModuleCold import cost (rough)
numpy + pandas~0.4 s
plotly~0.3 s
simweave (no extras)~0.1 s on top of numpy
simweave[viz]adds plotly to the above
torch (CPU build)3 – 6 s
torch + torchvision5 – 8 s
sklearn~1.5 s

For an interactive "tweak → run → look" loop on a numpy-only graph, paying 5 s for torch on every run is a real UX hit, and that gap will only widen as the core library grows.

Goals​

  1. Pay only for what you use. A graph that doesn't reference any torch nodes should not import torch.
  2. Keep the SDK ergonomic. Authors should still write @register_block(...) at module top level — no manual registration bookkeeping.
  3. Hot reload still works for projects/demo_project/modules/user_blocks/.
  4. No cost shifted to graph-edit time. The user must not see node palette flicker or per-keystroke slow-downs in the frontend.
  5. Backwards compatibility. Existing .weave files keep loading; no data migration required.

Current behaviour (baseline)​

  • backend/core/nodes/ holds shipped node families (numpy_nodes.py, torch_nodes.py, plotly_nodes.py, etc.).
  • projects/demo_project/modules/user_blocks/ holds user-authored families, hot-reloaded at edit time via backend/core/reloader.py.
  • At backend / worker startup, NODE_REGISTRY is populated by importing every module under those folders. @register_block side-effects fire and registry entries are inserted.
  • The worker then runs the graph; per-node handlers are simple callables stored on the registry entry — no further importing needed.

So the "expensive" thing happens entirely at module-import time; once a module is loaded, dispatching to its handlers is essentially free.

Constraints that shape the design​

  • Registration is a side effect of import. We can't enumerate the node types a module provides without importing it, unless we add an out-of-band manifest.
  • Workers are cheap to spawn but hot to populate. Forking is not available on Windows, so we can't fork() from a pre-warmed parent.
  • Hot reload depends on watchdog re-importing modules at edit time. Whatever scheme we pick must coexist with that.

Proposed approach: a node manifest​

Add a build-time / startup-time manifest that records, for each module, the set of node-type strings it registers — without us having to import the module to find out.

File layout​

backend/core/nodes/
  numpy_nodes.py
  torch_nodes.py
  ...
  __manifest__.json   # generated; one entry per module file
projects/demo_project/modules/user_blocks/
  simweave_continuous.py
  simweave_discrete.py
  __manifest__.json   # regenerated on hot reload

Manifest schema (sketch)​

{
  "schema_version": 1,
  "generated_at": "2026-04-25T17:43:00Z",
  "modules": [
    {
      "module": "backend.core.nodes.numpy_nodes",
      "path": "backend/core/nodes/numpy_nodes.py",
      "provides": [
        "numpy_arange",
        "numpy_linspace",
        "numpy_random_normal"
      ],
      "import_cost_estimate_ms": 420
    },
    {
      "module": "backend.core.nodes.torch_nodes",
      "path": "backend/core/nodes/torch_nodes.py",
      "provides": [
        "torch_conv2d_layer",
        "torch_maxpool2d_layer",
        "torch_relu_layer",
        "torch_sequential",
        "torch_graph_model",
        "..."
      ],
      "import_cost_estimate_ms": 5200
    }
  ]
}

How it gets populated​

A small discovery pass. Two options, in increasing order of effort and robustness:

  1. Runtime discovery (simplest). A tools/build_manifest.py script that imports every module in a freshly-spawned subprocess, captures the registrations into a temporary registry, snapshots them to __manifest__.json, and exits. Run on:

    • Backend bootstrap if the manifest is missing / older than any module file.
    • Hot-reload events for the user_blocks folder.
    • CI, as a sanity check that all modules import cleanly.

  2. Static AST scan (later, if we need to skip runtime imports). Walk each .py, parse its AST, extract the string literal in every @register_block("name", ...). Brittle — anything that builds a node type at runtime (e.g. a factory loop) breaks. Use only as a fallback for modules that can't be imported safely on the build host (e.g. the manifest builder doesn't have CUDA but a cuda_nodes.py requires it).

Recommendation: start with option 1, fall back to option 2 only if a module fails to import in the build environment.

How the worker uses the manifest​

Worker boot sequence becomes:

1. Read the .weave being executed.
2. Collect the set S of node-type strings it references.
3. Read __manifest__.json files (cheap; ~few KB total).
4. For each manifest entry M:
       if M.provides intersects S:
           importlib.import_module(M.module)
5. Resolve handlers from NODE_REGISTRY, run the graph.

Steps 3–4 take a few ms; step 4's import calls are the only expensive piece, and they only happen for the modules the graph actually needs.

For a numpy-only graph, this saves the ~5 s torch import outright. For a torch graph, we still pay the import once, but we no longer additionally pay for sklearn / simweave / etc.

Hot reload coexistence​

backend/core/reloader.py watches the user_blocks folder. On change:

  1. Re-run the manifest builder against the user_blocks folder only.
  2. Patch the user_blocks portion of __manifest__.json.
  3. Invalidate the affected modules in sys.modules (existing logic).
  4. Next graph run lazily re-imports them via the new manifest.

No frontend change required — the palette already pulls node types from NODE_REGISTRY over the existing REST endpoint.

Open questions​

  1. Manifest-out-of-sync story. If a developer edits a node file without re-running the manifest builder, the worker will silently fail to import the module on next run. Mitigations:
    • Stat the source files and rebuild on mtime mismatch (cheap; the manifest builder runs in milliseconds for small modules).
    • Refuse to start the worker if any manifest entry is older than its module file.
  2. Frontend palette population. Today the palette enumerates NODE_REGISTRY. With selective loading, the backend FastAPI process still needs to know every node type for tooltips, the palette, and validation, but the worker subprocess doesn't. Easiest split: the FastAPI process keeps importing everything (paid once at backend startup); the worker uses the manifest. That already gets us 90 % of the win — graph runs, not the always-running backend, are the expensive thing.
  3. Per-node handler-side imports. Some nodes (torch_*) carry their imports at module top. Others could be refactored to import inside the handler so even an "imported" module doesn't drag the heavy library in. Lower priority — manifest-based loading already covers the typical case.
  4. Multiple workers. If we ever pre-warm a small pool of workers, each warm worker would have only a subset of modules loaded. We'd either dispatch graphs to a "compatible" worker or accept that a worker may need to import an extra module mid-run. The manifest already supports the latter.
  5. Plugins / third-party node packs. Same shape: each package ships a manifest, the worker loads only what's referenced. Drop-in.

Estimated payoff vs. effort​

PhaseEffortSaving on numpy-only graphSaving on torch graph
Build the manifest tool~1 dayn/an/a
Wire worker to use the manifest~1 day~5 s per run~0.5 s per run
Integrate with hot reload~half dayunchangedunchanged
Static-AST fallback (optional)~1 daydittoditto

Net: the basic implementation is small (2 – 3 days) and removes a visible UX wart that will only get worse as the core library grows. A reasonable PR plan:

  1. Add tools/build_manifest.py and a CLI test (no behavioural change).
  2. Wire the worker to read the manifest and import selectively, but keep the old "import everything" path behind a config flag for one release.
  3. Flip the default once we're satisfied; remove the old path.
  • backend/core/reloader.py — hot-reload mechanism for user_blocks.
  • sdk/decorators.py — where @register_block populates NODE_REGISTRY.
  • backend/utils/parse_drawflow.py — already does the set-of-node-types collection; the manifest worker would reuse this.