From 4a798f4bce29302c9e981c877753e083f80fe533 Mon Sep 17 00:00:00 2001 From: rob-maron <132852777+rob-maron@users.noreply.github.com> Date: Thu, 30 Jul 2026 16:29:27 -0400 Subject: [PATCH] improve polling for FLUX3 video gen (#75010) * wait between polls --- tests/tools/test_flux3_video_tool.py | 92 +++++++++++++++++++++++++--- tools/flux3_video_tool.py | 92 +++++++++++++++++++++++++--- 2 files changed, 167 insertions(+), 17 deletions(-) diff --git a/tests/tools/test_flux3_video_tool.py b/tests/tools/test_flux3_video_tool.py index caab52aa28e..29e01872bc3 100644 --- a/tests/tools/test_flux3_video_tool.py +++ b/tests/tools/test_flux3_video_tool.py @@ -44,10 +44,14 @@ class _FakeResponse: class _FakeClient: - """Captures the one request each handler makes.""" + """Captures each request a handler makes. + + A list of responses is served in order, with the last one repeating, so a + poll that looks twice can be given a job that finishes between looks. + """ def __init__(self, response, sink): - self._response = response + self._responses = list(response) if isinstance(response, list) else [response] self._sink = sink async def __aenter__(self): @@ -58,9 +62,10 @@ class _FakeClient: async def request(self, method, url, headers=None, json=None): self._sink.append({"method": method, "url": url, "headers": headers or {}, "json": json}) - if isinstance(self._response, Exception): - raise self._response - return self._response + response = self._responses[min(len(self._sink) - 1, len(self._responses) - 1)] + if isinstance(response, Exception): + raise response + return response class _FakeStream: @@ -115,6 +120,13 @@ def _run(coro): return asyncio.run(coro) +def _record_sleep(sink): + async def _sleep(seconds): + sink.append(seconds) + + return _sleep + + def _call(handler, args, response, headers=None): """Invoke a handler with the transport stubbed; returns (parsed, requests).""" sink = [] @@ -249,18 +261,82 @@ class TestSubmitTransport: assert "error" in parsed +@pytest.fixture(autouse=True) +def _no_real_poll_wait(monkeypatch): + """Keep the in-call wait out of the test clock.""" + monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 0) + + class TestPollTransport: - def test_poll_gets_the_job_and_returns_guidance(self): - response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."}) + def test_a_terminal_status_returns_at_once_without_waiting(self): + response = _FakeResponse(200, {"id": "bfl_job_1", "status": "Error", "guidance": "The job is over."}) parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response) assert requests[0]["method"] == "GET" assert requests[0]["url"] == f"{BASE_URL}/generations/bfl_job_1" assert requests[0]["json"] is None - assert parsed["result"] == "Still going." + assert len(requests) == 1 + assert parsed["result"] == "The job is over." + + def test_a_running_job_is_waited_out_inside_the_call(self, monkeypatch): + # A model has no clock, so telling it to pause produced a burst of polls + # instead of a paced one. The wait lives here where it cannot be skipped. + monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 45.0) + running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."}) + + slept = [] + with patch.object(flux3.asyncio, "sleep", new=_record_sleep(slept)): + parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, running) + + assert len(requests) == 2, "should look again after waiting" + assert sum(slept) == 45.0 assert parsed["details"]["status"] == "Generating" + def test_the_wait_is_answerable_to_a_stop(self, monkeypatch): + # Nothing outside the tool can end a call that has already started — + # the executor only checks for an interrupt between tools — so /stop + # has to land inside the wait rather than at the end of it. + monkeypatch.setattr(flux3, "_POLL_FOLLOW_UP_WAIT_SECONDS", 45.0) + from tools import interrupt as interrupt_module + + running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."}) + looks = [] + + def _stop_after_one_slice(): + looks.append(True) + return len(looks) > 1 + + slept = [] + with patch.object(interrupt_module, "is_interrupted", _stop_after_one_slice), \ + patch.object(flux3.asyncio, "sleep", new=_record_sleep(slept)): + parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, running) + + assert slept == [flux3._POLL_WAIT_SLICE_SECONDS], "the rest of the wait is abandoned" + assert len(requests) == 1, "and so is the second look" + assert parsed["details"]["status"] == "Generating" + + def test_the_call_returns_as_soon_as_the_job_finishes(self): + # The point of waiting in here is that the caller gets the result on the + # wait it was already taking, not one round trip later. + running = _FakeResponse(200, {"id": "bfl_job_1", "status": "Generating", "guidance": "Still going."}) + done = _FakeResponse(200, {"id": "bfl_job_1", "status": "Error", "guidance": "That job failed."}) + + parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, [running, done]) + + assert len(requests) == 2 + assert parsed["result"] == "That job failed." + + def test_a_refusal_is_returned_immediately_rather_than_waited_on(self): + # A 429 carries its own retry guidance; sleeping on it would only delay + # showing the model what to do, and spend the poll budget twice. + response = _FakeResponse(429, {"error": {"message": "Too many polls. Wait 30 seconds."}}) + + parsed, requests = _call(flux3._handle_get_result, {"id": "bfl_job_1"}, response) + + assert len(requests) == 1 + assert "Too many polls" in parsed["error"] + def test_ready_saves_the_clip_and_never_returns_the_signed_url(self, tmp_path): # The signed URL is a bearer credential for the clip and it used to be # re-keyed into a shell command by hand, dropping characters. Neither diff --git a/tools/flux3_video_tool.py b/tools/flux3_video_tool.py index 49dda1c86e9..b6415d703a0 100644 --- a/tools/flux3_video_tool.py +++ b/tools/flux3_video_tool.py @@ -179,6 +179,58 @@ _MEDIA_FIELDS = { _MAX_IMAGES = 10 +# How long one get_result call waits before its second look. The bound that +# matters is the whole call rather than this number: model_tools' async bridge +# abandons a tool at 300s, and one call spends two polls (each bounded +# server-side at 45s), this wait, and on Ready the download of the clip. There +# is room to roughly double this; the reason not to is that a finished job is +# only noticed at the next look, so the wait is also the notice delay. +_POLL_FOLLOW_UP_WAIT_SECONDS = 45.0 +# Taken in slices so the wait is answerable. Nothing outside a tool can end a +# call that has already started — the executor only checks for an interrupt +# between tools — so a tool that blocks this long has to watch the flag itself. +_POLL_WAIT_SLICE_SECONDS = 1.0 + +# Mirrors the gateway's BFL statuses +_TERMINAL_POLL_STATUSES = frozenset( + {"Ready", "Error", "Request Moderated", "Content Moderated", "Task not found"} +) + + +def _poll_is_finished(raw: str) -> bool: + """True when there is nothing to wait for: done, refused, or unreadable.""" + try: + payload = json.loads(raw) + except Exception: + return True + if not isinstance(payload, dict) or "error" in payload: + # A refusal carries its own guidance (a wait, a limit, a dead job). + # Sleeping on it would only delay showing the model what to do. + return True + details = payload.get("details") + status = details.get("status") if isinstance(details, dict) else None + return not isinstance(status, str) or status in _TERMINAL_POLL_STATUSES + + +async def _wait_before_second_look() -> bool: + """Hold the call open between looks; False if the user interrupted. + + Counted down rather than clock-driven: this paces polling, so a slice + that runs long changes nothing, and the loop stays testable without a + fake clock. + """ + from tools.interrupt import is_interrupted + + remaining = _POLL_FOLLOW_UP_WAIT_SECONDS + while remaining > 0: + if is_interrupted(): + return False + this_slice = min(_POLL_WAIT_SLICE_SECONDS, remaining) + await asyncio.sleep(this_slice) + remaining -= this_slice + return True + + def _warm_nous_token() -> None: """Refresh the Nous token once, before any parallel upload needs it. @@ -471,8 +523,25 @@ async def _handle_get_result(args: dict, **kwargs) -> str: return _error("BFL video generation is not available in this build.") from urllib.parse import quote - raw = await _call_gateway("GET", f"{endpoints['base_url']}/generations/{quote(job_id.strip(), safe='')}") - return await _save_if_ready(raw, (args or {}).get("save_to")) + url = f"{endpoints['base_url']}/generations/{quote(job_id.strip(), safe='')}" + save_to = (args or {}).get("save_to") + + raw = await _call_gateway("GET", url) + if _poll_is_finished(raw): + return await _save_if_ready(raw, save_to) + + # Still running, so absorb the wait here instead of asking the model to + # take it. A model has no clock: told to wait it emits "I'll wait a minute" + # and its next action lands immediately, so the guidance produced a burst of + # polls rather than a paced one. Waiting inside the call cannot be skipped, + # needs no shell, and works the same on every platform. One call therefore + # covers a couple of minutes and returns as soon as a look finds it done. + if not await _wait_before_second_look(): + # Interrupted mid-wait: hand back the status we already have rather + # than spending a round trip the user has just asked us to stop for. + return await _save_if_ready(raw, save_to) + raw = await _call_gateway("GET", url) + return await _save_if_ready(raw, save_to) async def _handle_prompting_guide(args: dict, **kwargs) -> str: @@ -682,8 +751,9 @@ GET_RESULT_SCHEMA = { "name": "bfl_flux3_get_result", "description": ( "Poll a FLUX 3 video job by the job id a generate tool returned. Generation takes minutes " - "and a long Generating phase is normal. Every response states the job's status, how long " - "to wait before polling again, and what you may do next — follow it rather than guessing. " + "and a long Generating phase is normal. This call waits for you while the job runs, so it " + "may take a couple of minutes; if it returns still generating, just call it again. Do not " + "sleep between calls. " "On Ready the clip is downloaded for you and the response gives its local path; your only " "remaining step is to deliver that file as the response describes." ), @@ -813,11 +883,15 @@ auto-downscaled and output tops out at 720p. Submit returns a job id immediately — the video does not exist yet. Poll bfl_flux3_get_result with that id; generation takes several minutes and a long Generating phase is normal, not a stall. Nothing reaches disk before the job is -Ready, so checking folders mid-run tells you nothing. Every response -states the job status, how long to wait before polling again, and what you may -do next — follow it literally rather than guessing an interval. A job survives -client restarts: re-poll the same id rather than resubmitting, which would only -spend your budgets on duplicate work. +Ready, so checking folders mid-run tells you nothing. + +The waiting is not yours to do. bfl_flux3_get_result takes the pause itself +while a job is still running, so one call can occupy a couple of minutes and +comes back the moment the job finishes. If it returns still generating, just +call it again — no sleeping, no interval to judge, nothing to time. + +A job survives client restarts: re-poll the same id rather than resubmitting, +which would only spend your budgets on duplicate work. ## Save and deliver