| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | import os |
| 5 | |
| 6 | |
| 7 | class ParameterSweep(list["ParameterSweepItem"]): |
| 8 | @classmethod |
| 9 | def read_json(cls, filepath: os.PathLike): |
| 10 | with open(filepath, "rb") as f: |
| 11 | data = json.load(f) |
| 12 | |
| 13 | # Support both list and dict formats |
| 14 | if isinstance(data, dict): |
| 15 | return cls.read_from_dict(data) |
| 16 | |
| 17 | return cls.from_records(data) |
| 18 | |
| 19 | @classmethod |
| 20 | def read_from_dict(cls, data: dict[str, dict[str, object]]): |
| 21 | """Read parameter sweep from a dict format where keys are names. |
| 22 | |
| 23 | Example: |
| 24 | { |
| 25 | "experiment1": {"max_tokens": 100, "temperature": 0.7}, |
| 26 | "experiment2": {"max_tokens": 200, "temperature": 0.9} |
| 27 | } |
| 28 | |
| 29 | """ |
| 30 | records = [{"_benchmark_name": name, **params} for name, params in data.items()] |
| 31 | return cls.from_records(records) |
| 32 | |
| 33 | @classmethod |
| 34 | def from_records(cls, records: list[dict[str, object]]): |
| 35 | if not isinstance(records, list): |
| 36 | raise TypeError( |
| 37 | f"The parameter sweep should be a list of dictionaries, " |
| 38 | f"but found type: {type(records)}" |
| 39 | ) |
| 40 | |
| 41 | # Validate that all _benchmark_name values are unique if provided |
| 42 | names = [r["_benchmark_name"] for r in records if "_benchmark_name" in r] |
| 43 | if names and len(names) != len(set(names)): |
| 44 | duplicates = [name for name in names if names.count(name) > 1] |
| 45 | raise ValueError( |
| 46 | f"Duplicate _benchmark_name values found: {set(duplicates)}. " |
| 47 | f"All _benchmark_name values must be unique." |
| 48 | ) |
| 49 | |
| 50 | return cls(ParameterSweepItem.from_record(record) for record in records) |
| 51 | |
| 52 | |
| 53 | class ParameterSweepItem(dict[str, object]): |
| 54 | @classmethod |
| 55 | def from_record(cls, record: dict[str, object]): |
| 56 | if not isinstance(record, dict): |
| 57 | raise TypeError( |
| 58 | f"Each item in the parameter sweep should be a dictionary, " |
| 59 | f"but found type: {type(record)}" |
| 60 | ) |
| 61 | |
| 62 | return cls(record) |
| 63 | |
| 64 | def __or__(self, other: dict[str, object], /) -> "ParameterSweepItem": # type: ignore[override] |
| 65 | return ParameterSweepItem(super().__or__(other)) |
| 66 | |
| 67 | @property |
| 68 | def name(self) -> str: |
| 69 | """Get the name for this parameter sweep item. |
| 70 | |
| 71 | Returns the '_benchmark_name' field if present, otherwise returns a text |
| 72 | representation of all parameters. |
| 73 | """ |
| 74 | if "_benchmark_name" in self: |
| 75 | return str(self["_benchmark_name"]) |
| 76 | |
| 77 | return self.as_text(sep="-") |
| 78 | |
| 79 | # In JSON, we prefer "_" |
| 80 | def _iter_param_key_candidates(self, param_key: str): |
| 81 | # Inner config arguments are not converted by the CLI |
| 82 | if "." in param_key: |
| 83 | prefix, rest = param_key.split(".", 1) |
| 84 | for prefix_candidate in self._iter_param_key_candidates(prefix): |
| 85 | yield prefix_candidate + "." + rest |
| 86 | |
| 87 | return |
| 88 | |
| 89 | yield param_key |
| 90 | yield param_key.replace("-", "_") |
| 91 | yield param_key.replace("_", "-") |
| 92 | |
| 93 | # In CLI, we prefer "-" |
| 94 | def _iter_cmd_key_candidates(self, param_key: str): |
| 95 | for k in reversed(tuple(self._iter_param_key_candidates(param_key))): |
| 96 | yield "--" + k |
| 97 | |
| 98 | def _normalize_cmd_key(self, param_key: str): |
| 99 | return next(self._iter_cmd_key_candidates(param_key)) |
| 100 | |
| 101 | def has_param(self, param_key: str) -> bool: |
| 102 | return any(k in self for k in self._iter_param_key_candidates(param_key)) |
| 103 | |
| 104 | def _normalize_cmd_kv_pair(self, k: str, v: object) -> list[str]: |
| 105 | """Normalize a key-value pair into command-line arguments. |
| 106 | |
| 107 | Returns a list containing either: |
| 108 | - A single element for boolean flags (e.g., ['--flag'] or ['--flag=true']) |
| 109 | - Two elements for key-value pairs (e.g., ['--key', 'value']) |
| 110 | """ |
| 111 | if isinstance(v, bool): |
| 112 | # For nested params (containing "."), use =true/false syntax |
| 113 | if "." in k: |
| 114 | return [f"{self._normalize_cmd_key(k)}={'true' if v else 'false'}"] |
| 115 | else: |
| 116 | return [self._normalize_cmd_key(k if v else "no-" + k)] |
| 117 | else: |
| 118 | return [self._normalize_cmd_key(k), str(v)] |
| 119 | |
| 120 | def apply_to_cmd(self, cmd: list[str]) -> list[str]: |
| 121 | cmd = list(cmd) |
| 122 | |
| 123 | for k, v in self.items(): |
| 124 | # Skip the '_benchmark_name' field, not a parameter |
| 125 | if k == "_benchmark_name": |
| 126 | continue |
| 127 | |
| 128 | # Serialize dict values as JSON |
| 129 | if isinstance(v, dict): |
| 130 | v = json.dumps(v) |
| 131 | |
| 132 | for k_candidate in self._iter_cmd_key_candidates(k): |
| 133 | try: |
| 134 | k_idx = cmd.index(k_candidate) |
| 135 | |
| 136 | # Replace existing parameter |
| 137 | normalized = self._normalize_cmd_kv_pair(k, v) |
| 138 | if len(normalized) == 1: |
| 139 | # Boolean flag |
| 140 | cmd[k_idx] = normalized[0] |
| 141 | else: |
| 142 | # Key-value pair |
| 143 | cmd[k_idx] = normalized[0] |
| 144 | cmd[k_idx + 1] = normalized[1] |
| 145 | |
| 146 | break |
| 147 | except ValueError: |
| 148 | continue |
| 149 | else: |
| 150 | # Add new parameter |
| 151 | cmd.extend(self._normalize_cmd_kv_pair(k, v)) |
| 152 | |
| 153 | return cmd |
| 154 | |
| 155 | def as_text(self, sep: str = ", ") -> str: |
| 156 | return sep.join(f"{k}={v}" for k, v in self.items() if k != "_benchmark_name") |