| 25 | |
| 26 | |
| 27 | def add_queue( |
| 28 | hq_env: HqEnv, |
| 29 | manager: ManagerType, |
| 30 | name: Optional[str] = None, |
| 31 | backlog=1, |
| 32 | max_workers_per_alloc=1, |
| 33 | additional_worker_args: List[str] = None, |
| 34 | additional_args=None, |
| 35 | additional_hq_args: List[str] = None, |
| 36 | time_limit="1h", |
| 37 | worker_time_limit: Optional[str] = None, |
| 38 | dry_run=False, |
| 39 | start_cmd: Optional[str] = None, |
| 40 | stop_cmd: Optional[str] = None, |
| 41 | wrap_cmd: Optional[str] = None, |
| 42 | **kwargs, |
| 43 | ) -> str: |
| 44 | args = [] |
| 45 | if additional_hq_args is not None: |
| 46 | args += additional_hq_args |
| 47 | |
| 48 | args += ["alloc", "add", manager] |
| 49 | if not dry_run: |
| 50 | args.append("--no-dry-run") |
| 51 | if name is not None: |
| 52 | args.extend(["--name", name]) |
| 53 | args.extend( |
| 54 | [ |
| 55 | "--backlog", |
| 56 | str(backlog), |
| 57 | "--max-workers-per-alloc", |
| 58 | str(max_workers_per_alloc), |
| 59 | ] |
| 60 | ) |
| 61 | if time_limit is not None: |
| 62 | args.extend(["--time-limit", time_limit]) |
| 63 | if worker_time_limit is not None: |
| 64 | args.extend(["--worker-time-limit", worker_time_limit]) |
| 65 | if start_cmd is not None: |
| 66 | args.extend(["--worker-start-cmd", start_cmd]) |
| 67 | if stop_cmd is not None: |
| 68 | args.extend(["--worker-stop-cmd", stop_cmd]) |
| 69 | if wrap_cmd is not None: |
| 70 | args.extend(["--worker-wrap-cmd", wrap_cmd]) |
| 71 | if additional_worker_args is not None: |
| 72 | args.extend(additional_worker_args) |
| 73 | if additional_args is not None: |
| 74 | args.append("--") |
| 75 | args.extend(additional_args.split(" ")) |
| 76 | |
| 77 | return hq_env.command(args, **kwargs) |
| 78 | |
| 79 | |
| 80 | def prepare_tasks(hq_env: HqEnv, count=1000, nodes=1): |