docs: keep pages that finish while a translate run is stopping

When the API rejects the credentials part-way through a parallel run, stop
starting pages but still collect the ones already in flight and write those
that succeed, then report usage and exit 2. Previously a page that had
already finished in the same batch could be dropped depending on set
iteration order. The test fake's call counter is now guarded by a lock.

Author: maxisbey

scripts/docs/translations.py

@@ -988,18 +988,23 @@ def produce(inputs: Inputs, job: Job) -> tuple[str | PageError, Usage]:
     usage, failed = Usage(), False
     queue = deque(work)
     running: dict[Future[tuple[str | PageError, Usage]], tuple[Language, Job]] = {}
+    stopped: ConfigError | None = None  # rejected credentials: start nothing more, keep what still lands
     pool = ThreadPoolExecutor(max_workers=args.jobs)
     try:
         # The window is refilled here rather than handing the pool every page up front, so
-        # once a page raises (rejected credentials) no further page starts; the ones already
-        # in flight run out unreported, and the process exits when they have.
-        while queue or running:
-            while queue and len(running) < args.jobs:
+        # that once the credentials are rejected no further page starts; pages already in
+        # flight are still collected, and written if they made it, before the run stops.
+        while (queue and stopped is None) or running:
+            while queue and stopped is None and len(running) < args.jobs:
                 inputs, job = queue.popleft()
                 running[pool.submit(produce, inputs, job)] = (inputs.language, job)
             for future in wait(running, return_when=FIRST_COMPLETED).done:
                 language, job = running.pop(future)
-                result, spent = future.result()
+                try:
+                    result, spent = future.result()
+                except ConfigError as exc:
+                    stopped = stopped or exc
+                    continue
                 usage.add(spent)
                 page = job.state.page
                 if isinstance(result, PageError):
@@ -1019,6 +1024,8 @@ def produce(inputs: Inputs, job: Job) -> tuple[str | PageError, Usage]:
     finally:
         pool.shutdown(wait=False)
     print(f"usage: {usage}")
+    if stopped is not None:
+        raise stopped
     return 1 if failed else 0
 
 

tests/docs/test_translations.py

@@ -900,7 +900,10 @@ def test_rejected_credentials_stop_the_run_with_exit_2(tmp_path: Path, capsys: p
     assert (code, out, err) == snapshot(
         (
             2,
-            "translating 4 pages (ja), 1 at a time, with test-model\n",
+            """\
+translating 4 pages (ja), 1 at a time, with test-model
+usage: 0 input / 0 output / 0 cache-write / 0 cache-read tokens
+""",
             "translations: the API rejected the credentials: invalid x-api-key\n",
         )
     )
@@ -957,30 +960,37 @@ def complete(self, *, model: str, system: str, messages: Sequence[t.Message], ma
     assert run(capsys, root, "status") == (0, "ja (日本語): 0 missing, 0 outdated, 4 current, 0 removable\n", "")
 
 
-def test_rejected_credentials_with_pages_in_flight_start_no_further_page(
+def test_rejected_credentials_with_pages_in_flight_start_no_further_page_but_a_finished_one_is_kept(
     tmp_path: Path, capsys: pytest.CaptureFixture[str]
 ) -> None:
-    """Tool-defined: with `--jobs 2`, when the API rejects the credentials (here for both pages in flight,
-    once both are waiting) the run stops with exit 2 and the two remaining pages are never requested."""
+    """Tool-defined: with `--jobs 2`, when the API rejects the credentials for one page while another is in
+    flight (here both are waiting before either is answered), the other page still lands and is written,
+    the two remaining pages are never requested, usage is reported, and the run stops with exit 2."""
     root = make_repo(tmp_path)
 
-    class Rejected:
-        both, calls = threading.Barrier(2, timeout=5), 0
+    class RejectsHome:
+        def __init__(self) -> None:
+            self.both, self.lock, self.calls = threading.Barrier(2, timeout=5), threading.Lock(), 0
 
         def complete(self, *, model: str, system: str, messages: Sequence[t.Message], max_tokens: int) -> t.Completion:
-            type(self).calls += 1
+            with self.lock:
+                self.calls += 1
             self.both.wait()
-            raise t.ConfigError("the API rejected the credentials: invalid bearer token")
+            if "# Home" in messages[0].content:
+                raise t.ConfigError("the API rejected the credentials: invalid bearer token")
+            return t.Completion(TOOLS_JA, t.Usage(10, 4, 0, 9))
 
-    code, out, err = translate(capsys, root, "--lang", "ja", jobs=2, translator=Rejected())
+    fake = RejectsHome()
 
-    assert (code, Rejected.calls, err) == (
-        2,
-        2,
-        "translations: the API rejected the credentials: invalid bearer token\n",
-    )
-    assert out == "translating 4 pages (ja), 2 at a time, with test-model\n"
-    assert not (root / "i18n" / "ja" / "pages").exists()
+    code, out, err = translate(capsys, root, "--lang", "ja", jobs=2, translator=fake)
+
+    assert (code, fake.calls, err) == (2, 2, "translations: the API rejected the credentials: invalid bearer token\n")
+    assert out == snapshot("""\
+translating 4 pages (ja), 2 at a time, with test-model
+ja: translated tools.md (3 of 3 sections)
+usage: 10 input / 4 output / 0 cache-write / 9 cache-read tokens
+""")
+    assert sorted(path.name for path in (root / "i18n" / "ja" / "pages").glob("*.md")) == ["tools.md"]
 
 
 def test_translate_without_lang_works_through_every_language_with_that_languages_own_prompt(