Configuration option for number of concurrent requests
This commit is contained in:
1 parent
94ca8b8256
commit
f43b713590
8 files changed
+151
-16
No files matched your search
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user