better path handling

Author: mivanit

spd/clustering/clustering_pipeline.py

@@ -9,6 +9,43 @@
 
 from spd.clustering.merge_run_config import RunConfig
 
+PIPELINE_PATHS: dict[str, str] = {
+    "run_record_path": "run_record.json",
+    "histories_dir": "merge_histories",
+    "dataset_dir": "dataset",
+    "ensemble_dir": "ensemble",
+    "distances_dir": "distances",
+}
+
+
+class PipelinePaths:
+    def __init__(self, config: RunConfig) -> None:
+        self.config: RunConfig = config
+
+    @property
+    def run_path(self) -> Path:
+        return self.config.base_path / self.config.config_identifier
+
+    @property
+    def run_record_path(self) -> Path:
+        return self.run_path / PIPELINE_PATHS["run_record_path"]
+
+    @property
+    def histories_dir(self) -> Path:
+        return self.run_path / PIPELINE_PATHS["histories_dir"]
+
+    @property
+    def dataset_dir(self) -> Path:
+        return self.run_path / PIPELINE_PATHS["dataset_dir"]
+
+    @property
+    def ensemble_dir(self) -> Path:
+        return self.run_path / PIPELINE_PATHS["ensemble_dir"]
+
+    @property
+    def distances_dir(self) -> Path:
+        return self.run_path / PIPELINE_PATHS["distances_dir"]
+
 
 def main(config: RunConfig) -> None:
     from spd.clustering.math.merge_distances import (
@@ -21,39 +58,33 @@ def main(config: RunConfig) -> None:
     from spd.clustering.s3_normalize_histories import normalize_and_save
     from spd.clustering.s4_compute_distances import create_clustering_report
 
-    # TODO: factor these out into dataclass or something
-    run_path: Path = config.base_path / config.config_identifier
-    run_record_path: Path = run_path / "run_record.json"
-    histories_dir: Path = run_path / "merge_histories"
-    dataset_dir: Path = run_path / "dataset"
-    ensemble_dir: Path = run_path / "ensemble"
-    distances_dir: Path = run_path / "distances"
+    paths: PipelinePaths = PipelinePaths(config=config)
 
-    print(f"Run record saved to {run_record_path}")
-    run_record_path.write_text(config.model_dump_json(indent=2))
+    print(f"Run record saved to {paths.run_record_path}")
+    paths.run_record_path.write_text(config.model_dump_json(indent=2))
 
     print(f"Splitting dataset into {config.n_batches} batches...")
-    data_files: list[Path] = split_and_save_dataset(config=config, output_dir=dataset_dir)
+    data_files: list[Path] = split_and_save_dataset(config=config, output_dir=paths.dataset_dir)
 
     print(
         f"Processing {len(data_files)} batches with {config.workers_per_device} workers per device..."
     )
     results: list[ClusteringResult] = process_batches_parallel(
         data_files=data_files,
         config=config,
-        output_dir=histories_dir,
+        output_dir=paths.histories_dir,
         workers_per_device=config.workers_per_device,
         devices=config.devices,
     )
 
     normalized_merge_array: MergesArray = normalize_and_save(
         history_paths=[r.history_save_path for r in results],
-        output_dir=ensemble_dir,
+        output_dir=paths.ensemble_dir,
     )
 
     distances: DistancesArray = compute_and_save_distances(
         normalized_merge_array=normalized_merge_array,
-        output_dir=distances_dir,
+        output_dir=paths.distances_dir,
     )
 
     wandb_urls: list[str] = [r.wandb_url for r in results if r.wandb_url]  # Gross - clean up

spd/clustering/pipeline/__init__.py

@@ -39,7 +39,8 @@ def test_merge_with_range_sampler(self):
         # Check results
         assert history is not None
         assert len(history.merges.k_groups) > 0
-        assert history.merges.k_groups[0].item() == n_components
+        # First entry is after first merge, so should be n_components - 1
+        assert history.merges.k_groups[0].item() == n_components - 1
         # After iterations, should have fewer groups (merges reduce count)
         # Exact count depends on early stopping conditions
         assert history.merges.k_groups[-1].item() < n_components
@@ -75,7 +76,8 @@ def test_merge_with_mcmc_sampler(self):
         # Check results
         assert history is not None
         assert len(history.merges.k_groups) > 0
-        assert history.merges.k_groups[0].item() == n_components
+        # First entry is after first merge, so should be n_components - 1
+        assert history.merges.k_groups[0].item() == n_components - 1
         # Should have fewer groups after iterations
         assert history.merges.k_groups[-1].item() < n_components
         assert history.merges.k_groups[-1].item() >= 2