Skills › Scientific Agent Skills › get-available-resources › plan_workload.py
plan_workload.py
318 lines · skills/get-available-resources/scripts/plan_workload.py at 5f2bc03
1#!/usr/bin/env python3
2"""Build a conservative worker and memory plan from a validated snapshot."""
4from __future__ import annotations
6import argparse
7import math
8from typing import Any, Mapping, Sequence
10from _common import (
11 MAX_BYTES,
12 MAX_TASKS,
13 MAX_WORKERS,
14 SCHEMA_VERSION,
15 ResourceToolError,
16 argparse_nonnegative_int,
17 argparse_positive_int,
18 bounded_number,
19 cli_error,
20 emit_json,
21 read_json_file,
22)
23from snapshot_tools import validate_snapshot
25MIB = 1024**2
28def _optional_number(value: Any) -> float | None:
29 return bounded_number(value, maximum=float(MAX_BYTES))
32def _accelerator_decision(
33 snapshot: Mapping[str, Any],
34 requested: str,
35) -> dict[str, Any]:
36 if requested == "none":
37 return {
38 "candidate_count_upper_bound": 0,
39 "requested": "none",
40 "status": "not_requested",
41 }
42 accelerators = snapshot["accelerators"]
43 bounds = accelerators.get("candidate_upper_bounds", {})
44 counts = accelerators.get("candidate_counts", {})
45 if requested == "any":
46 candidates = [
47 value
48 for backend, value in bounds.items()
49 if backend in {"cuda", "metal", "rocm"} and isinstance(value, int)
50 ]
51 query_count = sum(
52 value
53 for backend, value in counts.items()
54 if backend in {"cuda", "metal", "rocm"} and isinstance(value, int)
55 )
56 upper_bound = sum(candidates) if candidates else None
57 else:
58 query_count = counts.get(requested, 0)
59 upper_bound = bounds.get(requested)
60 if not query_count:
61 status = "no_management_visible_candidate"
62 elif upper_bound == 0:
63 status = "restricted_by_allocation_or_visibility"
64 else:
65 status = "diagnostic_required"
66 return {
67 "candidate_count_upper_bound": upper_bound,
68 "management_query_count": query_count,
69 "requested": requested,
70 "status": status,
71 "usability": "not_verified",
72 }
75def build_plan(
76 snapshot: Mapping[str, Any],
77 *,
78 workload: str = "cpu",
79 task_count: int | None = None,
80 requested_workers: int | None = None,
81 memory_per_worker_mib: int | None = None,
82 reserve_memory_mib: int | None = None,
83 accelerator: str = "none",
84) -> dict[str, Any]:
85 """Return a bounded deterministic plan without executing work."""
86 validation = validate_snapshot(snapshot)
87 if not validation["valid"]:
88 raise ResourceToolError("snapshot does not satisfy schema 1.1")
89 if workload not in {"cpu", "io", "mixed"}:
90 raise ResourceToolError("workload must be cpu, io, or mixed")
91 if accelerator not in {"none", "any", "cuda", "metal", "rocm"}:
92 raise ResourceToolError(
93 "accelerator must be none, any, cuda, metal, or rocm"
94 )
95 for label, value, maximum in (
96 ("task_count", task_count, MAX_TASKS),
97 ("requested_workers", requested_workers, MAX_WORKERS),
98 ("memory_per_worker_mib", memory_per_worker_mib, 1_048_576),
99 ):
100 if value is not None and (
101 isinstance(value, bool)
102 or not isinstance(value, int)
103 or not 1 <= value <= maximum
104 ):
105 raise ResourceToolError(f"{label} is outside its safe range")
106 if reserve_memory_mib is not None and (
107 isinstance(reserve_memory_mib, bool)
108 or not isinstance(reserve_memory_mib, int)
109 or not 0 <= reserve_memory_mib <= 1_048_576
110 ):
111 raise ResourceToolError("reserve_memory_mib is outside its safe range")
113 effective_cpu = snapshot["cpu"]["effective"]
114 cpu_capacity = _optional_number(effective_cpu.get("capacity_cores"))
115 cpu_ceiling = effective_cpu.get("worker_ceiling", 1)
116 if not isinstance(cpu_ceiling, int):
117 cpu_ceiling = 1
118 cpu_ceiling = max(1, min(MAX_WORKERS, cpu_ceiling))
119 if cpu_capacity is not None:
120 cpu_ceiling = min(cpu_ceiling, max(1, math.floor(cpu_capacity)))
121 if workload == "io":
122 workload_ceiling = min(MAX_WORKERS, 32, max(2, cpu_ceiling * 2))
123 else:
124 workload_ceiling = cpu_ceiling
126 available_memory = _optional_number(
127 snapshot["memory"]["effective"].get("available_bytes")
128 )
129 if reserve_memory_mib is None and available_memory is not None:
130 reserve_bytes = min(
131 int(available_memory * 0.5),
132 max(256 * MIB, int(available_memory * 0.1)),
133 )
134 elif reserve_memory_mib is not None:
135 reserve_bytes = reserve_memory_mib * MIB
136 else:
137 reserve_bytes = None
139 memory_ceiling: int | None = None
140 planner_warnings: list[dict[str, str]] = []
141 if memory_per_worker_mib is not None:
142 per_worker_bytes = memory_per_worker_mib * MIB
143 if available_memory is None or reserve_bytes is None:
144 memory_ceiling = 1
145 planner_warnings.append(
146 {
147 "code": "MEMORY_UNKNOWN",
148 "message": (
149 "Memory per worker was supplied but effective available "
150 "memory is unknown; one worker is recommended."
151 ),
152 }
153 )
154 else:
155 usable = max(0, int(available_memory) - reserve_bytes)
156 memory_ceiling = min(MAX_WORKERS, usable // per_worker_bytes)
157 if usable < per_worker_bytes:
158 planner_warnings.append(
159 {
160 "code": "MEMORY_REQUIREMENT_EXCEEDS_BUDGET",
161 "message": (
162 "No worker fits the post-reserve memory budget; "
163 "reduce the requirement or use out-of-core processing."
164 ),
165 }
166 )
168 ceilings: list[tuple[str, int]] = [("workload_cpu", workload_ceiling)]
169 if memory_ceiling is not None:
170 ceilings.append(("memory", memory_ceiling))
171 if task_count is not None:
172 ceilings.append(("tasks", min(task_count, MAX_WORKERS)))
173 if requested_workers is not None:
174 ceilings.append(("user_request", requested_workers))
175 suggested = min(value for _, value in ceilings)
176 binding = sorted(source for source, value in ceilings if value == suggested)
178 if workload == "io":
179 planner_warnings.append(
180 {
181 "code": "IO_OVERSUBSCRIPTION_HEURISTIC",
182 "message": (
183 "The I/O worker count is a bounded heuristic; benchmark the "
184 "real workload without exceeding allocations."
185 ),
186 }
187 )
188 if snapshot["scheduler"].get("detected"):
189 planner_warnings.append(
190 {
191 "code": "SCHEDULER_SCOPE",
192 "message": (
193 "Remain inside the scheduler allocation even when host "
194 "inventory is larger."
195 ),
196 }
197 )
198 if cpu_capacity is None:
199 planner_warnings.append(
200 {
201 "code": "CPU_CAPACITY_UNKNOWN",
202 "message": "Effective CPU capacity is unknown; the plan remains conservative.",
203 }
204 )
206 threads_per_worker = 1 if suggested else 0
207 if suggested and workload in {"cpu", "mixed"} and cpu_capacity is not None:
208 threads_per_worker = max(1, math.floor(cpu_capacity / suggested))
209 planner_warnings.sort(key=lambda item: (item["code"], item["message"]))
210 return {
211 "accelerator": _accelerator_decision(snapshot, accelerator),
212 "inputs": {
213 "accelerator": accelerator,
214 "memory_per_worker_bytes": (
215 memory_per_worker_mib * MIB
216 if memory_per_worker_mib is not None
217 else None
218 ),
219 "requested_workers": requested_workers,
220 "reserve_memory_bytes": reserve_bytes,
221 "task_count": task_count,
222 "workload": workload,
223 },
224 "kind": "workload_plan",
225 "limits": {
226 "cpu_capacity_cores": cpu_capacity,
227 "cpu_worker_ceiling": cpu_ceiling,
228 "memory_worker_ceiling": memory_ceiling,
229 "workload_worker_ceiling": workload_ceiling,
230 },
231 "recommendation": {
232 "binding_limits": binding,
233 "status": (
234 "insufficient_memory" if not suggested else
235 "review_required" if cpu_capacity is None or available_memory is None else
236 "provisional"
237 ),
238 "suggested_workers": suggested,
239 "threads_per_worker": threads_per_worker,
240 },
241 "schema_version": SCHEMA_VERSION,
242 "warnings": planner_warnings,
243 }
246def build_parser() -> argparse.ArgumentParser:
247 parser = argparse.ArgumentParser(
248 description="Plan bounded workers and memory from a resource snapshot"
249 )
250 parser.add_argument("snapshot", help="validated snapshot JSON file")
251 parser.add_argument(
252 "--workload",
253 choices=("cpu", "io", "mixed"),
254 default="cpu",
255 help="workload shape; default: cpu",
256 )
257 parser.add_argument(
258 "--tasks",
259 type=argparse_positive_int,
260 help="number of independent tasks",
261 )
262 parser.add_argument(
263 "--workers",
264 type=lambda value: argparse_positive_int(value, maximum=MAX_WORKERS),
265 help="explicit upper bound on workers",
266 )
267 parser.add_argument(
268 "--memory-per-worker-mib",
269 type=lambda value: argparse_positive_int(value, maximum=1_048_576),
270 help="estimated MiB required by each worker",
271 )
272 parser.add_argument(
273 "--reserve-memory-mib",
274 type=lambda value: argparse_nonnegative_int(value, maximum=1_048_576),
275 help="explicit MiB to keep outside the worker budget",
276 )
277 parser.add_argument(
278 "--accelerator",
279 choices=("none", "any", "cuda", "rocm", "metal"),
280 default="none",
281 help="request a candidate backend without assuming runtime usability",
282 )
283 parser.add_argument(
284 "--output",
285 metavar="FILE.json",
286 help="write a private local JSON file instead of stdout",
287 )
288 parser.add_argument(
289 "--force",
290 action="store_true",
291 help="replace an existing explicit output file",
292 )
293 return parser
296def main(argv: Sequence[str] | None = None) -> int:
297 parser = build_parser()
298 args = parser.parse_args(argv)
299 if args.force and not args.output:
300 parser.error("--force requires --output")
301 try:
302 plan = build_plan(
303 read_json_file(args.snapshot),
304 workload=args.workload,
305 task_count=args.tasks,
306 requested_workers=args.workers,
307 memory_per_worker_mib=args.memory_per_worker_mib,
308 reserve_memory_mib=args.reserve_memory_mib,
309 accelerator=args.accelerator,
310 )
311 emit_json(plan, args.output, force=args.force)
312 except ResourceToolError as exc:
313 return cli_error(parser, exc)
314 return 0
317if __name__ == "__main__":
318 raise SystemExit(main())