Sweeps (pytorch#75)

* add zip
* add e2e test, add types
* add test_titan_sweep.sh  back, used as default command in script

Author: ganeshk-gc

scripts/titan_sweep.py

@@ -4,30 +4,49 @@
 #
 """
 Parameter Sweep Script for W&B style configs
+Usage: python titan_sweep.py <json_config_file_or_json_string>
+
 Config format:
  {
     "command": "./submit.sh", # Executable to run
+    "method": "grid", # Optional: "grid" (default) or "zip"
     "parameters": { # Single parameter set
         "param.name": {"values": [value1, value2]},
         "another.param": {"values": [value3, value4]}
     }
  }
 
- Or
+This will generate the following combinations for grid method:
+    ./submit.sh --param.name value1 --another.param value3
+    ./submit.sh --param.name value1 --another.param value4
+    ./submit.sh --param.name value2 --another.param value3
+    ./submit.sh --param.name value2 --another.param value4
+
+And for zip method:
+    ./submit.sh --param.name value1 --another.param value3
+    ./submit.sh --param.name value2 --another.param value4
+
+
+Or specify config with multiple parameter sets:
 
  {
     "command": "./submit.sh",
+    "method": "zip", # Optional: "grid" (default) or "zip"
     "parameters": [ # Multiple parameter sets
         {
             "param.name": {"values": [value1, value2]},
             "another.param": {"values": [value3, value4]}
         },
         {
+            "project.name": {"values": ["my_project"]},
             "param.name": {"values": [value5, value6]},
             "different.param": {"values": [value7]}
         }
     ]
  }
+
+Special handling for single-value parameters:
+If any parameter has a single value, it is added to all combinations in both grid and zip modes.
 """
 
 import json
@@ -36,9 +55,29 @@
 import subprocess
 import sys
 from pathlib import Path
+from typing import Dict, List, Union, Iterator, Tuple, Any, Optional, Literal, TypedDict
+
+SearchMethodOptions = Literal["grid", "zip"]
+
+# Type aliases for better readability
+ParamConfig = Dict[str, Dict[Literal["values"], List[Any]]]
+CommandList = List[Tuple[str, Union[str, int, float, bool]]]
+
 
-DEFAULT_CONFIG = {
+class SweepConfig(TypedDict):
+    command: str
+    method: SearchMethodOptions
+    parameters: List[ParamConfig]
+
+
+class ParamConfig(TypedDict):
+    name: str
+    values: List[str | float | int | bool]
+
+
+DEFAULT_CONFIG: Dict[str, Union[str, List[ParamConfig]]] = {
     "command": "tests/test_titan_sweep.sh",
+    "method": "grid",
     "parameters": [
         {
             "training.batch_size": {"values": [4, 8]},
@@ -55,7 +94,7 @@
 }
 
 
-def load_config(config_source=None):
+def load_config(config_source: Optional[str] = None) -> Dict[str, Any]:
     if not config_source:
         return DEFAULT_CONFIG
     try:
@@ -68,9 +107,20 @@ def load_config(config_source=None):
         sys.exit(1)
 
 
-def validate_parameters(params):
+def validate_parameters(params: ParamConfig, method: str = "grid") -> None:
+    """
+    Validate a single parameter configuration.
+
+    Args:
+        params: Parameter configuration dictionary
+        method: Combination method ("grid" or "zip")
+
+    Raises:
+        ValueError: If parameters are invalid
+    """
     if not isinstance(params, dict):
         raise ValueError("Parameters must be a dictionary")
+
     for param, param_config in params.items():
         if not isinstance(param, str) or "." not in param:
             raise ValueError(f"Invalid parameter name: {param}")
@@ -79,37 +129,106 @@ def validate_parameters(params):
         if not isinstance(param_config["values"], list):
             raise ValueError(f"Values must be a list for parameter: {param}")
 
+    if method == "zip":
+        # Check length consistency excluding single-value parameters
+        multi_value_params = {k: v for k, v in params.items() if len(v["values"]) > 1}
+        if multi_value_params:
+            lengths = [len(param["values"]) for param in multi_value_params.values()]
+            if not all(length == lengths[0] for length in lengths):
+                raise ValueError(
+                    f"For zip method, all multi-value parameters "
+                    f"must have {lengths[0]} values each"
+                )
 
-def validate_config(config):
+
+def validate_config(config: Dict[str, Any]) -> None:
+    """
+    Validate the complete configuration.
+
+    Args:
+        config: Complete configuration dictionary
+
+    Raises:
+        ValueError: If configuration is invalid
+    """
     if "command" not in config:
         raise ValueError("Missing 'command' in config")
     if not isinstance(config["command"], str):
         raise ValueError("'command' must be a string")
     if "parameters" not in config:
         raise ValueError("Missing 'parameters' in config")
 
-    # Handle both single dict and list of dicts
+    method = config.get("method", "grid")
+    if method not in ["grid", "zip"]:
+        raise ValueError("'method' must be either 'grid' or 'zip'")
+
     param_sets = config["parameters"]
     if isinstance(param_sets, dict):
-        validate_parameters(param_sets)
+        validate_parameters(param_sets, method)
     elif isinstance(param_sets, list):
         for param_set in param_sets:
-            validate_parameters(param_set)
+            validate_parameters(param_set, method)
     else:
         raise ValueError("'parameters' must be either dict or list of dicts")
 
 
-def generate_command_combinations(parameters):
-    param_names = list(parameters.keys())
-    param_values = [parameters[param]["values"] for param in param_names]
-    for values in product(*param_values):
-        yield [(param_names[i], val) for i, val in enumerate(values)]
+def generate_command_combinations(
+    parameters: ParamConfig, method: SearchMethodOptions = "grid"
+) -> Iterator[CommandList]:
+    """
+    Generate parameter combinations based on the specified method.
+
+    Args:
+        parameters: Parameter configuration dictionary
+        method: Combination method ("grid" or "zip")
+
+    Yields:
+        List of parameter-value tuples for each combination
+    """
+    # Get single-value parameters first
+    single_values = [
+        (k, v["values"][0]) for k, v in parameters.items() if len(v["values"]) == 1
+    ]
+
+    # Get multi-value parameters
+    multi_value_params = {k: v for k, v in parameters.items() if len(v["values"]) > 1}
 
+    if method == "grid":
+        if not multi_value_params:
+            yield single_values
+            return
 
-def run_command(command, command_list):
+        param_names = list(multi_value_params.keys())
+        param_values = [multi_value_params[param]["values"] for param in param_names]
+        for values in product(*param_values):
+            yield single_values + [(param_names[i], val) for i, val in enumerate(values)]
+    elif method == "zip":
+        if not multi_value_params:
+            yield single_values
+            return
+
+        names = list(multi_value_params.keys())
+        for values in zip(*(v["values"] for v in multi_value_params.values())):
+            yield single_values + list(zip(names, values))
+    else:
+        raise ValueError(f"Invalid method: {method}")
+
+
+def run_command(command: str, command_list: CommandList) -> None:
+    """
+    Execute the command with the given parameters.
+
+    Args:
+        command: Base command to execute
+        command_list: List of parameter-value tuples to append to command
+
+    Raises:
+        subprocess.CalledProcessError: If command execution fails
+    """
     cmd = command.split()
     if isinstance(cmd, str):
         cmd = [command]
+
     for param, value in command_list:
         if isinstance(value, bool):
             if value:
@@ -128,7 +247,7 @@ def run_command(command, command_list):
         print(f"Error: {e}")
 
 
-def main():
+def main() -> None:
     parser = argparse.ArgumentParser(description="Parameter sweep script")
     group = parser.add_mutually_exclusive_group()
     group.add_argument("config_file", nargs="?", help="Config file path")
@@ -141,14 +260,14 @@ def main():
     try:
         validate_config(config)
         command = config["command"]
+        method = config.get("method", "grid")
 
-        # Handle both single dict and list of dicts
         param_sets = config["parameters"]
         if isinstance(param_sets, dict):
             param_sets = [param_sets]
 
         for param_set in param_sets:
-            for combo in generate_command_combinations(param_set):
+            for combo in generate_command_combinations(param_set, method):
                 run_command(command, combo)
     except Exception as e:
         print(f"Error: {e}")

tests/test_titan_sweep.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2025 Graphcore Ltd. All rights reserved.
+#
+from typing import List, Dict, Any
+import json
+import tempfile
+import subprocess
+import sys
+import os
+import pytest
+
+
+@pytest.fixture
+def echo_script(tmp_path) -> str:
+    """
+    Create a temporary shell script that echoes its arguments.
+    Args:
+        tmp_path: Pytest-provided temporary directory path
+    Returns:
+        Path to the temporary shell script
+    """
+    script_path = tmp_path / "echo_script.sh"
+    script_content = """#!/bin/bash
+echo "$@"
+"""
+    script_path.write_text(script_content)
+    script_path.chmod(0o755)
+    return str(script_path)
+
+
+def run_sweep_and_capture(config_dict: Dict[str, Any]) -> List[str]:
+    """Run the sweep script with given config and return output lines."""
+    with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".json") as tf:
+        json.dump(config_dict, tf)
+        config_path = tf.name
+    try:
+        result = subprocess.run(
+            [sys.executable, "scripts/titan_sweep.py", config_path],
+            capture_output=True,
+            text=True,
+        )
+        return result.stdout.strip().split("\n")
+    finally:
+        os.unlink(config_path)
+
+
+def test_grid_sweep(echo_script: str) -> None:
+    """Test grid method parameter sweep."""
+    config = {
+        "command": echo_script,
+        "method": "grid",
+        "parameters": {
+            "wandb.project": {"values": ["my_project"]},
+            "training.steps": {"values": [1000, 2000]},
+            "training.batch_size": {"values": [16, 32, 64]},
+        },
+    }
+    outputs = run_sweep_and_capture(config)
+    expected_outputs = [
+        "--wandb.project my_project --training.steps 1000 --training.batch_size 16",
+        "--wandb.project my_project --training.steps 1000 --training.batch_size 32",
+        "--wandb.project my_project --training.steps 1000 --training.batch_size 64",
+        "--wandb.project my_project --training.steps 2000 --training.batch_size 16",
+        "--wandb.project my_project --training.steps 2000 --training.batch_size 32",
+        "--wandb.project my_project --training.steps 2000 --training.batch_size 64",
+    ]
+    assert sorted(outputs) == sorted(expected_outputs)
+
+
+def test_zip_sweep(echo_script: str) -> None:
+    """Test zip method parameter sweep with a single-value parameter."""
+    config = {
+        "command": echo_script,
+        "method": "zip",
+        "parameters": {
+            "wandb.project": {"values": ["my_project"]},
+            "model.type": {"values": ["a", "b"]},
+            "learning.rate": {"values": [0.1, 0.2]},
+        },
+    }
+    outputs = run_sweep_and_capture(config)
+    expected_outputs = [
+        "--wandb.project my_project --model.type a --learning.rate 0.1",
+        "--wandb.project my_project --model.type b --learning.rate 0.2",
+    ]
+    assert sorted(outputs) == sorted(expected_outputs)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])