From daba334920e1efaac9c8a565e4c226ed57da6c39 Mon Sep 17 00:00:00 2001 From: xraymemory Date: Tue, 14 Jul 2026 15:54:00 -0400 Subject: [PATCH 1/2] refactor(config): distinguish model names from objects --- GRID_SEARCH.md | 2 +- run_grid_search.py | 37 ++++++------ src/sampleworks/cli/guidance.py | 2 +- .../utils/guidance_script_arguments.py | 59 +++++++++++------- .../utils/guidance_script_utils.py | 6 +- tests/cli/test_guidance_cli.py | 30 +++++----- .../integration/test_mismatch_integration.py | 2 +- tests/test_run_grid_search.py | 54 +++++++++++++++++ tests/utils/conftest.py | 2 +- tests/utils/test_guidance_script_arguments.py | 60 +++++++++++++++++-- tests/utils/test_guidance_script_utils.py | 12 ++-- 11 files changed, 196 insertions(+), 70 deletions(-) create mode 100644 tests/test_run_grid_search.py diff --git a/GRID_SEARCH.md b/GRID_SEARCH.md index 8b98400d..f258ebf9 100644 --- a/GRID_SEARCH.md +++ b/GRID_SEARCH.md @@ -159,7 +159,7 @@ see what trial is actually being run by looking for lines like this one: ```log 2026-03-10 02:03:23.664 | INFO | sampleworks.utils.guidance_script_utils:run_guidance_job_queue:589 - Running job 1/15: GuidanceConfig(protein='5I MV_1.0occB', structure='/data/inputs/processed/5IMV/5IMV_single_001_density_input.cif', density='/data/inputs/occ_sweeps/1.0occB/density_maps/5IMV_1.0 -occB_1.00A.ccp4', model='boltz2', guidance_type='pure_guidance', log_path='/data/results/5IMV_1.0occB/boltz2_X-RAY_DIFFRACTION/pure_guidance/ens8_gw0. +occB_1.00A.ccp4', model_name='boltz2', guidance_type='pure_guidance', log_path='/data/results/5IMV_1.0occB/boltz2_X-RAY_DIFFRACTION/pure_guidance/ens8_gw0. 1/run.log', output_dir='/data/results/5IMV_1.0occB/boltz2_X-RAY_DIFFRACTION/pure_guidance/ens8_gw0.1', partial_diffusion_step=120, loss_order=2, resol ution=1.0, device='cuda:0', gradient_normalization=True, em=False, guidance_start=-1, augmentation=True, align_to_input=True) ``` diff --git a/run_grid_search.py b/run_grid_search.py index d3f77df7..b3336417 100755 --- a/run_grid_search.py +++ b/run_grid_search.py @@ -27,7 +27,7 @@ class GridSearchConfig: """Serializable summary of the grid-search dimensions and output location.""" - model: str + model_name: str scalers: list[str] ensemble_sizes: list[int] gradient_weights: list[float] @@ -180,7 +180,7 @@ def build_args_for_process_pool( protein=job.protein, structure=job.structure_path, density=job.density_path, - model=job.model, + model_name=job.model_name, guidance_type=job.scaler, log_path=job.log_path, output_dir=job.output_dir, @@ -266,7 +266,7 @@ def run_grid_search( futures = {} for worker_num, job_queue_path in enumerate(job_queue_paths): - model = worker_job_queues[worker_num][0].model + model = worker_job_queues[worker_num][0].model_name future = executor.submit( run_guidance_queue_script, job_queue_path, @@ -286,13 +286,13 @@ def run_grid_search( if r.status == "success": successful += 1 log.info( - f"SUCCESS ({r.protein}, {r.model}, {r.method}, {r.scaler} " + f"SUCCESS ({r.protein}, {r.model_name}, {r.method}, {r.scaler} " f"{r.runtime_seconds:.1f}s): {r.log_path}" ) else: failed += 1 log.error( - f"FAILED ({r.protein}, {r.model}, {r.method}, {r.scaler} " + f"FAILED ({r.protein}, {r.model_name}, {r.method}, {r.scaler} " f"exit={r.exit_code}): {r.log_path}" ) except Exception as e: @@ -447,7 +447,7 @@ def main(args: argparse.Namespace): return config = GridSearchConfig( - model=args.model, + model_name=args.model, scalers=args.scalers.split(), ensemble_sizes=[int(x) for x in args.ensemble_sizes.split()], gradient_weights=[float(x) for x in args.gradient_weights.split()], @@ -506,7 +506,7 @@ def generate_jobs(args: argparse.Namespace) -> list[JobConfig]: structure_path=structure, density_path=density, resolution=resolution, - model=model, + model_name=model, scaler=scaler, ensemble_size=ens, gradient_weight=gw, @@ -533,7 +533,7 @@ def generate_jobs(args: argparse.Namespace) -> list[JobConfig]: structure_path=structure, density_path=density, resolution=resolution, - model=model, + model_name=model, scaler=scaler, ensemble_size=ens, gradient_weight=gw, @@ -573,7 +573,7 @@ def save_results( new_run_keys = { ( r.protein, - r.model, + r.model_name, r.method, r.scaler, r.ensemble_size, @@ -585,17 +585,20 @@ def save_results( merged_runs = [asdict(r) for r in results] for existing_run in existing_runs: + normalized_run = existing_run.copy() + if "model" in normalized_run: + normalized_run.setdefault("model_name", normalized_run.pop("model")) key = ( - existing_run.get("protein"), - existing_run.get("model"), - existing_run.get("method"), - existing_run.get("scaler"), - existing_run.get("ensemble_size"), - existing_run.get("gradient_weight"), - existing_run.get("gd_steps"), + normalized_run.get("protein"), + normalized_run.get("model_name"), + normalized_run.get("method"), + normalized_run.get("scaler"), + normalized_run.get("ensemble_size"), + normalized_run.get("gradient_weight"), + normalized_run.get("gd_steps"), ) if key not in new_run_keys: - merged_runs.append(existing_run) + merged_runs.append(normalized_run) output = { "config": asdict(config), diff --git a/src/sampleworks/cli/guidance.py b/src/sampleworks/cli/guidance.py index 737757b7..ef945a77 100644 --- a/src/sampleworks/cli/guidance.py +++ b/src/sampleworks/cli/guidance.py @@ -18,7 +18,7 @@ def main(argv: list[str] | None = None) -> int: device, model_wrapper = get_model_and_device( config.device, getattr(config, "model_checkpoint", None), - config.model, + config.model_name, config=config, ) result = run_guidance(config, config.guidance_type, model_wrapper, device) diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index c1f318cd..66a959de 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -216,7 +216,7 @@ class GuidanceConfig: protein: str structure: Path | str # actually a path to a structure file density: Path | str - model: str | StructurePredictor + model_name: str | StructurePredictor guidance_type: str | GuidanceType log_path: str output_dir: str = "output" @@ -242,23 +242,23 @@ def add_argument(self, name: str, default: Any = None, **kwargs): def from_cli( cls, argv: list[str] | None = None, - model: str | None = None, + model_name: str | None = None, guidance_type: str | None = None, ) -> GuidanceConfig: """Parse CLI arguments and return a fully populated GuidanceConfig. - When *model* and *guidance_type* are provided (e.g. from legacy + When *model_name* and *guidance_type* are provided (e.g. from legacy scripts), they are used directly and ``--model`` / ``--guidance-type`` are not required on the command line. Otherwise they are parsed as required CLI arguments. """ model_choices = [m.value for m in StructurePredictor] guidance_choices = [g.value for g in GuidanceType] - model_preset = model is not None + model_preset = model_name is not None guidance_preset = guidance_type is not None - if model_preset and model not in model_choices: - raise ValueError(f"Unknown model type: {model}") + if model_preset and model_name not in model_choices: + raise ValueError(f"Unknown model type: {model_name}") if guidance_preset and guidance_type not in guidance_choices: raise ValueError(f"Unknown guidance type: {guidance_type}") @@ -268,6 +268,7 @@ def from_cli( if not model_preset: pre.add_argument( "--model", + dest="model_name", type=str, required=True, choices=model_choices, @@ -282,17 +283,21 @@ def from_cli( help="Guidance method", ) pre_args, _ = pre.parse_known_args(argv) - model = model or pre_args.model + model_name = model_name or pre_args.model_name guidance_type = guidance_type or pre_args.guidance_type + assert model_name is not None + assert guidance_type is not None + # -- full parser ----------------------------------------------------- parser = argparse.ArgumentParser( - description=f"Run {guidance_type} guidance with {model}", + description=f"Run {guidance_type} guidance with {model_name}", ) parser.add_argument( "--model", + dest="model_name", type=str, - default=model, + default=model_name, choices=model_choices, help=argparse.SUPPRESS if model_preset else "Structure prediction model", ) @@ -310,14 +315,14 @@ def from_cli( help="Protein identifier (must match naming used in grid search / evaluation)", ) add_generic_args(parser) - _MODEL_ARG_ADDERS[model](parser) + _MODEL_ARG_ADDERS[model_name](parser) _GUIDANCE_ARG_ADDERS[guidance_type](parser) args = parser.parse_args(argv) - if model_preset and args.model != model: + if model_preset and args.model_name != model_name: parser.error( - f"This script is fixed to --model {model}." + f"This script is fixed to --model {model_name}." f" Use sampleworks-guidance for other models." ) if guidance_preset and args.guidance_type != guidance_type: @@ -330,7 +335,7 @@ def from_cli( protein=args.protein, structure=args.structure, density=args.density, - model=model, + model_name=model_name, guidance_type=guidance_type, log_path=getattr(args, "log_path", None) or "", output_dir=args.output_dir, @@ -363,9 +368,9 @@ def __post_init__(self): raise ValueError(f"Unknown guidance type: {self.guidance_type}") try: - _MODEL_ARG_ADDERS[self.model](self) + _MODEL_ARG_ADDERS[self.model_name](self) except KeyError: - raise ValueError(f"Unknown model type: {self.model}") + raise ValueError(f"Unknown model type: {self.model_name}") def populate_config_for_guidance_type(self, job: JobConfig, args: argparse.Namespace): """Apply per-job grid-search values onto this guidance configuration.""" @@ -374,13 +379,13 @@ def populate_config_for_guidance_type(self, job: JobConfig, args: argparse.Names self.model_checkpoint = checkpoint elif not getattr(self, "model_checkpoint", None): # Auto-resolve from baked-in /checkpoints/ or legacy fallback paths - model_key = str(self.model).lower().replace("structurepredictor.", "") + model_key = str(self.model_name).lower().replace("structurepredictor.", "") self.model_checkpoint = _resolve_checkpoint(model_key) - if job.model == StructurePredictor.BOLTZ_2 and job.method: + if job.model_name == StructurePredictor.BOLTZ_2 and job.method: self.method = job.method - if job.model == StructurePredictor.RF3: + if job.model_name == StructurePredictor.RF3: self.disable_chiral_features = getattr(args, "disable_chiral_features", False) self.track_chiral_features = getattr(args, "track_chiral_features", False) @@ -410,6 +415,13 @@ def as_dict(self) -> dict[str, Any]: output["log_path"] = _remap_container_path(str(self.log_path)) return output + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore state while migrating legacy pickles from ``model``.""" + migrated = state.copy() + if "model" in migrated: + migrated.setdefault("model_name", migrated.pop("model")) + self.__dict__.update(migrated) + def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig): """Add CLI arguments shared by all models and guidance methods.""" @@ -648,7 +660,7 @@ class JobConfig: structure_path: Path | str density_path: Path | str resolution: float - model: str + model_name: str scaler: str ensemble_size: int gradient_weight: float @@ -663,7 +675,7 @@ class JobResult: """Serializable status record produced after a guidance job finishes.""" protein: str - model: str + model_name: str method: str | None scaler: str ensemble_size: int @@ -694,3 +706,10 @@ def as_dict(self) -> dict[str, Any]: output["output_dir"] = _remap_container_path(str(self.output_dir)) output["log_path"] = _remap_container_path(str(self.log_path)) return output + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore state while migrating legacy pickles from ``model``.""" + migrated = state.copy() + if "model" in migrated: + migrated.setdefault("model_name", migrated.pop("model")) + self.__dict__.update(migrated) diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index 0f511a61..8b51f4c0 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/src/sampleworks/utils/guidance_script_utils.py @@ -715,7 +715,7 @@ def get_job_result( end_time = epoch_seconds(ended_at) result = JobResult( protein=args.protein, - model=args.model, + model_name=args.model_name, method=getattr(args, "method", None), scaler=args.guidance_type, ensemble_size=args.ensemble_size, @@ -740,7 +740,7 @@ def run_guidance_job_queue(job_queue_path: str) -> list[JobResult]: template_job = job_queue[0] if template_job.model_checkpoint is None or template_job.model_checkpoint == "": # Auto-resolve from baked-in /checkpoints/ or legacy fallback paths - model_key = str(template_job.model).lower().replace("structurepredictor.", "") + model_key = str(template_job.model_name).lower().replace("structurepredictor.", "") resolved = _resolve_checkpoint(model_key) # will raise if not found template_job.model_checkpoint = resolved # Propagate to all jobs in the queue @@ -751,7 +751,7 @@ def run_guidance_job_queue(job_queue_path: str) -> list[JobResult]: device, model_wrapper = get_model_and_device( str(template_job.device), template_job.model_checkpoint, - template_job.model, # this is not actually the model, it's the model name, e.g. boltz2 + template_job.model_name, config=template_job, ) job_results = [] diff --git a/tests/cli/test_guidance_cli.py b/tests/cli/test_guidance_cli.py index ea415edb..a916a4da 100644 --- a/tests/cli/test_guidance_cli.py +++ b/tests/cli/test_guidance_cli.py @@ -29,7 +29,7 @@ class TestFromCliUnified: def test_boltz2_pure_guidance(self): argv = ["--model", "boltz2", "--guidance-type", "pure_guidance"] + COMMON_ARGS config = GuidanceConfig.from_cli(argv) - assert config.model == "boltz2" + assert config.model_name == "boltz2" assert config.guidance_type == "pure_guidance" assert config.protein == "1VME" assert config.structure == "test.cif" @@ -39,7 +39,7 @@ def test_boltz2_pure_guidance(self): def test_rf3_fk_steering(self): argv = ["--model", "rf3", "--guidance-type", "fk_steering"] + COMMON_ARGS config = GuidanceConfig.from_cli(argv) - assert config.model == "rf3" + assert config.model_name == "rf3" assert config.guidance_type == "fk_steering" assert hasattr(config, "num_particles") assert hasattr(config, "fk_lambda") @@ -123,10 +123,10 @@ class TestFromCliLegacyScripts: def test_preset_model_and_guidance_type(self): config = GuidanceConfig.from_cli( COMMON_ARGS, - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", ) - assert config.model == "boltz2" + assert config.model_name == "boltz2" assert config.guidance_type == "pure_guidance" assert config.protein == "1VME" @@ -134,10 +134,10 @@ def test_no_model_guidance_type_on_cli_required(self): """Legacy scripts should not require --model/--guidance-type on CLI.""" config = GuidanceConfig.from_cli( COMMON_ARGS, - model="protenix", + model_name="protenix", guidance_type="fk_steering", ) - assert config.model == "protenix" + assert config.model_name == "protenix" assert config.guidance_type == "fk_steering" @pytest.mark.parametrize("model", ["boltz1", "boltz2", "protenix", "rf3"]) @@ -145,10 +145,10 @@ def test_no_model_guidance_type_on_cli_required(self): def test_all_model_guidance_combos(self, model, guidance_type): config = GuidanceConfig.from_cli( COMMON_ARGS, - model=model, + model_name=model, guidance_type=guidance_type, ) - assert config.model == model + assert config.model_name == model assert config.guidance_type == guidance_type @@ -272,12 +272,12 @@ def test_msa_path_rejected_for_boltz1(self): def test_conflicting_model_in_preset_mode(self): argv = ["--model", "boltz2"] + COMMON_ARGS with pytest.raises(SystemExit): - GuidanceConfig.from_cli(argv, model="rf3", guidance_type="fk_steering") + GuidanceConfig.from_cli(argv, model_name="rf3", guidance_type="fk_steering") def test_conflicting_guidance_type_in_preset_mode(self): argv = ["--guidance-type", "fk_steering"] + COMMON_ARGS with pytest.raises(SystemExit): - GuidanceConfig.from_cli(argv, model="boltz1", guidance_type="pure_guidance") + GuidanceConfig.from_cli(argv, model_name="boltz1", guidance_type="pure_guidance") def test_fk_args_rejected_for_pure_guidance(self): argv = [ @@ -332,12 +332,12 @@ class TestPresetMatchingAccepted: def test_matching_model_accepted(self): argv = ["--model", "boltz1"] + COMMON_ARGS - config = GuidanceConfig.from_cli(argv, model="boltz1", guidance_type="fk_steering") - assert config.model == "boltz1" + config = GuidanceConfig.from_cli(argv, model_name="boltz1", guidance_type="fk_steering") + assert config.model_name == "boltz1" def test_matching_guidance_type_accepted(self): argv = ["--guidance-type", "pure_guidance"] + COMMON_ARGS - config = GuidanceConfig.from_cli(argv, model="boltz2", guidance_type="pure_guidance") + config = GuidanceConfig.from_cli(argv, model_name="boltz2", guidance_type="pure_guidance") assert config.guidance_type == "pure_guidance" @@ -583,11 +583,11 @@ def test_help_exits_zero(self): def test_invalid_preset_model_raises(self): with pytest.raises(ValueError, match="Unknown model type"): - GuidanceConfig.from_cli(COMMON_ARGS, model="typo", guidance_type="fk_steering") + GuidanceConfig.from_cli(COMMON_ARGS, model_name="typo", guidance_type="fk_steering") def test_invalid_preset_guidance_type_raises(self): with pytest.raises(ValueError, match="Unknown guidance type"): - GuidanceConfig.from_cli(COMMON_ARGS, model="boltz1", guidance_type="typo") + GuidanceConfig.from_cli(COMMON_ARGS, model_name="boltz1", guidance_type="typo") class TestAlignmentReverseDiffusion: diff --git a/tests/integration/test_mismatch_integration.py b/tests/integration/test_mismatch_integration.py index 01393d4b..9daf8f7b 100644 --- a/tests/integration/test_mismatch_integration.py +++ b/tests/integration/test_mismatch_integration.py @@ -1016,7 +1016,7 @@ def test_save_with_model_template(self, tmp_path: Path): protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(tmp_path), diff --git a/tests/test_run_grid_search.py b/tests/test_run_grid_search.py new file mode 100644 index 00000000..2eda9fee --- /dev/null +++ b/tests/test_run_grid_search.py @@ -0,0 +1,54 @@ +"""Tests for grid-search result serialization.""" + +import json + +from run_grid_search import GridSearchConfig, save_results +from sampleworks.utils.guidance_script_arguments import JobResult + + +def test_save_results_normalizes_legacy_model_key(tmp_path) -> None: + """Existing result records migrate to ``model_name`` without duplicates.""" + legacy_run = { + "protein": "1abc", + "model": "boltz2", + "method": None, + "scaler": "pure_guidance", + "ensemble_size": 1, + "gradient_weight": 0.1, + "gd_steps": 1, + "status": "failed", + } + (tmp_path / "results.json").write_text(json.dumps({"runs": [legacy_run]})) + result = JobResult( + protein="1abc", + model_name="boltz2", + method=None, + scaler="pure_guidance", + ensemble_size=1, + gradient_weight=0.1, + gd_steps=1, + status="success", + exit_code=0, + runtime_seconds=1.0, + started_at="2026-07-13T00:00:00", + finished_at="2026-07-13T00:00:01", + log_path="run.log", + output_dir="output", + ) + config = GridSearchConfig( + model_name="boltz2", + scalers=["pure_guidance"], + ensemble_sizes=[1], + gradient_weights=[0.1], + gd_steps=[1], + method="", + proteins_file="proteins.csv", + output_dir=str(tmp_path), + ) + + save_results([result], config, str(tmp_path), total_time=1.0) + + saved = json.loads((tmp_path / "results.json").read_text()) + assert len(saved["runs"]) == 1 + assert saved["runs"][0]["model_name"] == "boltz2" + assert "model" not in saved["runs"][0] diff --git a/tests/utils/conftest.py b/tests/utils/conftest.py index d626d2fb..eb86734c 100644 --- a/tests/utils/conftest.py +++ b/tests/utils/conftest.py @@ -15,7 +15,7 @@ def guidance_job_result(tmp_path: Path) -> JobResult: """Successful JobResult populated with paths under ``tmp_path``.""" return JobResult( protein="1l63", - model="boltz2", + model_name="boltz2", method=None, scaler="pure_guidance", ensemble_size=8, diff --git a/tests/utils/test_guidance_script_arguments.py b/tests/utils/test_guidance_script_arguments.py index f18d5aa1..0b2f14e3 100644 --- a/tests/utils/test_guidance_script_arguments.py +++ b/tests/utils/test_guidance_script_arguments.py @@ -2,6 +2,7 @@ from __future__ import annotations +import pickle from argparse import Namespace from pathlib import Path from unittest.mock import patch @@ -13,6 +14,7 @@ get_checkpoint, GuidanceConfig, JobConfig, + JobResult, validate_model_checkpoint, ) @@ -54,7 +56,7 @@ def _build_job(model: StructurePredictor) -> JobConfig: structure_path="/tmp/structure.cif", density_path="/tmp/density.mrc", resolution=2.0, - model=model, + model_name=model, scaler=GuidanceType.PURE_GUIDANCE, ensemble_size=1, gradient_weight=0.1, @@ -75,7 +77,7 @@ def test_populate_config_resolves_checkpoint_when_none_provided(_mock_resolve, m protein="protein", structure="/tmp/structure.cif", density="/tmp/density.mrc", - model=model_wrapper_type, + model_name=model_wrapper_type, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/tmp/output/run.log", ) @@ -98,7 +100,7 @@ def test_populate_config_uses_model_checkpoint_argument(model_wrapper_type): protein="protein", structure="/tmp/structure.cif", density="/tmp/density.mrc", - model=model_wrapper_type, + model_name=model_wrapper_type, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/tmp/output/run.log", ) @@ -272,7 +274,7 @@ def test_as_dict_remaps_all_four_path_fields(monkeypatch): protein="1abc", structure="/data/inputs/structures/1abc.cif", density="/data/inputs/maps/1abc.ccp4", - model=StructurePredictor.BOLTZ_2, + model_name=StructurePredictor.BOLTZ_2, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/data/results/1abc/boltz2/run.log", output_dir="/data/results/1abc/boltz2", @@ -297,7 +299,7 @@ def test_as_dict_unchanged_when_no_env_vars(monkeypatch): protein="1abc", structure="/data/inputs/structures/1abc.cif", density="/data/inputs/maps/1abc.ccp4", - model=StructurePredictor.BOLTZ_2, + model_name=StructurePredictor.BOLTZ_2, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/data/results/1abc/run.log", output_dir="/data/results/1abc", @@ -317,7 +319,7 @@ def test_as_dict_remaps_with_catchall_host_dir(monkeypatch): protein="1abc", structure="/data/inputs/structures/1abc.cif", density="/data/inputs/maps/1abc.ccp4", - model=StructurePredictor.BOLTZ_2, + model_name=StructurePredictor.BOLTZ_2, guidance_type=GuidanceType.PURE_GUIDANCE, log_path="/data/results/1abc/run.log", output_dir="/data/results/1abc", @@ -328,3 +330,49 @@ def test_as_dict_remaps_with_catchall_host_dir(monkeypatch): assert d["output_dir"] == "/mnt/storage/results/1abc" assert d["log_path"] == "/mnt/storage/results/1abc/run.log" assert d["protein"] == "1abc" + + +def test_guidance_config_migrates_legacy_model_pickle() -> None: + """Old job queues restore ``model`` state as ``model_name``.""" + config = GuidanceConfig( + protein="1abc", + structure="structure.cif", + density="density.ccp4", + model_name=StructurePredictor.BOLTZ_2, + guidance_type=GuidanceType.PURE_GUIDANCE, + log_path="run.log", + ) + config.__dict__["model"] = config.__dict__.pop("model_name") + + restored = pickle.loads(pickle.dumps(config)) + + assert restored.model_name == StructurePredictor.BOLTZ_2 + assert "model" not in restored.__dict__ + assert "model" not in restored.as_dict() + + +def test_job_result_migrates_legacy_model_pickle() -> None: + """Old result pickles restore ``model`` state as ``model_name``.""" + result = JobResult( + protein="1abc", + model_name="boltz2", + method=None, + scaler="pure_guidance", + ensemble_size=1, + gradient_weight=0.1, + gd_steps=1, + status="success", + exit_code=0, + runtime_seconds=1.0, + started_at="2026-07-13T00:00:00", + finished_at="2026-07-13T00:00:01", + log_path="run.log", + output_dir="output", + ) + result.__dict__["model"] = result.__dict__.pop("model_name") + + restored = pickle.loads(pickle.dumps(result)) + + assert restored.model_name == "boltz2" + assert "model" not in restored.__dict__ + assert "model" not in restored.as_dict() diff --git a/tests/utils/test_guidance_script_utils.py b/tests/utils/test_guidance_script_utils.py index a91a0fa5..f09176c0 100644 --- a/tests/utils/test_guidance_script_utils.py +++ b/tests/utils/test_guidance_script_utils.py @@ -44,7 +44,7 @@ def test_save_everything_uses_model_atom_array_for_mismatch(tmp_path: Path): protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(tmp_path), @@ -122,7 +122,7 @@ def test_write_job_metadata_with_job_result_appends_timing_and_status( protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(tmp_path), @@ -134,6 +134,8 @@ def test_write_job_metadata_with_job_result_appends_timing_and_status( # GuidanceConfig keys are preserved assert metadata["protein"] == "1l63" assert metadata["guidance_type"] == "pure_guidance" + assert metadata["model_name"] == "boltz2" + assert "model" not in metadata # JobResult-only keys are appended assert metadata["started_at"] == "2026-05-05T10:00:00" assert metadata["finished_at"] == "2026-05-05T10:00:12.340000" @@ -151,7 +153,7 @@ def test_write_job_metadata_creates_missing_output_dir( protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path="dummy", output_dir=str(nested), @@ -182,14 +184,14 @@ def test_write_job_metadata_remaps_job_result_paths_to_host( protein="1l63", structure=Path("dummy"), density=Path("dummy"), - model="boltz2", + model_name="boltz2", guidance_type="pure_guidance", log_path=container_log, output_dir=container_output, ) job_result = JobResult( protein="1l63", - model="boltz2", + model_name="boltz2", method=None, scaler="pure_guidance", ensemble_size=8, From fa6ac7a7c627a2c3ad4f89db44f3575a081cb56f Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 16 Jul 2026 15:45:16 -0400 Subject: [PATCH 2/2] refactor(config): drop legacy model-name migration --- run_grid_search.py | 19 +++---- .../utils/guidance_script_arguments.py | 18 +------ tests/test_run_grid_search.py | 54 ------------------- tests/utils/test_guidance_script_arguments.py | 48 ----------------- 4 files changed, 10 insertions(+), 129 deletions(-) delete mode 100644 tests/test_run_grid_search.py diff --git a/run_grid_search.py b/run_grid_search.py index b3336417..4987239c 100755 --- a/run_grid_search.py +++ b/run_grid_search.py @@ -585,20 +585,17 @@ def save_results( merged_runs = [asdict(r) for r in results] for existing_run in existing_runs: - normalized_run = existing_run.copy() - if "model" in normalized_run: - normalized_run.setdefault("model_name", normalized_run.pop("model")) key = ( - normalized_run.get("protein"), - normalized_run.get("model_name"), - normalized_run.get("method"), - normalized_run.get("scaler"), - normalized_run.get("ensemble_size"), - normalized_run.get("gradient_weight"), - normalized_run.get("gd_steps"), + existing_run.get("protein"), + existing_run.get("model_name"), + existing_run.get("method"), + existing_run.get("scaler"), + existing_run.get("ensemble_size"), + existing_run.get("gradient_weight"), + existing_run.get("gd_steps"), ) if key not in new_run_keys: - merged_runs.append(normalized_run) + merged_runs.append(existing_run) output = { "config": asdict(config), diff --git a/src/sampleworks/utils/guidance_script_arguments.py b/src/sampleworks/utils/guidance_script_arguments.py index 66a959de..27f68457 100644 --- a/src/sampleworks/utils/guidance_script_arguments.py +++ b/src/sampleworks/utils/guidance_script_arguments.py @@ -286,8 +286,8 @@ def from_cli( model_name = model_name or pre_args.model_name guidance_type = guidance_type or pre_args.guidance_type - assert model_name is not None - assert guidance_type is not None + if model_name is None or guidance_type is None: + raise RuntimeError("CLI parsing did not resolve a model name and guidance type") # -- full parser ----------------------------------------------------- parser = argparse.ArgumentParser( @@ -415,13 +415,6 @@ def as_dict(self) -> dict[str, Any]: output["log_path"] = _remap_container_path(str(self.log_path)) return output - def __setstate__(self, state: dict[str, Any]) -> None: - """Restore state while migrating legacy pickles from ``model``.""" - migrated = state.copy() - if "model" in migrated: - migrated.setdefault("model_name", migrated.pop("model")) - self.__dict__.update(migrated) - def add_generic_args(parser: argparse.ArgumentParser | GuidanceConfig): """Add CLI arguments shared by all models and guidance methods.""" @@ -706,10 +699,3 @@ def as_dict(self) -> dict[str, Any]: output["output_dir"] = _remap_container_path(str(self.output_dir)) output["log_path"] = _remap_container_path(str(self.log_path)) return output - - def __setstate__(self, state: dict[str, Any]) -> None: - """Restore state while migrating legacy pickles from ``model``.""" - migrated = state.copy() - if "model" in migrated: - migrated.setdefault("model_name", migrated.pop("model")) - self.__dict__.update(migrated) diff --git a/tests/test_run_grid_search.py b/tests/test_run_grid_search.py deleted file mode 100644 index 2eda9fee..00000000 --- a/tests/test_run_grid_search.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Tests for grid-search result serialization.""" - -import json - -from run_grid_search import GridSearchConfig, save_results -from sampleworks.utils.guidance_script_arguments import JobResult - - -def test_save_results_normalizes_legacy_model_key(tmp_path) -> None: - """Existing result records migrate to ``model_name`` without duplicates.""" - legacy_run = { - "protein": "1abc", - "model": "boltz2", - "method": None, - "scaler": "pure_guidance", - "ensemble_size": 1, - "gradient_weight": 0.1, - "gd_steps": 1, - "status": "failed", - } - (tmp_path / "results.json").write_text(json.dumps({"runs": [legacy_run]})) - result = JobResult( - protein="1abc", - model_name="boltz2", - method=None, - scaler="pure_guidance", - ensemble_size=1, - gradient_weight=0.1, - gd_steps=1, - status="success", - exit_code=0, - runtime_seconds=1.0, - started_at="2026-07-13T00:00:00", - finished_at="2026-07-13T00:00:01", - log_path="run.log", - output_dir="output", - ) - config = GridSearchConfig( - model_name="boltz2", - scalers=["pure_guidance"], - ensemble_sizes=[1], - gradient_weights=[0.1], - gd_steps=[1], - method="", - proteins_file="proteins.csv", - output_dir=str(tmp_path), - ) - - save_results([result], config, str(tmp_path), total_time=1.0) - - saved = json.loads((tmp_path / "results.json").read_text()) - assert len(saved["runs"]) == 1 - assert saved["runs"][0]["model_name"] == "boltz2" - assert "model" not in saved["runs"][0] diff --git a/tests/utils/test_guidance_script_arguments.py b/tests/utils/test_guidance_script_arguments.py index 0b2f14e3..7b3c9258 100644 --- a/tests/utils/test_guidance_script_arguments.py +++ b/tests/utils/test_guidance_script_arguments.py @@ -2,7 +2,6 @@ from __future__ import annotations -import pickle from argparse import Namespace from pathlib import Path from unittest.mock import patch @@ -14,7 +13,6 @@ get_checkpoint, GuidanceConfig, JobConfig, - JobResult, validate_model_checkpoint, ) @@ -330,49 +328,3 @@ def test_as_dict_remaps_with_catchall_host_dir(monkeypatch): assert d["output_dir"] == "/mnt/storage/results/1abc" assert d["log_path"] == "/mnt/storage/results/1abc/run.log" assert d["protein"] == "1abc" - - -def test_guidance_config_migrates_legacy_model_pickle() -> None: - """Old job queues restore ``model`` state as ``model_name``.""" - config = GuidanceConfig( - protein="1abc", - structure="structure.cif", - density="density.ccp4", - model_name=StructurePredictor.BOLTZ_2, - guidance_type=GuidanceType.PURE_GUIDANCE, - log_path="run.log", - ) - config.__dict__["model"] = config.__dict__.pop("model_name") - - restored = pickle.loads(pickle.dumps(config)) - - assert restored.model_name == StructurePredictor.BOLTZ_2 - assert "model" not in restored.__dict__ - assert "model" not in restored.as_dict() - - -def test_job_result_migrates_legacy_model_pickle() -> None: - """Old result pickles restore ``model`` state as ``model_name``.""" - result = JobResult( - protein="1abc", - model_name="boltz2", - method=None, - scaler="pure_guidance", - ensemble_size=1, - gradient_weight=0.1, - gd_steps=1, - status="success", - exit_code=0, - runtime_seconds=1.0, - started_at="2026-07-13T00:00:00", - finished_at="2026-07-13T00:00:01", - log_path="run.log", - output_dir="output", - ) - result.__dict__["model"] = result.__dict__.pop("model_name") - - restored = pickle.loads(pickle.dumps(result)) - - assert restored.model_name == "boltz2" - assert "model" not in restored.__dict__ - assert "model" not in restored.as_dict()