Skip to content

Add flyte.is_control_plane_available() and Run.first_failure() for agent edit-and-relaunch loops - #1522

Merged
cosmicBboy merged 8 commits into
mainfrom
nielsb/agent-fork-sdk-utils
Sep 3, 2026
Merged

Add flyte.is_control_plane_available() and Run.first_failure() for agent edit-and-relaunch loops#1522
cosmicBboy merged 8 commits into
mainfrom
nielsb/agent-fork-sdk-utils

Conversation

@cosmicBboy

@cosmicBboy cosmicBboy commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

Why

The agent-mediated forking example (unionai/unionai-examples#308) — an agent task that launches a workflow, observes the failure, patches the workflow's source on disk, reloads it, and fork()s the failed run with the fixed code — had to reach into SDK internals:

  1. isinstance(get_controller(), RemoteController) to decide whether launching/forking real runs is possible, or whether to fall back to running the workflow inline.
  2. Hand-rolled Action.listall(...FAILED...) iteration to find which step failed and its error message.
  3. flyte._code_bundle.bundle.build_code_bundle.cache_clear() to make the next fork ship the edited working tree — no longer needed: fix(code-bundle): rebuild after source changes #1508 removed the code-bundle memoization entirely, so every launch re-bundles from disk. This PR originally shipped a flyte.refresh_code_bundle_cache() for that; it was dropped when merging main, since the cache it managed no longer exists.

What

flyte.is_control_plane_available() — True when the process can submit work to a control plane (launch real runs whose actions can be awaited and replayed/forked). Inside a task, the orchestration mode decides (remote/hybrid → True, local → False, even when a client is configured — flyte run --local configures one too); outside a task, a configured client decides. This replaces the isinstance-on-internal-controller probe with the TaskContext.mode the runtime already maintains.

Run.first_failure() and ActionDetails.error_message — the observation half of a repair loop: which step of a run failed, and why. first_failure() returns the ActionDetails of the first failed action in creation order, preferring a failed sub-action over the failed root (whose error usually just repeats the sub-action's); error_message is the failed action's message or "".

With these, the example's loop reduces to:

run = await flyte.run.aio(wf.main, n_records=n_records)
await run.wait.aio(quiet=True)
if failure := await run.first_failure.aio():
    patch_workflow_source(failure.task_name, failure.error_message)
    importlib.reload(sys.modules["workflow"])
    run = await fork.aio(run.name, task_template=wf.main)

Testing

  • New unit tests: is_control_plane_available() across uninitialized/client/local/remote/hybrid contexts; Run.first_failure() sub-action preference, root fallback, and no-failure; error_message.
  • Existing code_bundle, remote, deploy, and cli/test_run.py suites pass (the only failures are pre-existing on main: the two loaded_modules discovery tests).
  • ruff, mypy, and ty clean via pre-commit hooks.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NwQixBcyR5va6BC75jaQx3

…ted forking)

Three additions that make the "agent task edits a workflow, reloads it, and
forks the failed run" pattern native, instead of relying on SDK internals:

- flyte.refresh_code_bundle_cache(): forget in-process memoized code bundles
  so the next run/deploy/serve (or fork) re-bundles the working tree as it is
  on disk now. Also make skip_cache=True (what disable_run_cache threads
  through) bypass the in-process memo, not just the persistent SQLite cache.
  build_code_bundle*.cache_clear() is preserved for existing callers.

- flyte.control_plane_available(): True when the process can submit work to a
  control plane. Inside a task it is decided by the orchestration mode
  (remote/hybrid vs local) rather than isinstance checks on the internal
  controller; outside a task, by whether a client is configured.

- Run.first_failure() + ActionDetails.error_message: the observation half of a
  repair loop - which step of a run failed, and why - preferring the failed
  sub-action over the root action whose error just repeats it.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NwQixBcyR5va6BC75jaQx3
@cosmicBboy
cosmicBboy marked this pull request as ready for review September 2, 2026 16:55
Main removed the alru_cache on build_code_bundle entirely (rebuild after
source changes), which solves the stale-bundle problem at the root. Drop
this branch's refresh_code_bundle_cache() and skip_cache memo-bypass —
they managed a cache that no longer exists. control_plane_available()
and Run.first_failure() / ActionDetails.error_message remain.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NwQixBcyR5va6BC75jaQx3
@cosmicBboy cosmicBboy changed the title Add SDK utilities for in-process edit-and-relaunch loops (agent-mediated forking) Add flyte.control_plane_available() and Run.first_failure() for agent edit-and-relaunch loops Sep 2, 2026
@cosmicBboy cosmicBboy changed the title Add flyte.control_plane_available() and Run.first_failure() for agent edit-and-relaunch loops Add flyte.control_plane_is_available() and Run.first_failure() for agent edit-and-relaunch loops Sep 2, 2026
@cosmicBboy cosmicBboy changed the title Add flyte.control_plane_is_available() and Run.first_failure() for agent edit-and-relaunch loops Add flyte.is_control_plane_available() and Run.first_failure() for agent edit-and-relaunch loops Sep 3, 2026
@cosmicBboy
cosmicBboy merged commit 1a0d78f into main Sep 3, 2026
64 checks passed
@cosmicBboy
cosmicBboy deleted the nielsb/agent-fork-sdk-utils branch September 3, 2026 01:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant