Skills › Scientific Agent Skills › fluidsim › grid_resource_estimator.py
grid_resource_estimator.py
300 lines · skills/fluidsim/scripts/grid_resource_estimator.py at 5f2bc03
1#!/usr/bin/env python3
2"""Estimate bounded FluidSim grid memory and output storage envelopes."""
4from __future__ import annotations
6import argparse
7import math
8from typing import Any
10try:
11 from ._common import (
12 GIB,
13 MAX_RECORDS,
14 ToolError,
15 bounded_int,
16 checked_input,
17 emit_json,
18 fail_json,
19 finite_float,
20 load_json,
21 )
22 from ._schema import (
23 CONFIG_SOLVER_DIMENSIONS,
24 STATE_FIELD_COUNTS,
25 normalized_copy,
26 )
27except ImportError: # Direct script execution.
28 from _common import (
29 GIB,
30 MAX_RECORDS,
31 ToolError,
32 bounded_int,
33 checked_input,
34 emit_json,
35 fail_json,
36 finite_float,
37 load_json,
38 )
39 from _schema import (
40 CONFIG_SOLVER_DIMENSIONS,
41 STATE_FIELD_COUNTS,
42 normalized_copy,
43 )
46TOOL = "fluidsim-grid-resource-estimator"
49def _records(t_end: float, period: float) -> int:
50 if period <= 0:
51 return 0
52 ratio = t_end / period
53 if not math.isfinite(ratio) or ratio > MAX_RECORDS - 2:
54 raise ToolError("predicted output records exceed the planning limit; reduce duration or save frequency")
55 return int(math.floor(ratio + 1e-12)) + 2
58def estimate(
59 config: dict[str, Any],
60 *,
61 precision_bytes: int,
62 workspace_factor: float,
63 safety_factor: float,
64 compression_ratio: float,
65) -> dict[str, Any]:
66 solver = config["solver"]
67 dimension = CONFIG_SOLVER_DIMENSIONS[solver]
68 parameters = config["parameters"]
69 oper = parameters["oper"]
70 axes = ("x", "y", "z")[:dimension]
71 shape = [int(oper[f"n{axis}"]) for axis in axes]
72 grid_points = math.prod(shape)
73 # Sequential FluidFFT stores physical (y, x) / (z, y, x) and reduces x.
74 spectral_shape = [*reversed(shape[1:]), shape[0] // 2 + 1]
75 spectral_points = math.prod(spectral_shape)
77 real_fields, complex_fields = STATE_FIELD_COUNTS[solver]
78 real_state_bytes = grid_points * real_fields * precision_bytes
79 complex_state_bytes = spectral_points * complex_fields * 2 * precision_bytes
80 resident_state_bytes = real_state_bytes + complex_state_bytes
81 peak_bytes = math.ceil(
82 resident_state_bytes * workspace_factor * safety_factor
83 )
85 output = parameters.get("output", {})
86 periods = output.get("periods_save", {})
87 time = parameters["time_stepping"]
88 t_end = float(time["t_end"])
89 state_records = (
90 _records(t_end, float(periods.get("phys_fields", 0.0)))
91 if output.get("HAS_TO_SAVE")
92 else 0
93 )
94 spectra_records = (
95 _records(t_end, float(periods.get("spectra", 0.0)))
96 if output.get("HAS_TO_SAVE")
97 else 0
98 )
99 budget_records = (
100 _records(
101 t_end,
102 float(periods.get("spect_energy_budg", 0.0)),
103 )
104 if output.get("HAS_TO_SAVE")
105 else 0
106 )
107 scalar_records = (
108 _records(t_end, float(periods.get("spatial_means", 0.0)))
109 if output.get("HAS_TO_SAVE")
110 else 0
111 )
113 snapshot_bytes = math.ceil(real_state_bytes / compression_ratio)
114 state_storage = state_records * snapshot_bytes
115 bins = sum(size // 2 + 1 for size in shape)
116 spectra_storage = math.ceil(
117 spectra_records * max(1, bins) * max(4, real_fields) * 8
118 / compression_ratio
119 )
120 budget_storage = math.ceil(
121 budget_records * max(1, bins) * max(6, complex_fields) * 8
122 / compression_ratio
123 )
124 scalar_storage = scalar_records * 16 * 32
125 metadata_storage = 4 * 1024**2 if output.get("HAS_TO_SAVE") else 0
126 storage_bytes = (
127 state_storage
128 + spectra_storage
129 + budget_storage
130 + scalar_storage
131 + metadata_storage
132 )
133 estimated_files = (
134 state_records
135 + (2 if spectra_records else 0)
136 + (1 if budget_records else 0)
137 + (1 if scalar_records else 0)
138 + (4 if output.get("HAS_TO_SAVE") else 0)
139 )
141 resources = config["resources"]
142 ram_bytes = float(resources["ram_gib"]) * GIB
143 disk_bytes = float(resources["disk_gib"]) * GIB
144 ranks = int(resources["mpi_ranks"])
145 memory_ok = peak_bytes <= ram_bytes
146 storage_ok = storage_bytes <= disk_bytes
147 files_ok = estimated_files <= int(resources["max_output_files"])
148 unmodeled = sorted(
149 name for name, period in periods.items()
150 if period > 0 and name not in {"phys_fields", "spectra", "spect_energy_budg", "spatial_means"}
151 ) if output.get("HAS_TO_SAVE") else []
152 time_bounded = bool(time["USE_T_END"])
153 return {
154 "assumptions": {
155 "compression_ratio": compression_ratio,
156 "complex_fields": complex_fields,
157 "precision_bytes": precision_bytes,
158 "real_fields": real_fields,
159 "safety_factor": safety_factor,
160 "spectra_and_budget_layout_is_approximate": True,
161 "workspace_factor": workspace_factor,
162 "unmodeled_enabled_outputs": unmodeled,
163 },
164 "declared_bounds": {
165 "cpu_cores": resources["cpu_cores"],
166 "disk_gib": resources["disk_gib"],
167 "max_output_files": resources["max_output_files"],
168 "mpi_ranks": ranks,
169 "ram_gib": resources["ram_gib"],
170 "threads_per_rank": resources["threads_per_rank"],
171 "wall_time_minutes": resources["wall_time_minutes"],
172 },
173 "estimates": {
174 "estimated_output_files": estimated_files,
175 "idealized_peak_bytes_per_rank": math.ceil(peak_bytes / ranks),
176 "peak_memory_bytes": peak_bytes,
177 "peak_memory_gib": peak_bytes / GIB,
178 "resident_state_bytes": resident_state_bytes,
179 "state_snapshot_bytes": snapshot_bytes,
180 "state_snapshots": state_records,
181 "storage_bytes": storage_bytes,
182 "storage_gib": storage_bytes / GIB,
183 },
184 "grid": {
185 "dimensions": dimension,
186 "points": grid_points,
187 "shape": shape,
188 "spectral_points_estimate": spectral_points,
189 "spectral_shape_estimate": spectral_shape,
190 },
191 "limits": {
192 "files_within_bound": files_ok,
193 "memory_within_bound": memory_ok,
194 "storage_within_bound": storage_ok,
195 },
196 "notes": [
197 "This is an approximate planning model, not a hard memory or storage bound.",
198 "FFT decomposition, backend buffers, Python overhead, diagnostics, and MPI imbalance are backend-specific.",
199 "Runtime is not inferred from grid size; benchmark a tiny representative pilot on the target machine.",
200 "A resource fit does not establish numerical convergence or physical validity.",
201 ],
202 "ok": memory_ok and storage_ok and files_ok and time_bounded and not unmodeled,
203 "coverage": {
204 "output_model_complete": not unmodeled,
205 "time_horizon_bounded": time_bounded,
206 },
207 "runtime_estimated": False,
208 "solver": solver,
209 "tool": TOOL,
210 }
213def build_parser() -> argparse.ArgumentParser:
214 parser = argparse.ArgumentParser(
215 description=(
216 "Estimate an approximate FluidSim memory/storage envelope from strict "
217 "local JSON. No arrays are allocated and no package is imported."
218 )
219 )
220 parser.add_argument("--config", required=True)
221 parser.add_argument("--root", default=".")
222 parser.add_argument(
223 "--precision-bytes",
224 type=int,
225 choices=(4, 8),
226 default=8,
227 help="Real scalar bytes in planning model (default: 8).",
228 )
229 parser.add_argument(
230 "--workspace-factor",
231 type=float,
232 default=8.0,
233 help="Resident-state multiplier, 2..64 (default: 8).",
234 )
235 parser.add_argument(
236 "--safety-factor",
237 type=float,
238 default=1.5,
239 help="Additional margin, 1..10 (default: 1.5).",
240 )
241 parser.add_argument(
242 "--compression-ratio",
243 type=float,
244 default=1.0,
245 help="Optimistic output compression ratio, 1..20 (default: 1).",
246 )
247 return parser
250def main(argv: list[str] | None = None) -> int:
251 args = build_parser().parse_args(argv)
252 try:
253 bounded_int(
254 args.precision_bytes,
255 name="precision_bytes",
256 minimum=4,
257 maximum=8,
258 )
259 workspace = finite_float(
260 args.workspace_factor,
261 name="workspace_factor",
262 minimum=2.0,
263 maximum=64.0,
264 )
265 safety = finite_float(
266 args.safety_factor,
267 name="safety_factor",
268 minimum=1.0,
269 maximum=10.0,
270 )
271 compression = finite_float(
272 args.compression_ratio,
273 name="compression_ratio",
274 minimum=1.0,
275 maximum=20.0,
276 )
277 path = checked_input(
278 args.config,
279 root=args.root,
280 kind="file",
281 suffixes={".json"},
282 )
283 config = normalized_copy(load_json(path))
284 report = estimate(
285 config,
286 precision_bytes=args.precision_bytes,
287 workspace_factor=float(workspace),
288 safety_factor=float(safety),
289 compression_ratio=float(compression),
290 )
291 report["commands_executed"] = False
292 report["arrays_allocated"] = False
293 emit_json(report)
294 return 0 if report["ok"] else 2
295 except (OSError, ToolError, UnicodeError) as exc:
296 return fail_json(TOOL, exc)
299if __name__ == "__main__":
300 raise SystemExit(main())