Code simplifications for getting and validating user settings.

Author: choijon5

helion/runtime/settings.py

@@ -7,6 +7,7 @@
 import time
 from typing import TYPE_CHECKING
 from typing import Callable
+from typing import Collection
 from typing import Literal
 from typing import Protocol
 from typing import Sequence
@@ -36,6 +37,45 @@ def __call__(
         ) -> BaseAutotuner: ...
 
 
+def _validate_enum_setting(
+    value: object,
+    *,
+    name: str,
+    valid: Collection[str],
+    allow_none: bool = True,
+) -> str | None:
+    """Normalize and validate an enum setting.
+
+    Args:
+        value: The value to normalize and validate
+        name: Name of the setting
+        valid: Collection of valid settings
+        allow_none: If True, None and _NONE_VALUES strings return None. If False, they raise an error.
+    """
+    # String values that should be treated as None
+    _NONE_VALUES = frozenset({"", "0", "false", "none"})
+
+    # Normalize values
+    normalized: str | None
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+    else:
+        normalized = None
+
+    is_none_value = normalized is None or normalized in _NONE_VALUES
+    is_valid = normalized in valid if normalized else False
+
+    # Valid value (none or valid setting)
+    if is_none_value and allow_none:
+        return None
+    if is_valid:
+        return normalized
+
+    # Invalid value, raise error
+    valid_list = "', '".join(sorted(valid))
+    raise ValueError(f"{name} must be one of '{valid_list}', got {value!r}")
+
+
 _tls: _TLS = cast("_TLS", threading.local())
 
 
@@ -108,63 +148,6 @@ def default_autotuner_fn(
     return LocalAutotuneCache(autotuner_cls(bound_kernel, args, **kwargs))  # pyright: ignore[reportArgumentType]
 
 
-def _get_autotune_random_seed() -> int:
-    value = os.environ.get("HELION_AUTOTUNE_RANDOM_SEED")
-    if value is not None:
-        return int(value)
-    return int(time.time() * 1000) % 2**32
-
-
-def _get_autotune_max_generations() -> int | None:
-    value = os.environ.get("HELION_AUTOTUNE_MAX_GENERATIONS")
-    if value is not None:
-        return int(value)
-    return None
-
-
-def _get_autotune_rebenchmark_threshold() -> float | None:
-    value = os.environ.get("HELION_REBENCHMARK_THRESHOLD")
-    if value is not None:
-        return float(value)
-    return None  # Will use effort profile default
-
-
-def _normalize_autotune_effort(value: object) -> AutotuneEffort:
-    if isinstance(value, str):
-        normalized = value.lower()
-        if normalized in _PROFILES:
-            return cast("AutotuneEffort", normalized)
-    raise ValueError("autotune_effort must be one of 'none', 'quick', or 'full'")
-
-
-def _get_autotune_effort() -> AutotuneEffort:
-    return _normalize_autotune_effort(os.environ.get("HELION_AUTOTUNE_EFFORT", "full"))
-
-
-def _get_autotune_precompile() -> str | None:
-    value = os.environ.get("HELION_AUTOTUNE_PRECOMPILE")
-    if value is None:
-        return "spawn"
-    mode = value.strip().lower()
-    if mode in {"", "0", "false", "none"}:
-        return None
-    if mode in {"spawn", "fork"}:
-        return mode
-    raise ValueError(
-        "HELION_AUTOTUNE_PRECOMPILE must be 'spawn', 'fork', or empty to disable precompile"
-    )
-
-
-def _get_autotune_precompile_jobs() -> int | None:
-    value = os.environ.get("HELION_AUTOTUNE_PRECOMPILE_JOBS")
-    if value is None or value.strip() == "":
-        return None
-    jobs = int(value)
-    if jobs <= 0:
-        raise ValueError("HELION_AUTOTUNE_PRECOMPILE_JOBS must be a positive integer")
-    return jobs
-
-
 @dataclasses.dataclass
 class _Settings:
     # see __slots__ below for the doc strings that show up in help(Settings)
@@ -182,33 +165,45 @@ class _Settings:
         os.environ.get("HELION_AUTOTUNE_COMPILE_TIMEOUT", "60")
     )
     autotune_precompile: str | None = dataclasses.field(
-        default_factory=_get_autotune_precompile
+        default_factory=lambda: os.environ.get("HELION_AUTOTUNE_PRECOMPILE", "spawn")
     )
     autotune_precompile_jobs: int | None = dataclasses.field(
-        default_factory=_get_autotune_precompile_jobs
+        default_factory=lambda: int(v)
+        if (v := os.environ.get("HELION_AUTOTUNE_PRECOMPILE_JOBS"))
+        else None
     )
     autotune_random_seed: int = dataclasses.field(
-        default_factory=_get_autotune_random_seed
+        default_factory=lambda: (
+            int(v)
+            if (v := os.environ.get("HELION_AUTOTUNE_RANDOM_SEED"))
+            else int(time.time() * 1000) % 2**32
+        )
     )
     autotune_accuracy_check: bool = (
         os.environ.get("HELION_AUTOTUNE_ACCURACY_CHECK", "1") == "1"
     )
     autotune_rebenchmark_threshold: float | None = dataclasses.field(
-        default_factory=_get_autotune_rebenchmark_threshold
+        default_factory=lambda: float(v)
+        if (v := os.environ.get("HELION_REBENCHMARK_THRESHOLD"))
+        else None
     )
     autotune_progress_bar: bool = (
         os.environ.get("HELION_AUTOTUNE_PROGRESS_BAR", "1") == "1"
     )
     autotune_max_generations: int | None = dataclasses.field(
-        default_factory=_get_autotune_max_generations
+        default_factory=lambda: int(v)
+        if (v := os.environ.get("HELION_AUTOTUNE_MAX_GENERATIONS"))
+        else None
     )
     print_output_code: bool = os.environ.get("HELION_PRINT_OUTPUT_CODE", "0") == "1"
     force_autotune: bool = os.environ.get("HELION_FORCE_AUTOTUNE", "0") == "1"
     autotune_config_overrides: dict[str, object] = dataclasses.field(
         default_factory=dict
     )
     autotune_effort: AutotuneEffort = dataclasses.field(
-        default_factory=_get_autotune_effort
+        default_factory=lambda: cast(
+            "AutotuneEffort", os.environ.get("HELION_AUTOTUNE_EFFORT", "full")
+        )
     )
     allow_warp_specialize: bool = (
         os.environ.get("HELION_ALLOW_WARP_SPECIALIZE", "1") == "1"
@@ -220,35 +215,43 @@ class _Settings:
     autotuner_fn: AutotunerFunction = default_autotuner_fn
 
     def __post_init__(self) -> None:
-        def _is_bool(val: object) -> bool:
-            return isinstance(val, bool)
-
-        def _is_non_negative_int(val: object) -> bool:
-            return isinstance(val, int) and val >= 0
+        # Validate all user settings
+
+        self.autotune_effort = cast(
+            "AutotuneEffort",
+            _validate_enum_setting(
+                self.autotune_effort,
+                name="autotune_effort",
+                valid=_PROFILES.keys(),
+                allow_none=False,  # do not allow None as "none" is a non-default setting
+            ),
+        )
+        self.autotune_precompile = _validate_enum_setting(
+            self.autotune_precompile,
+            name="autotune_precompile",
+            valid={"spawn", "fork"},
+        )
 
-        # Validate user settings
         validators: dict[str, Callable[[object], bool]] = {
-            "autotune_log_level": _is_non_negative_int,
-            "autotune_compile_timeout": _is_non_negative_int,
-            "autotune_precompile": lambda v: v in (None, "spawn", "fork"),
-            "autotune_precompile_jobs": lambda v: v is None or _is_non_negative_int(v),
-            "autotune_accuracy_check": _is_bool,
-            "autotune_progress_bar": _is_bool,
-            "autotune_max_generations": lambda v: v is None or _is_non_negative_int(v),
-            "print_output_code": _is_bool,
-            "force_autotune": _is_bool,
-            "allow_warp_specialize": _is_bool,
-            "debug_dtype_asserts": _is_bool,
+            "autotune_log_level": lambda v: isinstance(v, int) and v >= 0,
+            "autotune_compile_timeout": lambda v: isinstance(v, int) and v > 0,
+            "autotune_precompile_jobs": lambda v: v is None
+            or (isinstance(v, int) and v > 0),
+            "autotune_accuracy_check": lambda v: isinstance(v, bool),
+            "autotune_progress_bar": lambda v: isinstance(v, bool),
+            "autotune_max_generations": lambda v: v is None
+            or (isinstance(v, int) and v >= 0),
+            "print_output_code": lambda v: isinstance(v, bool),
+            "force_autotune": lambda v: isinstance(v, bool),
+            "allow_warp_specialize": lambda v: isinstance(v, bool),
+            "debug_dtype_asserts": lambda v: isinstance(v, bool),
             "autotune_rebenchmark_threshold": lambda v: v is None
             or (isinstance(v, (int, float)) and v >= 0),
         }
 
-        normalized_effort = _normalize_autotune_effort(self.autotune_effort)
-        object.__setattr__(self, "autotune_effort", normalized_effort)
-
-        for field_name, checker in validators.items():
+        for field_name, validator in validators.items():
             value = getattr(self, field_name)
-            if not checker(value):
+            if not validator(value):
                 raise ValueError(f"Invalid value for {field_name}: {value!r}")
 
 

test/test_settings.py

@@ -7,12 +7,10 @@
 
 class TestSettingsValidation(unittest.TestCase):
     def test_autotune_effort_none_raises(self) -> None:
-        with self.assertRaisesRegex(
-            ValueError, "autotune_effort must be one of 'none', 'quick', or 'full'"
-        ):
+        with self.assertRaisesRegex(ValueError, "autotune_effort must be one of"):
             helion.Settings(autotune_effort=None)
 
-    def test_autotune_effort_quick_normalized(self) -> None:
+    def test_autotune_effort_case_insensitive(self) -> None:
         settings = helion.Settings(autotune_effort="Quick")
         self.assertEqual(settings.autotune_effort, "quick")
 
@@ -33,3 +31,7 @@ def test_autotune_max_generations_negative_raises(self) -> None:
             ValueError, r"Invalid value for autotune_max_generations: -1"
         ):
             helion.Settings(autotune_max_generations=-1)
+
+    def test_autotune_effort_invalid_raises(self) -> None:
+        with self.assertRaisesRegex(ValueError, "autotune_effort must be one of"):
+            helion.Settings(autotune_effort="super-fast")