Skills › Scientific Agent Skills › pathml › plan_inference.py
plan_inference.py
282 lines · skills/pathml/scripts/plan_inference.py at 5f2bc03
1#!/usr/bin/env python3
2"""Plan bounded inference batches from numbers or a JSON model card only."""
4from __future__ import annotations
6import argparse
7import math
8import re
9from pathlib import Path
10from typing import Any
12from _common import (
13 CliError,
14 MAX_JSON_BYTES,
15 emit_json,
16 finite_float,
17 load_json_object,
18 run_cli,
19 validate_keys,
20)
23DTYPE_BYTES = {
24 "uint8": 1,
25 "int8": 1,
26 "float16": 2,
27 "bfloat16": 2,
28 "float32": 4,
29 "float64": 8,
30}
31MODEL_SUFFIXES = {
32 ".bin",
33 ".ckpt",
34 ".joblib",
35 ".onnx",
36 ".pickle",
37 ".pkl",
38 ".pt",
39 ".pth",
40 ".safetensors",
41}
42SHA256_PATTERN = re.compile(r"^[0-9a-fA-F]{64}$")
45def _load_card(path: str | None, root: str) -> dict[str, Any]:
46 if path is None:
47 return {}
48 if Path(path).suffix.lower() in MODEL_SUFFIXES:
49 raise CliError(
50 "--model-card accepts JSON metadata only; model/checkpoint files "
51 "are intentionally never loaded"
52 )
53 card = load_json_object(path, root=root, max_bytes=MAX_JSON_BYTES)
54 validate_keys(
55 card,
56 allowed={
57 "schema_version",
58 "model_id",
59 "artifact_sha256",
60 "input_shape",
61 "dtype",
62 "output_elements_per_tile",
63 "activation_multiplier",
64 "parameter_bytes",
65 "runtime_workspace_mib",
66 },
67 required={"schema_version", "model_id", "input_shape", "dtype"},
68 context="model card",
69 )
70 if card["schema_version"] != "1.0":
71 raise CliError("model card schema_version must be '1.0'")
72 if not isinstance(card["model_id"], str) or not 1 <= len(card["model_id"]) <= 128:
73 raise CliError("model_id must be a nonempty string of at most 128 characters")
74 digest = card.get("artifact_sha256")
75 if digest is not None and (
76 not isinstance(digest, str) or not SHA256_PATTERN.fullmatch(digest)
77 ):
78 raise CliError("artifact_sha256 must be exactly 64 hexadecimal characters")
79 shape = card["input_shape"]
80 if (
81 not isinstance(shape, list)
82 or len(shape) != 3
83 or any(isinstance(value, bool) or not isinstance(value, int) for value in shape)
84 ):
85 raise CliError("input_shape must be [channels, height, width] integers")
86 return card
89def _select(cli_value: Any, card: dict[str, Any], card_key: str, default: Any) -> Any:
90 return cli_value if cli_value is not None else card.get(card_key, default)
93def make_plan(args: argparse.Namespace) -> dict[str, Any]:
94 card = _load_card(args.model_card, args.root)
95 shape = card.get("input_shape", [3, 256, 256])
96 channels = _select(args.channels, card, "channels", shape[0])
97 height = _select(args.height, card, "height", shape[1])
98 width = _select(args.width, card, "width", shape[2])
99 dtype = _select(args.dtype, card, "dtype", "float32")
100 output_elements = _select(
101 args.output_elements_per_tile,
102 card,
103 "output_elements_per_tile",
104 0,
105 )
106 activation_multiplier = _select(
107 args.activation_multiplier,
108 card,
109 "activation_multiplier",
110 8.0,
111 )
112 parameter_bytes = _select(
113 args.parameter_bytes,
114 card,
115 "parameter_bytes",
116 0,
117 )
118 runtime_workspace_mib = _select(
119 args.runtime_workspace_mib,
120 card,
121 "runtime_workspace_mib",
122 512.0,
123 )
125 integers = {
126 "tile_count": (args.tile_count, 1, 100_000_000),
127 "batch_size": (args.batch_size, 1, 1_000_000),
128 "channels": (channels, 1, 4096),
129 "height": (height, 1, 16384),
130 "width": (width, 1, 16384),
131 "output_elements_per_tile": (output_elements, 0, 10_000_000_000),
132 "parameter_bytes": (parameter_bytes, 0, 10_000_000_000_000),
133 }
134 checked: dict[str, int] = {}
135 for name, (value, minimum, maximum) in integers.items():
136 if isinstance(value, bool) or not isinstance(value, int):
137 raise CliError(f"{name} must be an integer")
138 if not minimum <= value <= maximum:
139 raise CliError(f"{name} must be between {minimum} and {maximum}")
140 checked[name] = value
141 if not isinstance(dtype, str) or dtype not in DTYPE_BYTES:
142 raise CliError(
143 "--dtype must be one of: " + ", ".join(sorted(DTYPE_BYTES))
144 )
145 activation_multiplier = finite_float(
146 activation_multiplier,
147 name="activation_multiplier",
148 minimum=0,
149 maximum=10_000,
150 )
151 runtime_workspace_mib = finite_float(
152 runtime_workspace_mib,
153 name="runtime_workspace_mib",
154 minimum=0,
155 maximum=1_000_000,
156 )
157 max_memory_mib = finite_float(
158 args.max_memory_mib,
159 name="max_memory_mib",
160 minimum=1,
161 maximum=10_000_000,
162 )
164 bytes_per_element = DTYPE_BYTES[dtype]
165 input_elements = checked["channels"] * checked["height"] * checked["width"]
166 input_bytes_per_tile = input_elements * bytes_per_element
167 output_bytes_per_tile = (
168 checked["output_elements_per_tile"] * bytes_per_element
169 )
170 activation_bytes_per_tile = math.ceil(
171 input_bytes_per_tile * activation_multiplier
172 )
173 variable_bytes_per_tile = (
174 input_bytes_per_tile + output_bytes_per_tile + activation_bytes_per_tile
175 )
176 persistent_bytes = checked["parameter_bytes"] + round(
177 runtime_workspace_mib * 1024**2
178 )
179 memory_budget_bytes = round(max_memory_mib * 1024**2)
180 available_variable_bytes = max(0, memory_budget_bytes - persistent_bytes)
181 recommended_max_batch = (
182 available_variable_bytes // variable_bytes_per_tile
183 if variable_bytes_per_tile
184 else checked["tile_count"]
185 )
186 recommended_max_batch = min(recommended_max_batch, checked["tile_count"])
187 planned_peak_bytes = persistent_bytes + (
188 variable_bytes_per_tile * checked["batch_size"]
189 )
190 within_memory = planned_peak_bytes <= memory_budget_bytes
191 number_of_batches = math.ceil(
192 checked["tile_count"] / checked["batch_size"]
193 )
194 final_batch_size = checked["tile_count"] % checked["batch_size"]
195 if final_batch_size == 0:
196 final_batch_size = min(checked["batch_size"], checked["tile_count"])
198 warnings = [
199 "Estimate excludes allocator fragmentation, framework caches, graph "
200 "workspace, postprocessing, stitching, and concurrent workers."
201 ]
202 if not within_memory:
203 warnings.append(
204 "Requested batch exceeds the planning memory budget; reduce batch "
205 "size or change the explicit budget."
206 )
207 if checked["output_elements_per_tile"] == 0:
208 warnings.append(
209 "No output tensor size was supplied; output memory is omitted."
210 )
211 if checked["parameter_bytes"] == 0:
212 warnings.append("No parameter size was supplied; model weights are omitted.")
214 return {
215 "model_card_used": bool(card),
216 "model_id": card.get("model_id"),
217 "artifact_sha256": card.get("artifact_sha256"),
218 "checkpoint_or_model_loaded": False,
219 "input_shape_chw": [
220 checked["channels"],
221 checked["height"],
222 checked["width"],
223 ],
224 "dtype": dtype,
225 "bytes_per_element": bytes_per_element,
226 "tile_count": checked["tile_count"],
227 "requested_batch_size": checked["batch_size"],
228 "number_of_batches": number_of_batches,
229 "final_batch_size": final_batch_size,
230 "input_bytes_per_tile": input_bytes_per_tile,
231 "output_bytes_per_tile": output_bytes_per_tile,
232 "activation_bytes_per_tile_estimate": activation_bytes_per_tile,
233 "persistent_bytes_estimate": persistent_bytes,
234 "planned_peak_memory_mib": round(planned_peak_bytes / 1024**2, 6),
235 "memory_budget_mib": max_memory_mib,
236 "within_memory_budget": within_memory,
237 "recommended_max_batch_under_estimate": int(recommended_max_batch),
238 "estimated_total_output_gib": round(
239 checked["tile_count"] * output_bytes_per_tile / 1024**3,
240 6,
241 ),
242 "warnings": warnings,
243 "note": (
244 "This planner reads numbers or strict JSON metadata only. It never "
245 "opens model/checkpoint artifacts or imports PathML/Torch/ONNX."
246 ),
247 }
250def build_parser() -> argparse.ArgumentParser:
251 parser = argparse.ArgumentParser(
252 description=(
253 "Plan bounded inference batches from dimensions or a strict JSON "
254 "model card. Model/checkpoint files are never loaded."
255 )
256 )
257 parser.add_argument("--tile-count", type=int, required=True)
258 parser.add_argument("--batch-size", type=int, required=True)
259 parser.add_argument("--model-card")
260 parser.add_argument("--root", default=".")
261 parser.add_argument("--channels", type=int)
262 parser.add_argument("--height", type=int)
263 parser.add_argument("--width", type=int)
264 parser.add_argument("--dtype", choices=sorted(DTYPE_BYTES))
265 parser.add_argument("--output-elements-per-tile", type=int)
266 parser.add_argument("--activation-multiplier", type=float)
267 parser.add_argument("--parameter-bytes", type=int)
268 parser.add_argument("--runtime-workspace-mib", type=float)
269 parser.add_argument("--max-memory-mib", type=float, default=4096.0)
270 parser.add_argument("--output")
271 parser.add_argument("--force", action="store_true")
272 return parser
275def main() -> None:
276 args = build_parser().parse_args()
277 report = make_plan(args)
278 emit_json(report, output=args.output, root=args.root, force=args.force)
281if __name__ == "__main__":
282 raise SystemExit(run_cli(main))