diff --git a/src/baskerville/helpers/stage_cache.py b/src/baskerville/helpers/stage_cache.py index 7f0334c..51e8e3a 100644 --- a/src/baskerville/helpers/stage_cache.py +++ b/src/baskerville/helpers/stage_cache.py @@ -137,6 +137,11 @@ def _iter_dir_files(root: str, *, filename: str | None = None): yield rel, full +def hash_text(text: str) -> str: + """SHA-256 hex of a string, e.g. a job command keying a run id.""" + return hashlib.sha256(text.encode()).hexdigest() + + def hash_dir(path: str, *, filename: str | None = None) -> str: """SHA256 of the sorted ``(relpath, file_sha256)`` manifest of a directory.""" manifest = hashlib.sha256() diff --git a/src/baskerville/scripts/hound_eval_folds.py b/src/baskerville/scripts/hound_eval_folds.py index c2ad7b3..ca3dff0 100755 --- a/src/baskerville/scripts/hound_eval_folds.py +++ b/src/baskerville/scripts/hound_eval_folds.py @@ -291,7 +291,7 @@ def eval_folds(args): data_names = [os.path.basename(d.rstrip("/")) for d in args.data_dirs] container_data_dirs = [f"{args.gcp_data_local}/{n}" for n in data_names] if not args.gcp_output_dir: - # Deterministic run id = hash(models, params, data URIs). Same inputs + # Deterministic run id = hash(models, params, data URIs, options). Same inputs # → same GCS location → per-shard resume (skip finished acc.txt). # Marker mode has no staged models_sha; key on the GCS models dir. data_uris = sorted(f"{args.gcp_data_dir}/{n}" for n in data_names) @@ -299,8 +299,13 @@ def eval_folds(args): models_key = ( models_sha or hashlib.sha256(gcs_models_dir.encode()).hexdigest() ) + # output-changing options too, so changing them starts fresh instead of + # resuming; spec's are keyed even without --spec so adding it resumes + opts_sha = stage_cache.hash_text( + eval_options(args, container_targets_gene) + spec_options(args) + ) run_id = stage_cache.build_run_id( - models_key, params_sha, data_uri_sha, deterministic=True + models_key, params_sha, data_uri_sha, opts_sha, deterministic=True ) args.gcp_output_dir = f"{stage_cache.output_prefix()}/eval/{run_id}" print("=" * 72) @@ -320,6 +325,7 @@ def eval_folds(args): else: params_file = os.path.abspath(args.params_file) data_dirs = [os.path.abspath(data_dir) for data_dir in args.data_dirs] + targets_gene = container_targets_gene if gcp_backend else args.targets_gene_file ####################################################### # prep work @@ -457,26 +463,8 @@ def eval_folds(args): cmd = cmd_base cmd += f" --dataset {di}" cmd += f" -o {eval_fold_dir}" - if args.rank_corr: - cmd += " --rank" - if args.rc: - cmd += " --rc" - if args.save: - cmd += " --save" - if args.shifts: - cmd += f" --shifts {args.shifts}" - if args.aggregate_genes: - cmd += " --aggregate_genes" - targets_gene = ( - container_targets_gene - if gcp_backend - else args.targets_gene_file - ) - if targets_gene: - cmd += f" --targets_gene_file {targets_gene}" cmd += f" --split fold{ei}" - if args.rank_corr or args.save: - cmd += f" --step {args.step}" + cmd += eval_options(args, targets_gene) cmd += f" {params_file}" cmd += f" {model_file}" cmd += f" {data_dirs[di]}" @@ -565,8 +553,8 @@ def eval_folds(args): cmd = spec_cmd_base cmd += f" --dataset {di}" cmd += f" -o {out_dir}" - cmd += f" --step {args.step}" cmd += f" --split fold{test_fold}" + cmd += spec_options(args) cmd += f" --ncpus {num_cpu}" if args.band is not None: cmd += f" --band {args.band}" @@ -576,10 +564,6 @@ def eval_folds(args): cmd += " --ram" if args.scratch_dir is not None: cmd += f" --scratch_dir {args.scratch_dir}" - if args.rc: - cmd += " --rc" - if args.shifts: - cmd += f" --shifts {args.shifts}" cmd += f" {params_file}" cmd += f" {model_file}" cmd += f" {data_dirs[di]}" @@ -663,6 +647,36 @@ def eval_folds(args): # their original private names so existing call sites and tests are unchanged. +def eval_options(args, targets_gene): + """hound_eval flags that change its output (not per-job paths/splits).""" + opts = "" + if args.rank_corr: + opts += " --rank" + if args.rc: + opts += " --rc" + if args.save: + opts += " --save" + if args.shifts: + opts += f" --shifts {args.shifts}" + if args.aggregate_genes: + opts += " --aggregate_genes" + if targets_gene: + opts += f" --targets_gene_file {targets_gene}" + if args.rank_corr or args.save: + opts += f" --step {args.step}" + return opts + + +def spec_options(args): + """hound_eval_spec flags that change its output (not memory/IO knobs).""" + opts = f" --step {args.step}" + if args.rc: + opts += " --rc" + if args.shifts: + opts += f" --shifts {args.shifts}" + return opts + + def _link_test_metrics(eval_dir, test_fold): """Symlink test -> fold{test_fold} and test.out -> fold{test_fold}.out. diff --git a/src/baskerville/scripts/hound_eval_genes_folds.py b/src/baskerville/scripts/hound_eval_genes_folds.py index fa21140..d0a0a0b 100644 --- a/src/baskerville/scripts/hound_eval_genes_folds.py +++ b/src/baskerville/scripts/hound_eval_genes_folds.py @@ -275,8 +275,16 @@ def eval_genes_folds(args): models_key = ( models_sha or hashlib.sha256(gcs_models_dir.encode()).hexdigest() ) + # options too (and --valid, which changes the split written to the + # same evalg/ dir), so changed options start fresh instead of resuming + opts = evalg_options(args, container_targets) + f" valid={args.valid}" run_id = stage_cache.build_run_id( - models_key, params_sha, gtf_sha, data_uri_sha, deterministic=True + models_key, + params_sha, + gtf_sha, + data_uri_sha, + stage_cache.hash_text(opts), + deterministic=True, ) args.gcp_output_dir = f"{stage_cache.output_prefix()}/evalg/{run_id}" print("=" * 72) @@ -438,19 +446,7 @@ def eval_genes_folds(args): cmd += f" --head {di}" cmd += f" -o {eval_dir}" cmd += f" --split fold{split_fold}" - cmd += f" --seq_step {args.seq_step}" - if args.pseudo_qtl is not None: - cmd += f" --pseudo_qtl {args.pseudo_qtl:.2f}" - if args.rc: - cmd += " --rc" - if args.save_span: - cmd += " --save_span" - if args.shifts: - cmd += f" --shifts {args.shifts}" - if args.span: - cmd += " --span" - if targets_file: - cmd += f" -t {targets_file}" + cmd += evalg_options(args, targets_file) cmd += f" {params_file}" cmd += f" {model_file}" cmd += f" {data_dirs[di]}" @@ -536,6 +532,24 @@ def eval_genes_folds(args): print(f"[gcs] full run at {args.gcp_output_dir}") +def evalg_options(args, targets_file): + """hound_eval_genes flags that change its output (not paths/splits).""" + opts = f" --seq_step {args.seq_step}" + if args.pseudo_qtl is not None: + opts += f" --pseudo_qtl {args.pseudo_qtl:.2f}" + if args.rc: + opts += " --rc" + if args.save_span: + opts += " --save_span" + if args.shifts: + opts += f" --shifts {args.shifts}" + if args.span: + opts += " --span" + if targets_file: + opts += f" -t {targets_file}" + return opts + + ################################################################################ # __main__ ################################################################################ diff --git a/src/baskerville/scripts/hound_ism_bed_folds.py b/src/baskerville/scripts/hound_ism_bed_folds.py index 43caab0..85027c4 100755 --- a/src/baskerville/scripts/hound_ism_bed_folds.py +++ b/src/baskerville/scripts/hound_ism_bed_folds.py @@ -305,7 +305,11 @@ def ism_bed_folds(args): models_key = ( models_sha or hashlib.sha256(gcs_models_dir.encode()).hexdigest() ) - run_id = stage_cache.build_run_id(bed_sha, models_key, deterministic=True) + # job command too, so changed options start fresh instead of resuming + cmd_sha = stage_cache.hash_text(build_ism_bed_cmd(args, "MODEL", "OUT")) + run_id = stage_cache.build_run_id( + bed_sha, models_key, cmd_sha, deterministic=True + ) args.gcp_output_dir = f"{stage_cache.output_prefix()}/ism_bed/{run_id}" print("=" * 72) print(f"[gcp] run output dir: {args.gcp_output_dir}") diff --git a/src/baskerville/scripts/hound_ism_snp_folds.py b/src/baskerville/scripts/hound_ism_snp_folds.py index 972a2ae..4efe71f 100755 --- a/src/baskerville/scripts/hound_ism_snp_folds.py +++ b/src/baskerville/scripts/hound_ism_snp_folds.py @@ -300,7 +300,11 @@ def ism_snp_folds(args): models_key = ( models_sha or hashlib.sha256(gcs_models_dir.encode()).hexdigest() ) - run_id = stage_cache.build_run_id(vcf_sha, models_key, deterministic=True) + # job command too, so changed options start fresh instead of resuming + cmd_sha = stage_cache.hash_text(build_ism_snp_cmd(args, "MODEL", "OUT")) + run_id = stage_cache.build_run_id( + vcf_sha, models_key, cmd_sha, deterministic=True + ) args.gcp_output_dir = f"{stage_cache.output_prefix()}/ism_snp/{run_id}" print("=" * 72) print(f"[gcp] run output dir: {args.gcp_output_dir}") diff --git a/src/baskerville/scripts/hound_snp_folds.py b/src/baskerville/scripts/hound_snp_folds.py index 80dcfca..8e1e9e2 100755 --- a/src/baskerville/scripts/hound_snp_folds.py +++ b/src/baskerville/scripts/hound_snp_folds.py @@ -356,9 +356,15 @@ def snp_folds(args): if getattr(args, "genes_gtf", None): _, args.genes_gtf = stage_cache.stage_file(args.genes_gtf, "gtf") if not args.gcp_output_dir: - # Deterministic run_id so a re-run with the same VCF+models lands - # in the same output prefix, enabling per-shard resume. - run_id = stage_cache.build_run_id(vcf_sha, models_sha, deterministic=True) + # Deterministic run_id so a re-run with the same VCF, models and job + # command (options, job size) lands in the same output prefix, enabling + # per-shard resume; any change starts fresh instead of reusing shards. + cmd_sha = stage_cache.hash_text( + build_snp_cmd(args, "MODEL", *job_bounds[0], "OUT", "FOLD") + ) + run_id = stage_cache.build_run_id( + vcf_sha, models_sha, cmd_sha, deterministic=True + ) args.gcp_output_dir = f"{stage_cache.output_prefix()}/snp/{run_id}" args.gcp_data_dir = stage_cache.cache_prefix() args.gcp_data_local = stage_cache.CONTAINER_CACHE_MOUNT diff --git a/src/tests/test_eval_folds_gcp.py b/src/tests/test_eval_folds_gcp.py index a05de1c..881ebb8 100644 --- a/src/tests/test_eval_folds_gcp.py +++ b/src/tests/test_eval_folds_gcp.py @@ -390,3 +390,21 @@ def test_run_marker_roundtrip(tmp_path): assert marker["gcp_zone"] is None # absent marker → None assert hef.stage_cache.read_run_marker(str(tmp_path / "nope")) is None + + +def _gcp_output_dir(tmp_path, monkeypatch, **kwargs): + out_dir = _make_models_dir(tmp_path, num_folds=2, test_fold=0, valid_fold=1) + captured = _patch_gcp(monkeypatch, num_folds=2) + hef.eval_folds(MockArgs(out_dir=out_dir, test_only=True, **kwargs)) + return captured[0].kwargs["output_dir_gcs"] + + +def test_gcp_output_dir_keyed_on_options(tmp_path, monkeypatch): + """Output-changing options start fresh; memory/IO knobs still resume.""" + base = _gcp_output_dir(tmp_path, monkeypatch) + assert _gcp_output_dir(tmp_path, monkeypatch) == base + assert _gcp_output_dir(tmp_path, monkeypatch, rc=True) != base + assert _gcp_output_dir(tmp_path, monkeypatch, shifts="0,1") != base + assert _gcp_output_dir(tmp_path, monkeypatch, band=8, ram=True) == base + # adding --spec to a finished eval run resumes its eval jobs + assert _gcp_output_dir(tmp_path, monkeypatch, spec=True) == base diff --git a/src/tests/test_eval_genes_folds_gcp.py b/src/tests/test_eval_genes_folds_gcp.py index 0027e8e..7d62e8b 100644 --- a/src/tests/test_eval_genes_folds_gcp.py +++ b/src/tests/test_eval_genes_folds_gcp.py @@ -360,3 +360,18 @@ def test_marker_mode_inherits_gcp_config(tmp_path, monkeypatch): assert args.gcp_image == "img:tag" assert args.gcp_project == "proj" assert args.gcp_region == "us-west1" + + +def _gcp_output_dir(tmp_path, monkeypatch, **kwargs): + out_dir = _make_models_dir(tmp_path, num_folds=4, test_fold=0, valid_fold=3) + captured = _patch_gcp(monkeypatch, num_folds=4) + hegf.eval_genes_folds(MockArgs(out_dir=out_dir, fold_subset=1, **kwargs)) + return captured[0].kwargs["output_dir_gcs"] + + +def test_gcp_output_dir_keyed_on_options(tmp_path, monkeypatch): + """Changed options (or --valid, same evalg/ dir) must not resume another run.""" + base = _gcp_output_dir(tmp_path, monkeypatch) + assert _gcp_output_dir(tmp_path, monkeypatch) == base + assert _gcp_output_dir(tmp_path, monkeypatch, rc=True) != base + assert _gcp_output_dir(tmp_path, monkeypatch, valid=True) != base diff --git a/src/tests/test_snp_folds.py b/src/tests/test_snp_folds.py index 21d9eb5..bb2a513 100644 --- a/src/tests/test_snp_folds.py +++ b/src/tests/test_snp_folds.py @@ -396,3 +396,70 @@ def test_collect_scores_gene_mode_index_remapping(): if __name__ == "__main__": pytest.main([__file__]) + + +def _gcp_output_dir(monkeypatch, tmp_path, script, *options): + """Run a *_folds script with --backend gcp up to its first resume check and + return the output dir it chose. Staging, image lookup, the runner and GCS are + stubbed out.""" + import importlib + import sys + + from baskerville.helpers import stage_cache + + module = importlib.import_module(f"baskerville.scripts.{script}") + model_dir = tmp_path / "models" / "f0c0" / "train" + model_dir.mkdir(parents=True, exist_ok=True) + (model_dir / "model_best.pth").touch() + + class Resumed(Exception): + pass + + def resume_check(gcs_uri, status): + raise Resumed(gcs_uri) + + staged = lambda path, type_, **kw: ("a" * 64, f"/cache/{type_}") + for name, stub in [ + ("resolve_project", lambda project: "project"), + ("resolve_image_arg", lambda *a, **kw: None), + ("announce_gcp_image", lambda *a, **kw: None), + ("make_runner", lambda args, **kw: None), + ]: + monkeypatch.setattr(module, name, stub, raising=False) + monkeypatch.setattr(stage_cache, "read_run_marker", lambda path: None) + monkeypatch.setattr(stage_cache, "cache_prefix", lambda: "gs://cache") + monkeypatch.setattr(stage_cache, "output_prefix", lambda: "gs://out") + monkeypatch.setattr(stage_cache, "stage_file", staged) + monkeypatch.setattr(stage_cache, "stage_dir", staged) + monkeypatch.setattr(stage_cache, "check_progress_h5_gcs", resume_check) + + test_dir = pathlib.Path(__file__).parent + if script == "hound_ism_bed_folds": + data_file = tmp_path / "regions.bed" + data_file.write_text("chrI\t1000\t1100\n") + else: + data_file = test_dir / "data" / "sc3_snps.vcf" + argv = [script, "--backend", "gcp", "-f", f"{test_dir}/data/sc3.fa.gz"] + argv += ["-t", f"{test_dir}/data/targets_sc3_me.txt", *options] + argv += [f"{test_dir}/data/params_sc3.json", str(tmp_path / "models")] + argv += [str(data_file)] + monkeypatch.setattr(sys, "argv", argv) + with pytest.raises(Resumed) as resumed: + module.main() + return str(resumed.value).rsplit("/f0c0/", 1)[0] + + +@pytest.mark.parametrize( + "script,option", + [ + ("hound_snp_folds", ["-j", "2"]), + ("hound_snp_folds", ["-m", "bfloat16"]), + ("hound_ism_snp_folds", ["-m", "bfloat16"]), + ("hound_ism_bed_folds", ["-m", "bfloat16"]), + ], +) +def test_folds_gcp_output_dir_keyed_on_command(monkeypatch, tmp_path, script, option): + """Changing scoring options or job size must not resume another run's shards.""" + base = _gcp_output_dir(monkeypatch, tmp_path, script) + assert _gcp_output_dir(monkeypatch, tmp_path, script) == base + assert _gcp_output_dir(monkeypatch, tmp_path, script, *option) != base