✨ engine.run/submit of WorkGraphs

[GeigerJ2]

Requested quite a few people for review, as I think this is quite important. Not really needed that everybody goes through everything. But, in any case, good to get a few opinions on this. If you don't have the capacity to review, please feel free to remove the request.

I went through a few iterations before settling on the current approach. Documenting them here for context.

What we landed on

• common/workgraph.py centralises the optional-dependency handling. It exports a WORKGRAPH_INSTALLED sentinel and an is_workgraph_instance() helper, so launch.py itself never needs from aiida_workgraph imports and the accompanying try/except boilerplate.
• launch.py uses a positive guard clause (if WORKGRAPH_INSTALLED and is_workgraph_instance(process)) that branches into the WorkGraph path and returns early. The standard launch path is the fallthrough. Handling the WorkGraph path first will also be consistent with WorkGraph becoming the recommended way to write workflows in the future. workgraph: t.Any = process is needed to keep mypy happy since WorkGraph can't appear in the type union as an optional dependency.
• WorkGraph owns its own logic. All the business logic (input merging, metadata handling) lives in WorkGraph.prepare_for_launch() and WorkGraph.update_after_launch() on the aiida-workgraph side (PR Drop WG.submit/run for core engine.submit/run aiida-workgraph#768). aiida-core just calls these methods without knowing the internals.
• engine.submit() gains timeout — a new optional parameter that raises TimeoutError if the process doesn't terminate within the given seconds (only effective
  when wait=True). This ports functionality that previously lived in WorkGraph.submit() into the general-purpose launcher, where it benefits all process types (could be moved to separate PR).

Alternatives I considered and rejected

1. Direct imports in launch.py with try/except

Put from aiida_workgraph import WorkGraph directly in launch.py behind try/except ImportError.
Rejected: pollutes launch.py with optional-dependency import boilerplate, and every new call site would need its own guard.

2. Helper functions in launch.py (_resolve_workgraph_for_run, prepare_workgraph_inputs)

Kept the WorkGraph-specific business logic (input merging, metadata separation, collision checks) in aiida-core helper functions.
Rejected: violates Information Expert — WorkGraph should own its own preparation logic. Replaced by WorkGraph.prepare_for_launch() on the workgraph side (now PR aiidateam/aiida-workgraph#768 on aiida-workgraph).

3. Single is_workgraph_instance that catches ImportError internally

Combined the availability check and the isinstance check into one function that returns False when workgraph isn't installed.
Rejected: semantically dishonest — answering "is this a WorkGraph?" with "no, because workgraph isn't importable" conflates two different questions. Split into WORKGRAPH_INSTALLED sentinel + is_workgraph_instance() that assumes availability.

4. as_workgraph() with walrus operator

A function returning Optional[WorkGraph] used as if wg := as_workgraph(process):.
Rejected: still conflates the two questions (availability vs type), and the naming felt off.

5. Inverted guard clause: if not (WORKGRAPH_INSTALLED and is_workgraph_instance(process)): return runner.run(...)

Early-return for the non-workgraph path, leaving the workgraph path as the fallthrough.
Rejected: reads backwards — the negated condition just duplicates the fallthrough path and makes the control flow harder to follow.

Tests

The WorkGraph tests in test_launch.py are mock-based and serve as regression guards for the dispatch logic. They verify that when a WorkGraph is passed to run/run_get_node/run_get_pk, the launch functions correctly enter the WorkGraph code path (rather than the standard one), that the WorkGraph is asked to convert itself into a process class and inputs via prepare_for_launch, that the resulting process is actually launched through the runner, and that the process node is handed back to the WorkGraph via update_after_launch. They also check that both input styles (positional inputs dict and **kwargs) are forwarded correctly to prepare_for_launch.

Since aiida-workgraph is not installed in aiida-core's CI, these tests mock the detection logic and runner. They can't catch integration issues (e.g., whether prepare_for_launch produces inputs the engine actually accepts). The real end-to-end coverage comes from the companion PR on aiida-workgraph (aiidateam/aiida-workgraph#768), where the full test suite passes against this branch.

Companion PR on aiida-workgraph

The workgraph code changes live in aiidateam/aiida-workgraph#768, with the bulk of that PR being mechanical: replacing wg.run()/wg.submit() with engine.run(wg)/engine.submit(wg) across docs and tests. The files worth reviewing:

• src/aiida_workgraph/workgraph.py — the only real source change:

  • Adds prepare_for_launch(inputs, **kwargs): validates inputs, separates metadata from task inputs, checks for naming collisions, and returns (WorkGraphEngine, engine_inputs).
  • Adds update_after_launch(node): stores the process node on the WorkGraph and calls self.update() so the in-memory state reflects the launched process.
  • Removes the old WorkGraph.run() and WorkGraph.submit() methods.
  • timeout (parameter of the old wg.submit()) is now handled natively by engine.submit(wg, timeout=...).
  • interval (parameter of the old wg.submit()) raises a ValueError pointing to engine.submit(wg, wait=True, wait_interval=...).

• pyproject.toml / uv.lock — temporarily points aiida-core to the feature branch on my fork. Will revert to a release pin before merging.

• .github/workflows/ci.yaml — temporarily comments out the aiida-core version override. Same reason.

Full test suite passes. Only RTD fails because it can't resolve the git dependency.

[codecov]

Codecov Report

❌ Patch coverage is 78.00000% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.86%. Comparing base (5974f8b) to head (5d016c1).

Files with missing lines	Patch %	Lines
src/aiida/engine/launch.py	80.00%	8 Missing ⚠️
src/aiida/common/workgraph.py	70.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7261      +/-   ##
==========================================
+ Coverage   79.85%   79.86%   +0.01%     
==========================================
  Files         566      567       +1     
  Lines       43962    44005      +43     
==========================================
+ Hits        35102    35138      +36     
- Misses       8860     8867       +7     

☔ 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.