Spaces:
Sleeping
Sleeping
| configfile: "config.yaml" | |
| import os | |
| import mlflow | |
| config["output_dir"] = os.environ.get("OUTPUT_DIR", config["output_dir"]) | |
| config["cache_dir"] = os.environ.get("CACHE_DIR", config["cache_dir"]) | |
| # mlruns.db lives at output_dir level — shared across experiments | |
| config["mlflow_uri"] = f"sqlite:///{config['output_dir']}/mlruns.db" | |
| # all stage outputs and sentinels live under experiment_dir | |
| config["experiment_dir"] = os.path.join(config["output_dir"], config["mlflow_experiment"]) | |
| # eval.out is derived from experiment_dir so it doesn't need to be in config.yaml | |
| config.setdefault("eval", {})["out"] = f"{config['experiment_dir']}/eval_metrics/eval_metrics.parquet" | |
| # Pre-initialize the MLflow DB and experiment here (single-threaded, before | |
| # any parallel jobs start) so workers don't race to CREATE TABLE. | |
| os.makedirs(config["experiment_dir"], exist_ok=True) | |
| mlflow.set_tracking_uri(config["mlflow_uri"]) | |
| mlflow.set_experiment(config["mlflow_experiment"]) | |
| wildcard_constraints: | |
| output_dir = r"[^/]+", | |
| run_key = r"[^/]+", | |
| include: "rules/clean.smk" | |
| include: "rules/primary.smk" | |
| include: "rules/transform.smk" | |
| include: "rules/correlations.smk" | |
| include: "rules/ultrametricity.smk" | |
| include: "rules/pair_figures.smk" | |
| include: "rules/merge.smk" | |
| include: "rules/eval.smk" | |
| def _run_key(run): | |
| """Convert a {model, revision} dict to the output directory name.""" | |
| model = run["model"] | |
| rev = run.get("revision") or None | |
| return f"{model}_{rev}" if rev else model | |
| def _all_targets(): | |
| ed = config["experiment_dir"] | |
| targets = [f"{ed}/done/{_run_key(r)}.transform.done" for r in config["runs"]] | |
| if config.get("merge", {}).get("cross_model", {}).get("enabled", False): | |
| targets.append(f"{ed}/done/cross_model.merge.done") | |
| for cp_model in config.get("merge", {}).get("checkpoints", []): | |
| targets.append(f"{ed}/done/{cp_model}_checkpoints.merge.done") | |
| return targets | |
| rule all: | |
| """Default target: primary + transform (+ cross-model merge) for all configured runs.""" | |
| input: | |
| _all_targets(), | |
| # Convenience targets for optional stages (not in default `all`). | |
| rule correlations_all: | |
| """Optional target: correlations for all configured runs.""" | |
| input: | |
| expand(config["experiment_dir"] + "/done/{run_key}.correlations.done", | |
| run_key=[_run_key(r) for r in config["runs"]]), | |
| rule ultrametricity_all: | |
| """Optional target: ultrametricity diagnostics for all configured runs.""" | |
| input: | |
| expand(config["experiment_dir"] + "/done/{run_key}.ultrametricity.done", | |
| run_key=[_run_key(r) for r in config["runs"]]), | |
| rule pair_figures_all: | |
| """Optional target: pair figures for all configured runs.""" | |
| input: | |
| expand(config["experiment_dir"] + "/done/{run_key}.pair_figures.done", | |
| run_key=[_run_key(r) for r in config["runs"]]), | |
| rule eval_all: | |
| """Optional target: eval + merge for all configured runs.""" | |
| input: | |
| config.get("eval", {}).get("out", "outputs/eval_metrics/eval_metrics.parquet"), | |
| # ── Named family targets ─────────────────────────────────────────────────────── | |
| # Each entry in config["targets"] becomes a snakemake target_<name> rule that | |
| # drives only the runs in that subset: | |
| # snakemake target_gpt2_family -j4 | |
| # snakemake target_pythia_family -j4 | |
| # snakemake target_llama_family -j4 | |
| # snakemake target_all -j8 | |
| # template rule — one instance per entry in config["targets"] | |
| for _target_name, _target_runs in config.get("targets", {}).items(): | |
| rule: | |
| name: f"target_{_target_name}" | |
| input: | |
| expand(config["experiment_dir"] + "/done/{run_key}.transform.done", | |
| run_key=[_run_key(r) for r in _target_runs]), | |
| rule target_all: | |
| """Run primary + transform for all model families (composes all family targets).""" | |
| input: | |
| rules.target_gpt2_family.input, | |
| rules.target_pythia_family.input, | |
| rules.target_llama_family.input, | |
| rules.target_olmo2_family.input, | |
| rules.target_llama32_family.input, | |
| rules.target_gemma2_family.input, | |
| rules.target_smollm2_family.input, | |
| # ── Pythia checkpoint sweep targets ─────────────────────────────────────────── | |
| # Each entry in config["pythia_steps"] generates target_pythia_{shortname}_steps | |
| # covering all PYTHIA_REVISIONS for that model. | |
| # Enable in config.yaml under pythia_steps:, then add the per-step runs to runs:. | |
| # snakemake target_pythia_70m_steps -j4 | |
| from transformer_analysis.model_registry import PYTHIA_REVISIONS | |
| # template rule — one instance per entry in config["pythia_steps"] | |
| for _ps in config.get("pythia_steps", []): | |
| _ps_model = _ps["model"] | |
| _ps_name = _ps["shortname"] | |
| rule: | |
| name: f"target_pythia_{_ps_name}_steps" | |
| input: | |
| expand(config["experiment_dir"] + "/done/{run_key}.transform.done", | |
| run_key=[f"{_ps_model}_{rev}" for rev in PYTHIA_REVISIONS]), | |