Add join sweep method (#173)

Allows adding extra agents to a sweep.

Author: alan-cooney

sparse_autoencoder/train/join_sweep.py

@@ -0,0 +1,23 @@
+"""Join an existing Weights and Biases sweep, as a new agent."""
+import argparse
+
+from sparse_autoencoder.train.sweep import sweep
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments.
+
+    Returns:
+        argparse.Namespace: Parsed command line arguments.
+    """
+    parser = argparse.ArgumentParser(description="Join an existing W&B sweep.")
+    parser.add_argument(
+        "--id", type=str, default=None, help="Sweep ID for the existing sweep.", required=True
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+
+    sweep(sweep_id=args.id)

sparse_autoencoder/train/sweep.py

@@ -347,9 +347,28 @@ def train() -> None:
         sys.exit(1)
 
 
-def sweep(sweep_config: SweepConfig) -> None:
-    """Main function to run the training pipeline with wandb hyperparameter sweep."""
-    sweep_id = wandb.sweep(sweep_config.to_dict(), project="sparse-autoencoder")
+def sweep(sweep_config: SweepConfig | None = None, sweep_id: str | None = None) -> None:
+    """Run the training pipeline with wandb hyperparameter sweep.
+
+    Warning:
+        Either sweep_config or sweep_id must be specified, but not both.
+
+    Args:
+        sweep_config: The sweep configuration.
+        sweep_id: The sweep id for an existing sweep.
+
+    Raises:
+        ValueError: If neither sweep_config nor sweep_id is specified.
+    """
+    if sweep_id is not None:
+        wandb.agent(sweep_id, train, project="sparse-autoencoder")
+
+    elif sweep_config is not None:
+        sweep_id = wandb.sweep(sweep_config.to_dict(), project="sparse-autoencoder")
+        wandb.agent(sweep_id, train)
+
+    else:
+        error_message = "Either sweep_config or sweep_id must be specified."
+        raise ValueError(error_message)
 
-    wandb.agent(sweep_id, train)
     wandb.finish()