Review round 2: fail loudly on bad inputs, report partial results

docs/cli.md

@@ -163,9 +163,10 @@ The summary contains:
 | `status` | `ok`, `partial`, `error` or `dry_run` |
 | `exit_code` | The process exit code (see below) |
 | `command`, `output_file` | What ran and where results were written |
-| `device`, `model`, `config` | The resolved device, the model, and the full config (Hugging Face token masked) |
+| `device`, `model`, `config` | The resolved device, the model (`type`, `name`, and `path` when a checkpoint was loaded), and the full config (Hugging Face token masked) |
 | `n_input`, `n_returned` | Number of input structures and of returned structures |
-| `structures` | Per returned structure: `index`, `n_atoms`, `charge`, `multiplicity`, `energy_ev`, `comment` |
+| `n_generated` | `ensemble` only: conformers generated (pruning can leave fewer than requested) |
+| `structures` | Per returned structure: `index`, `n_atoms`, `charge`, `multiplicity`, `energy_ev`, `comment`, `converged` |
 | `diagnostics` | The engine diagnostic report (embedding and optimization events, convergence where known) |
 | `error` | `{"type", "message"}` when the command failed, else `null` |
 | `wall_time_s`, `gpuma_version`, `schema_version` | Run metadata |
@@ -183,7 +184,7 @@ The summary contains:
 | `1` | Runtime failure (for example an optimization crashed) |
 | `2` | Command-line usage error |
 | `3` | Invalid input, config value, model name or device, as detected by GPUMA's own checks |
-| `4` | Partial batch: some structures failed to optimize |
+| `4` | Partial result: some structures failed to optimize or stopped at `max_steps` without converging |
 
 ## Checking before running
 

docs/config.md

@@ -25,7 +25,8 @@ Unknown fields are preserved. **Always use a config file for CLI and API calls.*
     "multiplicity": 1,
 
     "force_convergence_criterion": 5e-2,
-    "energy_convergence_criterion": null
+    "energy_convergence_criterion": null,
+    "max_steps": 10000
   },
 
   "model": {
@@ -140,6 +141,7 @@ optimization:
 
   force_convergence_criterion: 5.0e-2
   energy_convergence_criterion: null
+  max_steps: 10000
 
 model:
   model_type: fairchem
@@ -180,12 +182,13 @@ technical:
 
 | Parameter | Default | Description |
 |---|---|---|
-| `batch_optimization_mode` | `"batch"` | `"sequential"` (ASE per structure) or `"batch"` (torch-sim GPU-accelerated) |
+| `batch_optimization_mode` | `"batch"` | `"sequential"` (ASE per structure) or `"batch"` (torch-sim GPU-accelerated). If a batch fails, its structures are retried one at a time, so one bad input fails only itself |
 | `batch_optimizer` | `"fire"` | Optimizer for both single and batch modes: `"fire"`, `"gradient_descent"`, `"lbfgs"`, or `"bfgs"`. In single-structure (ASE) mode, maps to ASE FIRE/BFGS/LBFGS (`"gradient_descent"` falls back to FIRE). Defaults to `"fire"` if unset; an unknown value is an error |
-| `charge` | `0` | Total charge of the system (inferred from SMILES, not overridden) |
-| `multiplicity` | `1` | Spin multiplicity of the system |
+| `charge` | `0` | Total charge of the system (inferred from SMILES, not overridden). Must be an integer |
+| `multiplicity` | `1` | Spin multiplicity of the system. A multiplicity the electron count can't have (e.g. a singlet radical) is an error |
 | `force_convergence_criterion` | `5e-2` | Force convergence threshold (eV/A). Used for both single and batch modes |
-| `energy_convergence_criterion` | `null` | Energy convergence threshold. Batch mode only; force takes precedence if both set |
+| `energy_convergence_criterion` | `null` | Energy convergence threshold (eV): the energy change of one step. Used only when `force_convergence_criterion` is `null`, in both modes |
+| `max_steps` | `10000` | Step limit per structure. A structure that reaches it is returned but marked not converged, and the CLI exits with code 4 |
 
 ### `model`
 
@@ -195,7 +198,7 @@ technical:
 | `model_name` | `"uma-s-1p2p1-ft-d4rev"` | Model identifier (see [Available Models](#available-models)) |
 | `model_modal` | `null` | SevenNet multi-modal fidelity selector. Valid values depend on the checkpoint (`7net-omni`: e.g. `"omol25_high"`, `"omol25_low"`, `"spice"`, `"qcml"`; `7net-mf-ompa`: `"mpa"`, `"omat24"`). Required for multi-modal checkpoints; ignored by single-modal checkpoints and by other backends |
 | `model_path` | `null` | Local checkpoint path (Fairchem and SevenNet; overrides `model_name`). A path that doesn't exist, or one set for ORB, is an error |
-| `model_cache_dir` | `null` | Directory for cached model downloads |
+| `model_cache_dir` | `null` | Directory for cached model downloads (Fairchem only; ORB and SevenNet use their own caches, so setting it for them is an error). A directory that can't be created is an error |
 | `huggingface_token` | `null` | HuggingFace token for model access |
 | `huggingface_token_file` | `null` | File path to read the HF token from |
 | `d3_correction` | `false` | Enable D3 dispersion correction. Supported for the ORB, Fairchem/UMA, and SevenNet backends; rejected for `uma-s-1p2p1-ft-d4rev` |
@@ -207,7 +210,7 @@ technical:
 | Parameter | Default | Description |
 |---|---|---|
 | `max_num_conformers` | `20` | Maximum number of conformers to generate from SMILES |
-| `conformer_seed` | `42` | Random seed for conformer generation |
+| `conformer_seed` | `42` | Random seed for conformer generation, used by every command that embeds SMILES |
 
 ### `technical`
 
@@ -218,8 +221,8 @@ technical:
 | `memory_scaling_factor` | `1.75` | Factor by which the autobatcher grows the probe size during calibration. Larger = faster calibration but coarser final batch size; smaller = slower but tighter. A job runs as one batch without probing only if its forward-pass peak times this factor fits in `max_memory_padding` of free memory. Must be > 1 |
 | `max_atoms_to_try` | `100000` | Upper bound on the autobatcher's calibration probe size (atoms) |
 | `memory_scales_with` | `"auto"` | Per-structure memory metric for batch sizing: `"n_atoms"`, `"n_edges"` or `"auto"`. `"auto"` uses `n_atoms` for UMA, whose large fixed memory per structure makes edge counts overcharge small molecules, and `n_edges` for ORB and SevenNet |
-| `steps_between_swaps` | `1` | Optimization steps between batch swaps in the in-flight autobatcher. `1` is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower |
-| `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"` |
+| `steps_between_swaps` | `1` | Optimization steps between batch swaps in the in-flight autobatcher. `1` is fastest on this codebase's screenings (uma-s/uma-m/orb); higher values are monotonically slower. Must be at least 1 |
+| `logging_level` | `"INFO"` | Logging verbosity: `"DEBUG"`, `"INFO"`, `"WARNING"`, `"ERROR"`, `"CRITICAL"` |
 
 > You can also control the GPU with the `CUDA_VISIBLE_DEVICES` environment variable.
 

examples/config.json

@@ -8,7 +8,8 @@
     "multiplicity": 1,
 
     "force_convergence_criterion": 5e-2,
-    "energy_convergence_criterion": null
+    "energy_convergence_criterion": null,
+    "max_steps": 10000
 
   },
 

examples/config_orb.json

@@ -8,7 +8,8 @@
     "multiplicity": 1,
 
     "force_convergence_criterion": 5e-2,
-    "energy_convergence_criterion": null
+    "energy_convergence_criterion": null,
+    "max_steps": 10000
 
   },
 

examples/config_sevennet.json

@@ -8,7 +8,8 @@
     "multiplicity": 1,
 
     "force_convergence_criterion": 5e-2,
-    "energy_convergence_criterion": null
+    "energy_convergence_criterion": null,
+    "max_steps": 10000
 
   },
 

src/gpuma/agent_skills.py

@@ -41,7 +41,7 @@ def install_skill(targets: Sequence[Path] | None = None, force: bool = False) ->
             elif dest.exists():
                 shutil.rmtree(dest)
             shutil.copytree(source, dest, ignore=shutil.ignore_patterns("__pycache__"))
-            (dest / VERSION_FILE).write_text(f"{gpuma_version()}\n")
+            (dest / VERSION_FILE).write_text(f"{gpuma_version() or 'unknown'}\n")
     return dests
 
 

src/gpuma/api.py

@@ -144,6 +144,15 @@ def optimize_ensemble_smiles(
             conformers are not valid.
 
     """
+    return _optimize_ensemble_smiles(smiles, output_file, config)[0]
+
+
+def _optimize_ensemble_smiles(
+    smiles: str,
+    output_file: str | None = None,
+    config: Config | None = None,
+) -> tuple[list[Structure], int]:
+    """As :func:`optimize_ensemble_smiles`, plus the number of conformers generated."""
     if config is None:
         config = load_config_from_file()
     multiplicity = int(config.optimization.multiplicity)
@@ -164,13 +173,13 @@ def optimize_ensemble_smiles(
     # conformer that failed to optimize must not renumber the ones after it.
     optimized = [(i, s) for i, s in enumerate(results) if s is not None]
 
-    if output_file:
+    if output_file and optimized:
         comments = [
             f"Optimized conformer {i + 1} from SMILES: {smiles}" for i, _ in optimized
         ]
         save_multi_xyz([s for _, s in optimized], output_file, comments)
 
-    return [s for _, s in optimized]
+    return [s for _, s in optimized], len(conformers)
 
 
 def optimize_batch_multi_xyz_file(
@@ -218,7 +227,7 @@ def optimize_batch_multi_xyz_file(
     # leaves a gap rather than shifting every later label onto the wrong input.
     optimized = [(i, s) for i, s in enumerate(results) if s is not None]
 
-    if output_file:
+    if output_file and optimized:
         comments = [
             f"Optimized structure {i + 1} from: {input_file}" for i, _ in optimized
         ]
@@ -266,9 +275,9 @@ def optimize_batch_xyz_directory(
     # Numbered by position in the directory listing; see the note above.
     optimized = [(i, s) for i, s in enumerate(results) if s is not None]
 
-    if output_file:
+    if output_file and optimized:
         comments = [
-            f"Optimized structure {i + 1} from batch input" for i, _ in optimized
+            f"Optimized from: {structures[i].metadata['source_file']}" for i, _ in optimized
         ]
         save_multi_xyz([s for _, s in optimized], output_file, comments)
 

src/gpuma/cli.py

@@ -27,9 +27,9 @@
 
 from .agent_skills import install_skill
 from .api import (
+    _optimize_ensemble_smiles,
     optimize_batch_multi_xyz_file,
     optimize_batch_xyz_directory,
-    optimize_ensemble_smiles,
     optimize_single_smiles,
     optimize_single_xyz_file,
 )
@@ -49,6 +49,7 @@
 from .config import (
     Config,
     check_d3_allowed,
+    check_model_cache_dir,
     check_model_path,
     load_config_from_file,
     resolve_model_type,
@@ -58,6 +59,7 @@
 from .doctor import recommended_device, run_checks
 from .models.base import AVAILABLE_FAIRCHEM_MODELS, AVAILABLE_SEVENNET_MODELS
 from .models.device import _parse_device_string
+from .structure import Structure, check_spin_state
 from .utils.io_handler import (
     read_multi_xyz,
     read_xyz,
@@ -587,14 +589,7 @@ def cmd_optimize(args, config: Config) -> CommandResult:
         optimized.energy if optimized.energy is not None else float("nan"),
     )
     logger.info("Optimized structure saved to %s", args.output)
-    return CommandResult(
-        "optimize",
-        "ok",
-        EXIT_OK,
-        output_file=args.output,
-        n_input=1,
-        structures=[structure_record(0, optimized)],
-    )
+    return _optimized_result("optimize", args, [optimized], n_input=1, n_expected=1)
 
 
 def cmd_ensemble(args, config: Config) -> CommandResult:
@@ -610,38 +605,72 @@ def cmd_ensemble(args, config: Config) -> CommandResult:
     logger.info("Generating %d conformers for SMILES: %s", num_conf, args.smiles)
     config.conformer_generation.max_num_conformers = num_conf
 
-    optimized_conformers = optimize_ensemble_smiles(
+    optimized_conformers, n_generated = _optimize_ensemble_smiles(
         smiles=args.smiles,
         output_file=args.output,
         config=config,
     )
 
     if not optimized_conformers:
-        raise RuntimeError("No conformers found")
+        raise RuntimeError("No conformers were optimized")
 
     logger.info(
-        "Ensemble optimization completed successfully! %d optimized conformers saved to %s",
+        "Ensemble optimization finished: %d optimized conformers saved to %s",
         len(optimized_conformers),
         args.output,
     )
-    return CommandResult(
+    return _optimized_result(
         "ensemble",
-        "ok",
-        EXIT_OK,
-        output_file=args.output,
+        args,
+        optimized_conformers,
         n_input=num_conf,
-        structures=[structure_record(i, s) for i, s in enumerate(optimized_conformers)],
+        n_expected=n_generated,
+        extra={"n_generated": n_generated},
+    )
+
+
+def _optimized_result(
+    command: str,
+    args,
+    structures: list[Structure],
+    n_input: int,
+    n_expected: int,
+    extra: dict | None = None,
+) -> CommandResult:
+    """Return ``ok``, or ``partial`` when a structure failed or didn't converge."""
+    n_failed = n_expected - len(structures)
+    n_unconverged = sum(s.converged is False for s in structures)
+    if n_failed:
+        logger.warning(
+            "Only %d of %d structures were optimized; see the log for the failures.",
+            len(structures),
+            n_expected,
+        )
+    if n_unconverged:
+        logger.warning("%d structure(s) did not converge within max_steps.", n_unconverged)
+    partial = n_failed or n_unconverged
+    return CommandResult(
+        command,
+        "partial" if partial else "ok",
+        EXIT_PARTIAL if partial else EXIT_OK,
+        output_file=args.output,
+        n_input=n_input,
+        structures=[structure_record(i, s) for i, s in enumerate(structures)],
+        extra=extra or {},
     )
 
 
 _DRY_RUN_COMMANDS = ("optimize", "smiles", "ensemble", "batch")
 
 
-def _check_smiles(smiles: str) -> None:
+def _check_smiles(smiles: str, multiplicity: int) -> None:
     from rdkit import Chem
 
-    if Chem.MolFromSmiles(smiles) is None:
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
         raise ValueError(f"Invalid SMILES: {smiles!r}")
+    symbols = [atom.GetSymbol() for atom in Chem.AddHs(mol).GetAtoms()]
+    check_spin_state(symbols, Chem.GetFormalCharge(mol), multiplicity)
 
 
 _DEFAULT_MODEL_NAMES = {
@@ -683,26 +712,28 @@ def _check_inputs(args, config: Config) -> int:
     """Validate inputs, model and device without loading a model; return ``n_input``."""
     # CLI overrides apply after Config validation, so re-check what they can change.
     check_model_path(config)
+    check_model_cache_dir(config)
     check_d3_allowed(config)
     _check_model(config)
     _parse_device_string(str(config.technical.device))
     charge = int(config.optimization.charge)
     multiplicity = int(config.optimization.multiplicity)
+    if args.command in ("optimize", "smiles", "ensemble") and args.smiles:
+        _check_smiles(args.smiles, multiplicity)
+        return _resolve_conformer_count(args, config) if args.command == "ensemble" else 1
     if args.command in ("optimize", "smiles"):
-        if args.smiles:
-            _check_smiles(args.smiles)
-        else:
-            read_xyz(args.xyz, charge=charge, multiplicity=multiplicity)
-        return 1
-    if args.command == "ensemble":
-        _check_smiles(args.smiles)
-        return _resolve_conformer_count(args, config)
-    if args.multi_xyz:
+        structures = [read_xyz(args.xyz, charge=charge, multiplicity=multiplicity)]
+    elif args.multi_xyz:
         structures = read_multi_xyz(args.multi_xyz, charge=charge, multiplicity=multiplicity)
     else:
         structures = read_xyz_directory(args.xyz_dir, charge=charge, multiplicity=multiplicity)
     if not structures:
         raise ValueError("No input structures found")
+    for i, s in enumerate(structures):
+        try:
+            check_spin_state(s.symbols, s.charge, s.multiplicity)
+        except ValueError as exc:
+            raise ValueError(f"structure {i + 1}: {exc}") from exc
     return len(structures)
 
 
@@ -751,32 +782,8 @@ def cmd_batch(args, config: Config) -> CommandResult:
     if not structures:
         raise RuntimeError("No structures were optimized")
 
-    records = [structure_record(i, s) for i, s in enumerate(structures)]
-    if len(structures) < n_input:
-        logger.warning(
-            "Only %d of %d structures were optimized; see the log for the failures.",
-            len(structures),
-            n_input,
-        )
-        return CommandResult(
-            "batch",
-            "partial",
-            EXIT_PARTIAL,
-            output_file=args.output,
-            n_input=n_input,
-            structures=records,
-        )
-
-    logger.info("Batch optimization completed successfully!")
     logger.info("Optimized structures saved to %s", args.output)
-    return CommandResult(
-        "batch",
-        "ok",
-        EXIT_OK,
-        output_file=args.output,
-        n_input=n_input,
-        structures=records,
-    )
+    return _optimized_result("batch", args, structures, n_input=n_input, n_expected=n_input)
 
 
 def cmd_convert(args, config: Config) -> CommandResult:
@@ -785,6 +792,7 @@ def cmd_convert(args, config: Config) -> CommandResult:
     This command generates a single 3D structure from SMILES without running
     any optimization. ``config`` selects the embedding device.
     """
+    _parse_device_string(str(config.technical.device))
     logger.info("Converting SMILES '%s' to XYZ without optimization", args.smiles)
     structure = smiles_to_xyz(args.smiles, config=config)
     save_xyz_file(structure, args.output)
@@ -805,6 +813,7 @@ def cmd_generate(args, config: Config) -> CommandResult:
     This command generates a conformer ensemble from a SMILES string without
     optimization.
     """
+    _parse_device_string(str(config.technical.device))
     num_conf = _resolve_conformer_count(args, config)
     logger.info(
         "Generating %d conformers for SMILES (no optimization): %s",