Skills › Scientific Agent Skills › geniml › embedding_plan.py
embedding_plan.py
482 lines · skills/geniml/scripts/embedding_plan.py at 5f2bc03
1#!/usr/bin/env python3
2"""Plan a bounded local Geniml embedding run without importing ML packages."""
4from __future__ import annotations
6import argparse
7import os
8import stat
9import sys
10from collections import Counter
11from pathlib import Path
13from _common import (
14 HARD_MAX_BYTES,
15 HARD_MAX_EPOCHS,
16 HARD_MAX_WORKERS,
17 SafetyError,
18 add_path_mode_argument,
19 display_path,
20 fail_json,
21 int_type,
22 local_path,
23 print_json,
24 sha256_file,
25 simple_yaml_mapping,
26)
29TOOL = "embedding-run-planner"
32def build_parser() -> argparse.ArgumentParser:
33 parser = argparse.ArgumentParser(
34 description=(
35 "Validate local inputs and emit a redacted Region2Vec, scEmbed, or "
36 "BEDspace run plan. No package, model, binary, or network call is executed."
37 )
38 )
39 parser.add_argument("--mode", choices=("region2vec", "scembed", "bedspace"), required=True)
40 parser.add_argument("--data", required=True, help="Token Parquet or BED directory.")
41 parser.add_argument("--metadata", help="BEDspace metadata CSV.")
42 parser.add_argument("--universe", required=True)
43 parser.add_argument("--output-dir", required=True)
44 parser.add_argument("--model-dir", help="Optional trusted local model bundle.")
45 parser.add_argument("--starspace-dir", help="BEDspace-only local StarSpace directory.")
46 parser.add_argument("--assembly", required=True)
47 parser.add_argument(
48 "--split-unit",
49 choices=("patient", "donor", "biological-replicate", "sample", "cell"),
50 default="patient",
51 )
52 parser.add_argument(
53 "--embedding-dim",
54 type=int_type(minimum=1, maximum=4096, label="embedding-dim"),
55 default=100,
56 )
57 parser.add_argument(
58 "--epochs",
59 type=int_type(minimum=1, maximum=HARD_MAX_EPOCHS, label="epochs"),
60 default=10,
61 )
62 parser.add_argument(
63 "--workers",
64 type=int_type(minimum=1, maximum=HARD_MAX_WORKERS, label="workers"),
65 default=4,
66 )
67 parser.add_argument(
68 "--window-size",
69 type=int_type(minimum=1, maximum=100_000, label="window-size"),
70 default=5,
71 )
72 parser.add_argument(
73 "--min-count",
74 type=int_type(minimum=1, maximum=10_000_000, label="min-count"),
75 default=10,
76 )
77 parser.add_argument(
78 "--batch-size",
79 type=int_type(minimum=1, maximum=1_000_000, label="batch-size"),
80 default=64,
81 )
82 parser.add_argument(
83 "--seed",
84 type=int_type(minimum=0, maximum=2**32 - 1, label="seed"),
85 default=42,
86 )
87 parser.add_argument("--pooling", choices=("mean", "max"), default="mean")
88 parser.add_argument("--geniml-version", default="0.8.4")
89 parser.add_argument("--gtars-version", default="0.10.0")
90 parser.add_argument(
91 "--max-file-bytes",
92 type=int_type(minimum=1, maximum=HARD_MAX_BYTES, label="max-file-bytes"),
93 default=2 * 1024**3,
94 )
95 parser.add_argument(
96 "--max-total-bytes",
97 type=int_type(minimum=1, maximum=HARD_MAX_BYTES, label="max-total-bytes"),
98 default=4 * 1024**3,
99 )
100 add_path_mode_argument(parser)
101 return parser
104def _parquet_magic(path: Path) -> bool:
105 flags = os.O_RDONLY
106 if hasattr(os, "O_NOFOLLOW"):
107 flags |= os.O_NOFOLLOW
108 descriptor = os.open(path, flags)
109 try:
110 info = os.fstat(descriptor)
111 if not stat.S_ISREG(info.st_mode) or info.st_size < 8:
112 return False
113 first = os.read(descriptor, 4)
114 os.lseek(descriptor, -4, os.SEEK_END)
115 last = os.read(descriptor, 4)
116 return first == b"PAR1" and last == b"PAR1"
117 finally:
118 os.close(descriptor)
121def _render(path: Path, label: str, mode: str) -> str:
122 if mode == "full":
123 return str(path)
124 if mode == "basename":
125 return path.name
126 return f"<{label}>"
129def plan(args: argparse.Namespace) -> tuple[dict, int]:
130 if not args.assembly.strip() or len(args.assembly) > 200:
131 raise SafetyError("assembly must be a nonempty value of at most 200 characters")
132 if args.mode != "bedspace" and (args.metadata or args.starspace_dir):
133 raise SafetyError("metadata and starspace-dir apply only to bedspace")
134 if args.mode == "bedspace" and not args.metadata:
135 raise SafetyError("bedspace requires --metadata")
136 if args.mode == "bedspace" and args.model_dir:
137 raise SafetyError("model-dir applies only to Region2Vec or scEmbed bundles")
139 data = local_path(args.data, kind="dir" if args.mode == "bedspace" else "file")
140 universe = local_path(args.universe, kind="file")
141 output_dir = local_path(args.output_dir, kind="dir")
142 model_dir = local_path(args.model_dir, kind="dir") if args.model_dir else None
143 metadata = local_path(args.metadata, kind="file") if args.metadata else None
144 starspace_dir = (
145 local_path(args.starspace_dir, kind="dir") if args.starspace_dir else None
146 )
148 errors: Counter[str] = Counter()
149 warnings: Counter[str] = Counter()
150 if args.split_unit == "cell":
151 errors["nonindependent_split_unit:cell"] += 1
152 elif args.split_unit == "sample":
153 warnings["verify_samples_are_independent_across_patients"] += 1
154 if any(output_dir.iterdir()):
155 warnings["output_directory_not_empty"] += 1
157 universe_digest, universe_size = sha256_file(
158 universe,
159 max_bytes=args.max_file_bytes,
160 )
161 input_summary: dict[str, object] = {
162 "data": display_path(data, 1, args.path_mode),
163 "universe": display_path(universe, 2, args.path_mode),
164 "universe_sha256": universe_digest,
165 "universe_size_bytes": universe_size,
166 "output_dir": display_path(output_dir, 3, args.path_mode),
167 }
169 if args.mode in {"region2vec", "scembed"}:
170 data_digest, data_size = sha256_file(data, max_bytes=args.max_file_bytes)
171 parquet_ok = _parquet_magic(data)
172 if not parquet_ok:
173 errors["data_not_parquet"] += 1
174 input_summary.update(
175 {
176 "data_sha256": data_digest,
177 "data_size_bytes": data_size,
178 "parquet_magic_valid": parquet_ok,
179 "parquet_schema_inspected": False,
180 }
181 )
182 else:
183 bed_entries = list(data.iterdir())
184 if len(bed_entries) > 100_000:
185 raise SafetyError("BEDspace input directory exceeds 100,000 entries")
186 symlinks = sum(1 for entry in bed_entries if entry.is_symlink())
187 special = sum(
188 1
189 for entry in bed_entries
190 if not entry.is_symlink() and not entry.is_file()
191 )
192 if symlinks:
193 errors["bedspace_data_symlinks"] = symlinks
194 if special:
195 errors["bedspace_non_file_entries"] = special
196 total_bed_bytes = 0
197 oversized_beds = 0
198 for entry in bed_entries:
199 if entry.is_symlink() or not entry.is_file():
200 continue
201 size = entry.lstat().st_size
202 total_bed_bytes += size
203 if size > args.max_file_bytes:
204 oversized_beds += 1
205 if total_bed_bytes > args.max_total_bytes:
206 raise SafetyError("BEDspace input directory exceeds max-total-bytes")
207 if oversized_beds:
208 errors["bedspace_files_exceed_byte_limit"] = oversized_beds
209 metadata_digest, metadata_size = sha256_file(
210 metadata,
211 max_bytes=args.max_file_bytes,
212 )
213 input_summary.update(
214 {
215 "bed_directory_entries": len(bed_entries),
216 "bed_directory_bytes": total_bed_bytes,
217 "metadata": display_path(metadata, 4, args.path_mode),
218 "metadata_sha256": metadata_digest,
219 "metadata_size_bytes": metadata_size,
220 "metadata_values_emitted": False,
221 }
222 )
224 model_summary = None
225 if model_dir:
226 config = local_path(str(model_dir / "config.yaml"), kind="file")
227 checkpoint = local_path(str(model_dir / "checkpoint.pt"), kind="file")
228 bundle_universe = local_path(str(model_dir / "universe.bed"), kind="file")
229 config_metadata = simple_yaml_mapping(
230 config,
231 max_bytes=min(args.max_file_bytes, 16 * 1024**2),
232 )
233 config_digest, config_size = sha256_file(
234 config,
235 max_bytes=min(args.max_file_bytes, 16 * 1024**2),
236 )
237 checkpoint_digest, checkpoint_size = sha256_file(
238 checkpoint,
239 max_bytes=args.max_file_bytes,
240 )
241 bundle_digest, bundle_size = sha256_file(
242 bundle_universe,
243 max_bytes=args.max_file_bytes,
244 )
245 if bundle_digest != universe_digest:
246 errors["model_universe_checksum_mismatch"] += 1
247 if config_metadata.get("embedding_dim", config_metadata.get("embedding_size")) != (
248 args.embedding_dim
249 ):
250 warnings["planned_embedding_dim_differs_from_model"] += 1
251 model_summary = {
252 "path": display_path(model_dir, 5, args.path_mode),
253 "config_sha256": config_digest,
254 "config_size_bytes": config_size,
255 "checkpoint_sha256": checkpoint_digest,
256 "checkpoint_size_bytes": checkpoint_size,
257 "bundle_universe_sha256": bundle_digest,
258 "bundle_universe_size_bytes": bundle_size,
259 "metadata": {
260 "vocab_size": config_metadata.get("vocab_size"),
261 "embedding_dim": config_metadata.get(
262 "embedding_dim",
263 config_metadata.get("embedding_size"),
264 ),
265 "pooling_method": config_metadata.get("pooling_method"),
266 },
267 "deserialized": False,
268 }
270 stages: list[dict] = [
271 {
272 "stage": "environment",
273 "executed": False,
274 "requirements": [
275 f"geniml=={args.geniml_version}",
276 f"gtars=={args.gtars_version}",
277 "anndata==0.12.19",
278 "scanpy==1.12.4",
279 "zarr==2.18.7",
280 "torch==2.14.1",
281 "gensim==4.4.0",
282 "huggingface-hub==2.0.0",
283 ],
284 "instruction": "resolve with uv and retain the generated lockfile",
285 },
286 {
287 "stage": "data_contract",
288 "executed": False,
289 "instruction": (
290 "verify 0-based half-open coordinates, one assembly/chromosome-sizes "
291 "digest, bounded token counts, and patient/donor-grouped splits"
292 ),
293 },
294 ]
296 if args.mode == "region2vec":
297 stages.extend(
298 [
299 {
300 "stage": "api_smoke",
301 "executed": False,
302 "imports": [
303 "geniml.region2vec.main.Region2VecExModel",
304 "geniml.region2vec.utils.Region2VecDataset",
305 "gtars.tokenizers.Tokenizer",
306 ],
307 "note": "do not use the broken 0.8.4 top-level region2vec export",
308 },
309 {
310 "stage": "train",
311 "executed": False,
312 "parameters": {
313 "embedding_dim": args.embedding_dim,
314 "epochs": args.epochs,
315 "window_size": args.window_size,
316 "min_count": args.min_count,
317 "num_cpus": args.workers,
318 "seed": args.seed,
319 "pooling_method": args.pooling,
320 },
321 },
322 ]
323 )
324 elif args.mode == "scembed":
325 stages.extend(
326 [
327 {
328 "stage": "api_smoke",
329 "executed": False,
330 "imports": [
331 "geniml.scembed.main.ScEmbed",
332 "geniml.region2vec.utils.Region2VecDataset",
333 "gtars.tokenizers.Tokenizer",
334 ],
335 "note": (
336 "Geniml 0.8.4 ScEmbed.encode/tokenize_anndata fail with Gtars "
337 "0.10.0; use the explicit token/projection recipe in references/scembed.md"
338 ),
339 },
340 {
341 "stage": "train",
342 "executed": False,
343 "parameters": {
344 "embedding_dim": args.embedding_dim,
345 "epochs": args.epochs,
346 "window_size": args.window_size,
347 "min_count": args.min_count,
348 "num_cpus": args.workers,
349 "seed": args.seed,
350 "pooling_method": args.pooling,
351 },
352 },
353 ]
354 )
355 else:
356 starspace_summary = None
357 if starspace_dir:
358 binary = local_path(str(starspace_dir / "starspace"), kind="file")
359 binary_digest, binary_size = sha256_file(
360 binary,
361 max_bytes=args.max_file_bytes,
362 )
363 starspace_summary = {
364 "directory": display_path(starspace_dir, 6, args.path_mode),
365 "binary_sha256": binary_digest,
366 "binary_size_bytes": binary_size,
367 "executed": False,
368 }
369 else:
370 warnings["starspace_binary_required"] += 1
371 preprocess_template = [
372 "geniml",
373 "bedspace",
374 "preprocess",
375 "--input",
376 "<LOCAL_BED_DIRECTORY>",
377 "--metadata",
378 "<LOCAL_METADATA_CSV>",
379 "--universe",
380 "<LOCAL_UNIVERSE_BED>",
381 "--output",
382 "<LOCAL_OUTPUT_DIRECTORY_WITH_TRAILING_SEPARATOR>",
383 ]
384 train_template = [
385 "geniml",
386 "bedspace",
387 "train",
388 "-s",
389 "<LOCAL_STARSPACE_DIRECTORY>",
390 "--input",
391 "<LOCAL_OUTPUT_DIRECTORY>/train_input.txt",
392 "--output",
393 "<LOCAL_OUTPUT_DIRECTORY>",
394 "--dim",
395 str(args.embedding_dim),
396 "--epochs",
397 str(args.epochs),
398 ]
399 stages.extend(
400 [
401 {
402 "stage": "legacy_dependency_review",
403 "executed": False,
404 "starspace": starspace_summary,
405 "note": (
406 "StarSpace is archived; Geniml has no compatibility pin and "
407 "hard-codes 20 training threads"
408 ),
409 },
410 {
411 "stage": "preprocess",
412 "executed": False,
413 "argv_template": preprocess_template,
414 },
415 {
416 "stage": "train",
417 "executed": False,
418 "argv_template": train_template,
419 "note": "the 0.8.4 bedspace search dispatcher is not usable",
420 },
421 ]
422 )
424 stages.append(
425 {
426 "stage": "export_and_verify",
427 "executed": False,
428 "instruction": (
429 "export to a fresh directory, copy the exact universe.bed, create "
430 "SHA-256 provenance, and inspect without deserialization"
431 ),
432 "expected_local_output": _render(
433 output_dir,
434 "LOCAL_OUTPUT_DIRECTORY",
435 args.path_mode,
436 ),
437 }
438 )
440 report = {
441 "ok": not errors,
442 "ready_to_execute": not errors
443 and (args.mode != "bedspace" or starspace_dir is not None),
444 "tool": TOOL,
445 "contract": {
446 "mode": args.mode,
447 "assembly": args.assembly,
448 "coordinate_system": "0-based-half-open",
449 "split_unit": args.split_unit,
450 "network_used": False,
451 "packages_imported": False,
452 "models_deserialized": False,
453 "commands_executed": False,
454 },
455 "inputs": input_summary,
456 "existing_model": model_summary,
457 "errors": dict(sorted(errors.items())),
458 "warnings": dict(sorted(warnings.items())),
459 "stages": stages,
460 "required_postconditions": [
461 "token IDs are within the pinned tokenizer vocabulary",
462 "universe bytes/order and special-token IDs match model config",
463 "all output artifacts have provenance and SHA-256 checksums",
464 "evaluation uses patients/donors excluded from fitting and model selection",
465 "logs contain aggregate/redacted values only",
466 ],
467 }
468 return report, 0 if not errors else 2
471def main() -> int:
472 args = build_parser().parse_args()
473 try:
474 report, exit_code = plan(args)
475 print_json(report)
476 return exit_code
477 except (OSError, SafetyError, UnicodeError) as exc:
478 return fail_json(TOOL, exc)
481if __name__ == "__main__":
482 sys.exit(main())