| """Single, thin command-line entry point for GNN4Colliders.""" |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import os |
| import sys |
| from importlib import resources |
|
|
| from hydra import compose, initialize_config_dir |
| from omegaconf import OmegaConf |
|
|
| from gnn4colliders.config.application import ( |
| export_onnx, |
| predict_or_evaluate, |
| prepare, |
| prepare_shard, |
| train, |
| write_shard_manifest, |
| ) |
|
|
| _COMMANDS = { |
| "prepare", |
| "prepare-manifest", |
| "prepare-shard", |
| "train", |
| "evaluate", |
| "predict", |
| "export", |
| } |
|
|
|
|
| def _is_main_process() -> bool: |
| return int(os.environ.get("RANK", "0")) == 0 |
|
|
|
|
| def _help(command: str | None = None) -> None: |
| if command: |
| print(f"Usage: gnn4colliders {command} [--config-name NAME] [key=value ...]") |
| print("Compose Hydra configuration with semantic overrides.") |
| else: |
| print( |
| "Usage: gnn4colliders <prepare|prepare-manifest|prepare-shard|train|" |
| "evaluate|predict|export> " |
| "[key=value ...]" |
| ) |
| print("Use gnn4colliders <command> --help for command help.") |
|
|
|
|
| def _config(overrides: list[str], *, config_name: str = "config"): |
| config_resource = resources.files("gnn4colliders.configs") |
| with resources.as_file(config_resource) as config_dir: |
| with initialize_config_dir(version_base=None, config_dir=str(config_dir)): |
| return compose(config_name=config_name, overrides=overrides) |
|
|
|
|
| def _config_name(args: list[str]) -> tuple[str, list[str]]: |
| """Extract the optional config name while preserving Hydra overrides.""" |
|
|
| remaining: list[str] = [] |
| config_name = "config" |
| index = 0 |
| while index < len(args): |
| argument = args[index] |
| if argument == "--config-name": |
| if index + 1 >= len(args): |
| raise ValueError("--config-name requires a config name") |
| config_name = args[index + 1] |
| index += 2 |
| continue |
| if argument.startswith("--config-name="): |
| config_name = argument.split("=", 1)[1] |
| if not config_name: |
| raise ValueError("--config-name requires a config name") |
| index += 1 |
| continue |
| remaining.append(argument) |
| index += 1 |
| return config_name, remaining |
|
|
|
|
| def _shard_options( |
| args: list[str], |
| ) -> tuple[int | None, int | None, str | None, list[str]]: |
| values: dict[str, int | str] = {} |
| remaining: list[str] = [] |
| names = { |
| "--shard-index": "index", |
| "--shard-count": "count", |
| "--output-dir": "output", |
| } |
| index = 0 |
| while index < len(args): |
| argument = args[index] |
| matched = next((name for name in names if argument == name), None) |
| if matched: |
| if index + 1 >= len(args): |
| raise ValueError(f"{matched} requires a value") |
| values[names[matched]] = args[index + 1] |
| index += 2 |
| continue |
| prefix = next((name for name in names if argument.startswith(name + "=")), None) |
| if prefix: |
| values[names[prefix]] = argument.split("=", 1)[1] |
| index += 1 |
| continue |
| remaining.append(argument) |
| index += 1 |
| shard_index = int(values["index"]) if "index" in values else None |
| shard_count = int(values["count"]) if "count" in values else None |
| output_dir = str(values["output"]) if "output" in values else None |
| return shard_index, shard_count, output_dir, remaining |
|
|
|
|
| def main(argv: list[str] | None = None) -> int: |
| args = list(sys.argv[1:] if argv is None else argv) |
| if not args or args[0] in {"-h", "--help"}: |
| _help() |
| return 0 |
| command = args.pop(0) |
| if command not in _COMMANDS: |
| print( |
| f"unknown command {command!r}; choose from {', '.join(sorted(_COMMANDS))}", |
| file=sys.stderr, |
| ) |
| return 2 |
| if "--help" in args or "-h" in args: |
| _help(command) |
| return 0 |
| try: |
| config_name, overrides = _config_name(args) |
| shard_index, shard_count, output_dir, overrides = _shard_options(overrides) |
| config = _config(overrides, config_name=config_name) |
| logging.basicConfig( |
| level=getattr(logging, str(config.logging.level).upper(), logging.INFO) |
| ) |
| if command == "prepare": |
| print(f"prepared cache: {prepare(config)}") |
| elif command == "prepare-manifest": |
| if shard_count is None or output_dir is None: |
| raise ValueError( |
| "prepare-manifest requires --shard-count and --output-dir" |
| ) |
| manifest = write_shard_manifest( |
| config, shard_count=shard_count, output_dir=output_dir |
| ) |
| print(f"wrote shard manifest: {manifest}") |
| elif command == "prepare-shard": |
| if shard_index is None or shard_count is None or output_dir is None: |
| raise ValueError( |
| "prepare-shard requires --shard-index, --shard-count, and " |
| "--output-dir" |
| ) |
| shard = prepare_shard( |
| config, |
| shard_index=shard_index, |
| shard_count=shard_count, |
| output_dir=output_dir, |
| ) |
| print(f"prepared event shard: {shard}") |
| elif command == "train": |
| |
| checkpoint = train(config) |
| if _is_main_process(): |
| print(f"saved checkpoint: {checkpoint}") |
| elif command == "export": |
| if _is_main_process(): |
| print(f"wrote ONNX model: {export_onnx(config)}") |
| elif command == "evaluate": |
| result = predict_or_evaluate(config, evaluate=True) |
| if _is_main_process(): |
| print(OmegaConf.to_yaml(result.metrics)) |
| else: |
| result = predict_or_evaluate(config) |
| if _is_main_process(): |
| print( |
| f"wrote predictions: {config.inference.output} " |
| f"({len(result.sample_ids)} events)" |
| ) |
| return 0 |
| except (ValueError, FileNotFoundError, KeyError, ImportError) as error: |
| print(f"gnn4colliders: {error}", file=sys.stderr) |
| return 2 |
|
|
|
|
| __all__ = ["main"] |
|
|
