Skills › Scientific Agent Skills › pufferlib › validate_plan.py
validate_plan.py
577 lines · skills/pufferlib/scripts/validate_plan.py at 5f2bc03
1#!/usr/bin/env python3
2"""Strict, dependency-free validator for PufferLib training plans."""
4from __future__ import annotations
6import argparse
7import copy
8import re
9from typing import Any
11try:
12 from ._common import (
13 LOGGER_CREDENTIAL_ENV,
14 MAX_ENVS,
15 MAX_EVAL_EPISODES,
16 MAX_STEPS,
17 MAX_WORKERS,
18 SOURCE_4_COMMIT,
19 STABLE_SDIST_SHA256,
20 UserInputError,
21 emit_json,
22 load_json_object,
23 require_keys,
24 secret_key_paths,
25 validate_slug,
26 )
27except ImportError: # Direct script execution.
28 from _common import (
29 LOGGER_CREDENTIAL_ENV,
30 MAX_ENVS,
31 MAX_EVAL_EPISODES,
32 MAX_STEPS,
33 MAX_WORKERS,
34 SOURCE_4_COMMIT,
35 STABLE_SDIST_SHA256,
36 UserInputError,
37 emit_json,
38 load_json_object,
39 require_keys,
40 secret_key_paths,
41 validate_slug,
42 )
44PROFILES = ("pypi-3.0.0", "source-4.0")
45_COMMIT = re.compile(r"^[0-9a-f]{40}$")
48def default_plan(profile: str = "pypi-3.0.0") -> dict[str, Any]:
49 """Return a bounded local plan that never starts training."""
50 if profile not in PROFILES:
51 raise UserInputError(f"profile must be one of {PROFILES}")
52 common: dict[str, Any] = {
53 "schema_version": 1,
54 "profile": profile,
55 "environment": {
56 "adapter": "synthetic",
57 "name": "synthetic",
58 "provenance_verified": True,
59 },
60 "training": {
61 "device": "cpu",
62 "horizon": 16,
63 "minibatch_size": 256,
64 "seed": 42,
65 "total_timesteps": 10_000,
66 },
67 "evaluation": {
68 "deterministic": True,
69 "episodes": 10,
70 "seed": 1_000_042,
71 "separate": True,
72 },
73 "logging": {
74 "backend": "none",
75 "disclosure_ack": False,
76 "external_opt_in": False,
77 "upload_checkpoints": False,
78 },
79 "checkpoint": {
80 "format": "state_dict",
81 "trusted_only": True,
82 },
83 }
84 if profile == "pypi-3.0.0":
85 common["package"] = {
86 "name": "pufferlib",
87 "sha256": STABLE_SDIST_SHA256,
88 "version": "3.0.0",
89 }
90 common["vectorization"] = {
91 "backend": "serial",
92 "batch_size": 4,
93 "num_envs": 4,
94 "num_workers": 1,
95 "start_method": "spawn",
96 "zero_copy": False,
97 }
98 else:
99 common["package"] = {
100 "commit": SOURCE_4_COMMIT,
101 "name": "pufferlib",
102 "version": "4.0.0-source",
103 }
104 common["vectorization"] = {
105 "backend": "torch",
106 "num_buffers": 2,
107 "num_threads": 1,
108 "start_method": "spawn",
109 "total_agents": 64,
110 }
111 return common
114def _mapping(value: Any, path: str, errors: list[str]) -> dict[str, Any]:
115 if not isinstance(value, dict):
116 errors.append(f"{path} must be an object")
117 return {}
118 return value
121def _integer(
122 value: Any,
123 *,
124 path: str,
125 minimum: int,
126 maximum: int,
127 errors: list[str],
128) -> int | None:
129 if isinstance(value, bool) or not isinstance(value, int):
130 errors.append(f"{path} must be an integer")
131 return None
132 if not minimum <= value <= maximum:
133 errors.append(f"{path} must be between {minimum} and {maximum}")
134 return None
135 return value
138def _boolean(value: Any, *, path: str, errors: list[str]) -> bool | None:
139 if type(value) is not bool:
140 errors.append(f"{path} must be boolean")
141 return None
142 return value
145def _validate_package(
146 package: dict[str, Any], profile: str, errors: list[str]
147) -> None:
148 if profile == "pypi-3.0.0":
149 errors.extend(
150 require_keys(
151 package,
152 allowed={"name", "version", "sha256"},
153 required={"name", "version", "sha256"},
154 path="$.package",
155 )
156 )
157 if package.get("version") != "3.0.0":
158 errors.append("$.package.version must equal 3.0.0")
159 if package.get("sha256") != STABLE_SDIST_SHA256:
160 errors.append("$.package.sha256 must match the published 3.0.0 sdist")
161 elif profile == "source-4.0":
162 errors.extend(
163 require_keys(
164 package,
165 allowed={"name", "version", "commit"},
166 required={"name", "version", "commit"},
167 path="$.package",
168 )
169 )
170 commit = package.get("commit")
171 if not isinstance(commit, str) or not _COMMIT.fullmatch(commit):
172 errors.append("$.package.commit must be a pinned 40-character commit")
173 if package.get("version") != "4.0.0-source":
174 errors.append("$.package.version must equal 4.0.0-source")
175 if package.get("name") != "pufferlib":
176 errors.append("$.package.name must equal pufferlib")
179def _validate_environment(environment: dict[str, Any], errors: list[str]) -> None:
180 errors.extend(
181 require_keys(
182 environment,
183 allowed={"adapter", "name", "provenance_verified"},
184 required={"adapter", "name", "provenance_verified"},
185 path="$.environment",
186 )
187 )
188 try:
189 validate_slug(environment.get("name"), name="environment.name")
190 except UserInputError as exc:
191 errors.append(str(exc))
192 if environment.get("adapter") not in {
193 "synthetic",
194 "gymnasium",
195 "pettingzoo",
196 "native-ocean",
197 }:
198 errors.append("$.environment.adapter is not allowlisted")
199 if environment.get("name") == "synthetic" and environment.get("adapter") != "synthetic":
200 errors.append("the synthetic environment must use the synthetic adapter")
201 if environment.get("name") != "synthetic" and environment.get("adapter") == "synthetic":
202 errors.append("a non-synthetic environment cannot use the synthetic adapter")
203 if environment.get("provenance_verified") is not True:
204 errors.append("$.environment.provenance_verified must be true")
207def _validate_training(training: dict[str, Any], errors: list[str]) -> None:
208 errors.extend(
209 require_keys(
210 training,
211 allowed={"device", "horizon", "minibatch_size", "seed", "total_timesteps"},
212 required={"device", "horizon", "minibatch_size", "seed", "total_timesteps"},
213 path="$.training",
214 )
215 )
216 if training.get("device") not in {"cpu", "cuda"}:
217 errors.append("$.training.device must be cpu or cuda")
218 _integer(
219 training.get("seed"),
220 path="$.training.seed",
221 minimum=0,
222 maximum=2**32 - 1,
223 errors=errors,
224 )
225 _integer(
226 training.get("total_timesteps"),
227 path="$.training.total_timesteps",
228 minimum=1,
229 maximum=MAX_STEPS,
230 errors=errors,
231 )
232 horizon = _integer(
233 training.get("horizon"),
234 path="$.training.horizon",
235 minimum=1,
236 maximum=65_536,
237 errors=errors,
238 )
239 minibatch = _integer(
240 training.get("minibatch_size"),
241 path="$.training.minibatch_size",
242 minimum=1,
243 maximum=16_777_216,
244 errors=errors,
245 )
246 if horizon and minibatch and minibatch % horizon:
247 errors.append("$.training.minibatch_size must be divisible by horizon")
250def _validate_vectorization(
251 vectorization: dict[str, Any],
252 profile: str,
253 training: dict[str, Any],
254 errors: list[str],
255) -> None:
256 if profile == "pypi-3.0.0":
257 allowed = {
258 "backend",
259 "batch_size",
260 "num_envs",
261 "num_workers",
262 "start_method",
263 "zero_copy",
264 }
265 errors.extend(
266 require_keys(
267 vectorization,
268 allowed=allowed,
269 required=allowed,
270 path="$.vectorization",
271 )
272 )
273 backend = vectorization.get("backend")
274 if backend not in {
275 "serial",
276 "multiprocessing",
277 "native",
278 }:
279 errors.append("$.vectorization.backend is invalid for PufferLib 3.0.0")
280 num_envs = _integer(
281 vectorization.get("num_envs"),
282 path="$.vectorization.num_envs",
283 minimum=1,
284 maximum=MAX_ENVS,
285 errors=errors,
286 )
287 workers = _integer(
288 vectorization.get("num_workers"),
289 path="$.vectorization.num_workers",
290 minimum=1,
291 maximum=MAX_WORKERS,
292 errors=errors,
293 )
294 batch = _integer(
295 vectorization.get("batch_size"),
296 path="$.vectorization.batch_size",
297 minimum=1,
298 maximum=MAX_ENVS,
299 errors=errors,
300 )
301 zero_copy = _boolean(
302 vectorization.get("zero_copy"),
303 path="$.vectorization.zero_copy",
304 errors=errors,
305 )
306 if vectorization.get("start_method") not in {"spawn", "forkserver"}:
307 errors.append("$.vectorization.start_method must be spawn or forkserver")
308 if num_envs and workers and num_envs % workers:
309 errors.append("$.vectorization.num_envs must be divisible by num_workers")
310 if num_envs and batch and batch > num_envs:
311 errors.append("$.vectorization.batch_size cannot exceed num_envs")
312 if num_envs and batch and zero_copy and num_envs % batch:
313 errors.append(
314 "$.vectorization.num_envs must be divisible by batch_size "
315 "when zero_copy is true"
316 )
317 if num_envs and workers and batch:
318 envs_per_worker = num_envs // workers if num_envs % workers == 0 else 0
319 if envs_per_worker and batch % envs_per_worker:
320 errors.append(
321 "$.vectorization.batch_size must be divisible by envs_per_worker"
322 )
323 if backend == "native" and num_envs != 1:
324 errors.append(
325 "$.vectorization.num_envs must equal 1 for the stable native backend"
326 )
327 else:
328 allowed = {
329 "backend",
330 "num_buffers",
331 "num_threads",
332 "start_method",
333 "total_agents",
334 }
335 errors.extend(
336 require_keys(
337 vectorization,
338 allowed=allowed,
339 required=allowed,
340 path="$.vectorization",
341 )
342 )
343 if vectorization.get("backend") not in {"native", "torch"}:
344 errors.append("$.vectorization.backend must be native or torch")
345 agents = _integer(
346 vectorization.get("total_agents"),
347 path="$.vectorization.total_agents",
348 minimum=1,
349 maximum=MAX_ENVS,
350 errors=errors,
351 )
352 buffers = _integer(
353 vectorization.get("num_buffers"),
354 path="$.vectorization.num_buffers",
355 minimum=1,
356 maximum=256,
357 errors=errors,
358 )
359 _integer(
360 vectorization.get("num_threads"),
361 path="$.vectorization.num_threads",
362 minimum=1,
363 maximum=MAX_WORKERS,
364 errors=errors,
365 )
366 if vectorization.get("start_method") != "spawn":
367 errors.append("$.vectorization.start_method must be spawn for 4.0")
368 if agents and buffers and agents % buffers:
369 errors.append("$.vectorization.total_agents must be divisible by num_buffers")
370 horizon = training.get("horizon")
371 minibatch = training.get("minibatch_size")
372 if (
373 isinstance(agents, int)
374 and isinstance(horizon, int)
375 and isinstance(minibatch, int)
376 and minibatch > agents * horizon
377 ):
378 errors.append(
379 "$.training.minibatch_size cannot exceed total_agents * horizon"
380 )
383def _validate_evaluation(
384 evaluation: dict[str, Any], training: dict[str, Any], errors: list[str]
385) -> None:
386 allowed = {"deterministic", "episodes", "seed", "separate"}
387 errors.extend(
388 require_keys(
389 evaluation,
390 allowed=allowed,
391 required=allowed,
392 path="$.evaluation",
393 )
394 )
395 _boolean(evaluation.get("deterministic"), path="$.evaluation.deterministic", errors=errors)
396 _boolean(evaluation.get("separate"), path="$.evaluation.separate", errors=errors)
397 _integer(
398 evaluation.get("episodes"),
399 path="$.evaluation.episodes",
400 minimum=1,
401 maximum=MAX_EVAL_EPISODES,
402 errors=errors,
403 )
404 _integer(
405 evaluation.get("seed"),
406 path="$.evaluation.seed",
407 minimum=0,
408 maximum=2**32 - 1,
409 errors=errors,
410 )
411 if evaluation.get("separate") is not True:
412 errors.append("$.evaluation.separate must be true")
413 if evaluation.get("seed") == training.get("seed"):
414 errors.append("training and evaluation seeds must differ")
417def _validate_logging(
418 logging: dict[str, Any], profile: str, errors: list[str]
419) -> None:
420 allowed = {
421 "backend",
422 "disclosure_ack",
423 "external_opt_in",
424 "upload_checkpoints",
425 }
426 errors.extend(
427 require_keys(logging, allowed=allowed, required=allowed, path="$.logging")
428 )
429 backend = logging.get("backend")
430 if backend not in LOGGER_CREDENTIAL_ENV:
431 errors.append("$.logging.backend must be none, wandb, or neptune")
432 return
433 if backend == "neptune":
434 errors.append("Neptune hosted service shut down on 2026-03-05; use local logging")
435 opt_in = _boolean(
436 logging.get("external_opt_in"),
437 path="$.logging.external_opt_in",
438 errors=errors,
439 )
440 disclosure = _boolean(
441 logging.get("disclosure_ack"),
442 path="$.logging.disclosure_ack",
443 errors=errors,
444 )
445 upload = _boolean(
446 logging.get("upload_checkpoints"),
447 path="$.logging.upload_checkpoints",
448 errors=errors,
449 )
450 if backend == "none" and any(value is True for value in (opt_in, disclosure, upload)):
451 errors.append("external logging flags must be false when backend is none")
452 if backend != "none" and (opt_in is not True or disclosure is not True):
453 errors.append(
454 "external logging requires external_opt_in=true and disclosure_ack=true"
455 )
456 if backend == "wandb" and upload is not True:
457 errors.append(
458 "reviewed upstream W&B training uploads model artifacts at close; "
459 "upload_checkpoints=true is required (no supported CLI upload opt-out)"
460 )
463def _validate_checkpoint(checkpoint: dict[str, Any], errors: list[str]) -> None:
464 allowed = {"format", "trusted_only"}
465 errors.extend(
466 require_keys(
467 checkpoint,
468 allowed=allowed,
469 required=allowed,
470 path="$.checkpoint",
471 )
472 )
473 if checkpoint.get("format") not in {"state_dict", "native-bin", "opaque"}:
474 errors.append("$.checkpoint.format is invalid")
475 if checkpoint.get("trusted_only") is not True:
476 errors.append("$.checkpoint.trusted_only must be true")
479def validate_plan(plan: Any) -> list[str]:
480 """Return all deterministic schema and safety errors."""
481 errors: list[str] = []
482 if not isinstance(plan, dict):
483 return ["$ must be an object"]
484 secret_paths = secret_key_paths(plan)
485 if secret_paths:
486 errors.append(
487 "credential-bearing keys are forbidden in plans: " + ", ".join(secret_paths)
488 )
489 top_allowed = {
490 "checkpoint",
491 "environment",
492 "evaluation",
493 "logging",
494 "package",
495 "profile",
496 "schema_version",
497 "training",
498 "vectorization",
499 }
500 errors.extend(
501 require_keys(
502 plan,
503 allowed=top_allowed,
504 required=top_allowed,
505 path="$",
506 )
507 )
508 if plan.get("schema_version") != 1:
509 errors.append("$.schema_version must equal 1")
510 profile = plan.get("profile")
511 if profile not in PROFILES:
512 errors.append(f"$.profile must be one of {PROFILES}")
513 return errors
515 package = _mapping(plan.get("package"), "$.package", errors)
516 environment = _mapping(plan.get("environment"), "$.environment", errors)
517 training = _mapping(plan.get("training"), "$.training", errors)
518 vectorization = _mapping(plan.get("vectorization"), "$.vectorization", errors)
519 evaluation = _mapping(plan.get("evaluation"), "$.evaluation", errors)
520 logging = _mapping(plan.get("logging"), "$.logging", errors)
521 checkpoint = _mapping(plan.get("checkpoint"), "$.checkpoint", errors)
523 _validate_package(package, profile, errors)
524 _validate_environment(environment, errors)
525 if profile == "source-4.0" and environment.get("adapter") in {"gymnasium", "pettingzoo"}:
526 errors.append("source-4.0 has no Gymnasium/PettingZoo emulation adapter")
527 _validate_training(training, errors)
528 _validate_vectorization(vectorization, profile, training, errors)
529 _validate_evaluation(evaluation, training, errors)
530 _validate_logging(logging, profile, errors)
531 _validate_checkpoint(checkpoint, errors)
532 return errors
535def build_parser() -> argparse.ArgumentParser:
536 parser = argparse.ArgumentParser(
537 description=(
538 "Validate a strict local JSON plan. With no --config, validate the "
539 "bounded built-in dry-run plan."
540 )
541 )
542 parser.add_argument("--config", help="Explicit JSON file beneath --root")
543 parser.add_argument("--root", default=".", help="Allowed local path root")
544 parser.add_argument("--profile", choices=PROFILES, default="pypi-3.0.0")
545 parser.add_argument("--compact", action="store_true")
546 return parser
549def main(argv: list[str] | None = None) -> int:
550 parser = build_parser()
551 args = parser.parse_args(argv)
552 try:
553 plan = (
554 load_json_object(args.config, root=args.root)
555 if args.config
556 else default_plan(args.profile)
557 )
558 errors = validate_plan(plan)
559 report = {
560 "errors": errors,
561 "network_used": False,
562 "plan": copy.deepcopy(plan),
563 "status": "valid" if not errors else "invalid",
564 }
565 except UserInputError as exc:
566 report = {
567 "errors": [str(exc)],
568 "network_used": False,
569 "plan": None,
570 "status": "invalid",
571 }
572 emit_json(report, pretty=not args.compact)
573 return 0 if report["status"] == "valid" else 1
576if __name__ == "__main__":
577 raise SystemExit(main())