From f43b71359010e567c32c816f7d5f155c8fe93dbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Miquel?= Date: Tue, 6 Oct 2026 11:00:16 +0200 Subject: [PATCH] Configuration option for number of concurrent requests --- copienator/commands/correction.py | 20 ++++++-- copienator/commands/gemini_for_labels.py | 11 ++++- copienator_gui/app.py | 8 ++- copienator_gui/workflow.py | 22 +++++++++ tests/test_gui_convenience.py | 2 +- tests/test_gui_core.py | 62 +++++++++++++++++++++--- tests/test_gui_refaire.py | 11 +++++ tests/test_label_timeouts.py | 31 +++++++++++- 8 files changed, 151 insertions(+), 16 deletions(-) diff --git a/copienator/commands/correction.py b/copienator/commands/correction.py index 5dde024..f01965d 100644 --- a/copienator/commands/correction.py +++ b/copienator/commands/correction.py @@ -736,8 +736,17 @@ def resolve_delayed_moves(): return new_tasks +def request_workers(args: argparse.Namespace) -> int: + workers = getattr(args, "workers", NB_THREADS) + if workers < 1: + raise CliError("Le nombre de requêtes simultanées doit être positif.", + ExitCode.INVALID_ARGUMENTS) + return workers + + def run_configured(args: argparse.Namespace) -> ExitCode: global client, tasks_to_process + workers = request_workers(args) if client is None: client = genai.Client(api_key=api_key) if args.refaire: @@ -950,15 +959,15 @@ def run_configured(args: argparse.Namespace) -> ExitCode: made_progress = True while tasks_to_process or made_progress: if tasks_to_process: - print(f"Starting processing on {len(tasks_to_process)} tasks with {NB_THREADS} threads...") - with concurrent.futures.ThreadPoolExecutor(max_workers=NB_THREADS) as executor: + print(f"Starting processing on {len(tasks_to_process)} tasks with {workers} threads...") + with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor: waiting_tasks = list(tasks_to_process) futures = {} def submit_available_tasks() -> None: while ( waiting_tasks - and len(futures) < NB_THREADS + and len(futures) < workers and not stop_requested.is_set() ): task = waiting_tasks.pop(0) @@ -1074,6 +1083,7 @@ def run( *, api_client=None, ) -> ExitCode: + request_workers(args) if args.reset: workspace.require_directories("Copies") reset_workspace(workspace) @@ -1116,6 +1126,10 @@ def build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--overwrite", action="store_true", help="Redo requests") parser.add_argument("--limit", type=int, help="Maximum Gemini Pro calls") + parser.add_argument( + "--workers", type=int, default=NB_THREADS, + help="Maximum simultaneous Gemini correction requests (positive integer; default: 12)", + ) parser.add_argument( "--refaire", action="store_true", diff --git a/copienator/commands/gemini_for_labels.py b/copienator/commands/gemini_for_labels.py index 141d0b7..51d12f2 100644 --- a/copienator/commands/gemini_for_labels.py +++ b/copienator/commands/gemini_for_labels.py @@ -534,6 +534,9 @@ def run( sleep: Sleep = time.sleep, max_workers: int = 12, ) -> ExitCode: + if max_workers < 1: + raise CliError("Le nombre de requêtes simultanées doit être positif.", + ExitCode.INVALID_ARGUMENTS) workspace.require_files("labels") images, warnings = selected_images(workspace, targets) for warning in warnings: @@ -628,6 +631,12 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Regenerate JSON outputs that already exist", ) + parser.add_argument( + "--workers", + type=int, + default=12, + help="Maximum simultaneous Gemini requests (positive integer; default: 12)", + ) return parser @@ -655,7 +664,7 @@ def main(argv: Sequence[str] | None = None) -> int: ExitCode.INVALID_ARGUMENTS, ) targets.append(resolved) - return run(workspace, targets, overwrite=args.overwrite) + return run(workspace, targets, overwrite=args.overwrite, max_workers=args.workers) return execute(parser, argv, handle) diff --git a/copienator_gui/app.py b/copienator_gui/app.py index 8c26bf8..4aea8a0 100644 --- a/copienator_gui/app.py +++ b/copienator_gui/app.py @@ -874,7 +874,7 @@ class CopienatorApp(tk.Tk): variant = self._current_variant() evaluation_arg = self._evaluation_arg() for spec in step.arguments: - if step.section == REFAIRE_SECTION: + if step.section == REFAIRE_SECTION and spec.name != "workers": continue if spec.variants and variant.id not in spec.variants: continue @@ -1318,7 +1318,11 @@ class CopienatorApp(tk.Tk): if self.current_step and self.current_step.section == REFAIRE_SECTION: if self.current_step.id == "refaire_selection" and self.refaire_panel: return self.refaire_panel.values() - return {"target": self._evaluation_arg(), "annotation_dir": self.state_store.step("refaire_selection").get("values", {}).get("annotation_dir", "BGnot")} + return { + "target": self._evaluation_arg(), + "annotation_dir": self.state_store.step("refaire_selection").get("values", {}).get("annotation_dir", "BGnot"), + **{name: variable.get() for name, variable in self.arg_vars.items()}, + } return {name: variable.get() for name, variable in self.arg_vars.items()} def _save_current_form(self) -> None: diff --git a/copienator_gui/workflow.py b/copienator_gui/workflow.py index bd7296c..4861265 100644 --- a/copienator_gui/workflow.py +++ b/copienator_gui/workflow.py @@ -74,6 +74,15 @@ def arg_target(help_text: str = "un dossier d’évaluation ou un fichier à tra ) +def arg_correction_workers(variants: tuple[str, ...] = ()) -> ArgumentSpec: + return ArgumentSpec( + "workers", "Requêtes simultanées", "int", "--workers", default=12, + help="Fixe le nombre maximal de requêtes Gemini de correction simultanées (au moins 1). " + "Réduisez cette valeur pour limiter la charge sur l’API.", + variants=variants, + ) + + def build_workflow(show_personal_steps: bool) -> list[StepDefinition]: def python(ident, label, script, **kwargs): supports_verbose = kwargs.pop("supports_verbose", True) @@ -317,6 +326,16 @@ def build_workflow(show_personal_steps: bool) -> list[StepDefinition]: (python("default", "Détection des labels", "labels"),), arguments=( arg_target(), + ArgumentSpec( + "workers", + "Requêtes simultanées", + "int", + "--workers", + default=12, + help="Fixe le nombre maximal de requêtes Gemini simultanées (au moins 1). " + "Réduisez cette valeur pour limiter la charge sur l’API. " + "Les pages d’une même copie sont traitées dans l’ordre.", + ), ArgumentSpec( "overwrite", "Régénérer les résultats", @@ -411,6 +430,7 @@ def build_workflow(show_personal_steps: bool) -> list[StepDefinition]: ), arguments=( arg_target("le dossier de l’évaluation ou une image Group_X.jpg"), + arg_correction_workers(("live", "hybrid", "refaire", "integrate")), ArgumentSpec( "overwrite", "Écraser les corrections existantes", @@ -796,6 +816,8 @@ def build_refaire_workflow() -> list[StepDefinition]: steps = [selection] for suffix, title, description, program, flags, optional in definitions: arguments = (arg_target(),) if program else () + if suffix == "correct": + arguments += (arg_correction_workers(),) if suffix == "merge": arguments += (ArgumentSpec( "annotation_dir", diff --git a/tests/test_gui_convenience.py b/tests/test_gui_convenience.py index 1cb98d9..48378e5 100644 --- a/tests/test_gui_convenience.py +++ b/tests/test_gui_convenience.py @@ -505,7 +505,7 @@ class GuiConvenienceTests(unittest.TestCase): self.app._target_selected_copy() command = self.app._make_command() - self.assertEqual(Path(command[-1]), copies / "Copie02.pdf") + self.assertEqual(Path(command[5]), copies / "Copie02.pdf") def button_texts(widget): return [ diff --git a/tests/test_gui_core.py b/tests/test_gui_core.py index 6617ccf..1f2d785 100644 --- a/tests/test_gui_core.py +++ b/tests/test_gui_core.py @@ -1479,7 +1479,7 @@ class StandardCliTests(unittest.TestCase): with tempfile.TemporaryDirectory() as directory: workspace = EvaluationWorkspace(Path(directory)) args = module.build_parser().parse_args( - [str(workspace.root), "--overwrite"] + [str(workspace.root), "--overwrite", "--workers", "1"] ) module.configure_runtime( workspace, @@ -1495,18 +1495,46 @@ class StandardCliTests(unittest.TestCase): (workspace.root / "saved-result").write_text(task[0]) return [] - with patch.object(module, "NB_THREADS", 1), patch.object( + with patch.object( + module.concurrent.futures, "ThreadPoolExecutor", wraps=ThreadPoolExecutor + ) as executor, patch.object( module, "process_single_task", side_effect=process ), patch.object( module, "resolve_delayed_moves", return_value=[] ): status = module.run_configured(args) + executor.assert_called_once_with(max_workers=1) saved_result = (workspace.root / "saved-result").read_text() self.assertEqual(status, module.ExitCode.INTERRUPTED) self.assertEqual(processed, ["group-1.jpg"]) self.assertEqual(saved_result, "group-1.jpg") + def test_correction_concurrency_defaults_and_cli_forwarding(self) -> None: + module = self.modules["correction"] + self.assertEqual(module.build_parser().parse_args(["Exam"]).workers, 12) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + (root / "Copies").mkdir() + with patch.object(module, "run", return_value=ExitCode.SUCCESS) as run: + self.assertEqual(module.main([str(root), "--workers", "2"]), 0) + self.assertEqual(run.call_args.args[2].workers, 2) + + def test_invalid_correction_concurrency_rejected_before_changes(self) -> None: + module = self.modules["correction"] + with tempfile.TemporaryDirectory() as directory: + workspace = EvaluationWorkspace(Path(directory)) + for workers in (0, -1): + args = module.build_parser().parse_args( + [str(workspace.root), "--workers", str(workers)] + ) + with self.subTest(workers=workers), \ + patch.object(module, "configure_runtime") as configure, \ + self.assertRaises(CliError) as caught: + module.run(workspace, [], args) + self.assertEqual(caught.exception.exit_code, ExitCode.INVALID_ARGUMENTS) + configure.assert_not_called() + def test_gemini_call_finishes_inflight_stream_but_does_not_retry(self) -> None: module = self.modules["correction"] self.addCleanup(module.stop_requested.clear) @@ -1562,7 +1590,7 @@ class StandardCliTests(unittest.TestCase): api_client.models.generate_content_stream.side_effect = stream args = module.build_parser().parse_args([str(workspace.root)]) - with patch.object(module, "NB_THREADS", 1), patch.object( + with patch.object(args, "workers", 1), patch.object( module.prompting, "generate_request", return_value=([], Mock()) ): module.configure_runtime( @@ -1636,7 +1664,7 @@ class StandardCliTests(unittest.TestCase): [str(workspace.root), "--overwrite"] ) - with patch.object(module, "NB_THREADS", 1), patch.object( + with patch.object(args, "workers", 1), patch.object( module.prompting, "generate_request", return_value=([], Mock()) ), patch.object( module, "handle_label_errors", side_effect=resolve_error @@ -2223,6 +2251,16 @@ class WorkflowTests(unittest.TestCase): variant = next(item for item in step.variants if item.id == variant_id) return build_command(REPOSITORY, step, variant, values, self.evaluation) + def test_label_request_concurrency_is_configurable(self) -> None: + for personal in (False, True): + step = next(step for step in build_workflow(personal) if step.id == "labels") + for values, expected in (({}, "12"), ({"workers": "2"}, "2")): + with self.subTest(personal=personal, values=values): + command = build_command( + REPOSITORY, step, step.variants[0], values, self.evaluation + ) + self.assertEqual(command[command.index("--workers") + 1], expected) + def test_personal_steps_are_configurable(self) -> None: standard_ids = {step.id for step in build_workflow(False)} personal_ids = {step.id for step in build_workflow(True)} @@ -2496,10 +2534,19 @@ class WorkflowTests(unittest.TestCase): ) self.assertEqual( command_arguments(command), - [self.evaluation, "--overwrite", "--limit", "0"], + [self.evaluation, "--workers", "12", "--overwrite", "--limit", "0"], ) self.assertEqual(command[4], "correct") + def test_correction_concurrency_applies_to_live_requests(self) -> None: + for variant in self.steps["correction"].variants: + with self.subTest(variant=variant.id): + command = self.command("correction", variant.id, {"workers": "2"}) + if variant.id in {"batch", "reset"}: + self.assertNotIn("--workers", command) + else: + self.assertEqual(command[command.index("--workers") + 1], "2") + def test_refaire_correction_does_not_offer_overwrite(self) -> None: step = self.steps["correction"] overwrite = next( @@ -2520,7 +2567,7 @@ class WorkflowTests(unittest.TestCase): "correction", "hybrid", {"target": self.evaluation, "batch_from": "Ex 4"} ) self.assertEqual( - command_arguments(command), [self.evaluation, "--batch-from", "Ex 4"] + command_arguments(command), [self.evaluation, "--workers", "12", "--batch-from", "Ex 4"] ) def test_only_multi_target_steps_offer_free_form_arguments(self) -> None: @@ -2624,7 +2671,8 @@ class WorkflowTests(unittest.TestCase): self.evaluation, verbose=True, ) - self.assertEqual(command_arguments(labels_command), [self.evaluation, "--verbose"]) + self.assertEqual(command_arguments(labels_command), + [self.evaluation, "--workers", "12", "--verbose"]) rotate = self.steps["rotate"] rotate_command = build_command( diff --git a/tests/test_gui_refaire.py b/tests/test_gui_refaire.py index d1785ba..48be102 100644 --- a/tests/test_gui_refaire.py +++ b/tests/test_gui_refaire.py @@ -126,6 +126,17 @@ class RefaireGuiTests(unittest.TestCase): self.app._populate_tree() self.assertTrue(self.app.tree.item(section, "open")) + def test_correction_concurrency_can_be_edited_and_saved(self): + self.save_selection() + self.select("refaire_correct") + self.assertEqual(self.app.arg_vars["workers"].get(), "12") + self.app.arg_vars["workers"].set("2") + command = self.app._make_command() + self.assertEqual(command[command.index("--workers") + 1], "2") + self.select("refaire_review") + self.select("refaire_correct") + self.assertEqual(self.app.arg_vars["workers"].get(), "2") + def test_dropdown_selection_saves_json_and_drives_all_commands(self): self.save_selection() self.assertEqual( diff --git a/tests/test_label_timeouts.py b/tests/test_label_timeouts.py index 9295659..ec0d6b3 100644 --- a/tests/test_label_timeouts.py +++ b/tests/test_label_timeouts.py @@ -1,13 +1,13 @@ import io import tempfile import unittest -from concurrent.futures import wait as real_wait +from concurrent.futures import ThreadPoolExecutor, wait as real_wait from contextlib import redirect_stdout from pathlib import Path from threading import Event from unittest.mock import Mock, patch -from copienator import EvaluationWorkspace, ExitCode +from copienator import CliError, EvaluationWorkspace, ExitCode from copienator.commands import gemini_for_labels as labels @@ -36,6 +36,33 @@ class LabelTimeoutTests(unittest.TestCase): sleep=lambda _: None, target_interval=0, ) + def test_cli_forwards_selected_concurrency(self): + for arguments, expected in (([], 12), (["--workers", "2"], 2)): + with self.subTest(arguments=arguments), \ + patch.object(labels, "run", return_value=ExitCode.SUCCESS) as run: + self.assertEqual(labels.main([str(self.workspace.root), *arguments]), 0) + run.assert_called_once_with( + unittest.mock.ANY, [self.workspace.root], + overwrite=False, max_workers=expected, + ) + + def test_selected_concurrency_controls_executor(self): + self.client.models.generate_content.return_value = self.response + with patch.object(labels, "selected_images", return_value=(self.images, [])), \ + patch.object(labels, "ThreadPoolExecutor", wraps=ThreadPoolExecutor) as executor: + result = labels.run(self.workspace, [], client=self.client, + sleep=lambda _: None, max_workers=1) + self.assertEqual(result, ExitCode.SUCCESS) + executor.assert_called_once_with(max_workers=1) + self.assertEqual(self.client.models.generate_content.call_count, 3) + + def test_nonpositive_concurrency_rejected_before_requests(self): + for workers in (0, -1): + with self.subTest(workers=workers), self.assertRaises(CliError) as caught: + labels.run(self.workspace, [], client=self.client, max_workers=workers) + self.assertEqual(caught.exception.exit_code, ExitCode.INVALID_ARGUMENTS) + self.client.models.generate_content.assert_not_called() + def test_timeout_recovery_uses_bounded_requests_and_saves_result(self): self.client.models.generate_content.side_effect = [TimeoutError(), self.response] self.assertEqual(self.group(self.images[:1]), 1)