feat: auto-enable torch.compile/AMP on CUDA via apply_hardware_defaults

tinyml-modelmaker/tests/test_hardware_defaults.py

@@ -0,0 +1,54 @@
+from unittest.mock import patch
+import pytest
+
+
+def _make_params():
+    """Build a minimal ConfigDict with the two training flags at their defaults."""
+    from tinyml_modelmaker.utils.config_dict import ConfigDict
+    return ConfigDict(dict(training=dict(compile_model=0, native_amp=False)))
+
+
+def test_auto_enables_both_on_cuda_with_no_explicit_keys():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, set())
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
+
+
+def test_respects_explicit_native_amp_false():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, {'native_amp'})
+    assert params.training.compile_model == 1   # still auto-enabled
+    assert params.training.native_amp is False  # explicit choice respected
+
+
+def test_respects_explicit_compile_model_zero():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, {'compile_model'})
+    assert params.training.compile_model == 0   # explicit choice respected
+    assert params.training.native_amp is True   # still auto-enabled
+
+
+def test_no_change_without_cuda():
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = _make_params()
+    with patch('torch.cuda.is_available', return_value=False):
+        apply_hardware_defaults(params, set())
+    assert params.training.compile_model == 0
+    assert params.training.native_amp is False
+
+
+def test_safe_when_params_lacks_flags():
+    """hasattr guards: calling on a params dict without compile_model/native_amp must not raise."""
+    from tinyml_modelmaker.utils.config_dict import ConfigDict
+    from tinyml_modelmaker.utils.hardware_defaults import apply_hardware_defaults
+    params = ConfigDict(dict(training=dict(batch_size=32)))
+    with patch('torch.cuda.is_available', return_value=True):
+        apply_hardware_defaults(params, set())  # must not raise
+    assert params.training.batch_size == 32

tinyml-modelmaker/tests/test_hardware_defaults_integration.py

@@ -0,0 +1,69 @@
+from unittest.mock import patch
+import pytest
+
+
+def test_init_params_auto_enables_on_cuda():
+    """init_params with no training overrides auto-enables both flags on CUDA."""
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    user_config = dict(common=dict(task_category='timeseries_classification'))
+    with patch('torch.cuda.is_available', return_value=True):
+        params = init_params(user_config)
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
+
+
+def test_init_params_respects_native_amp_false_override():
+    """Explicit native_amp: false in user config is not overridden.
+
+    Mirrors the real production call pattern in run_tinyml_modelmaker.py:
+    ModelRunner.init_params() is called with zero arguments (so hardware
+    defaults auto-enable via apply_hardware_defaults with an empty
+    explicitly_set), and the user's config is merged in afterward via
+    ConfigDict.update() -- a deep merge -- not passed into the ConfigDict
+    constructor. (The constructor's *args merge path has a separate,
+    pre-existing shallow-merge bug that this test must avoid triggering.)
+    """
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    user_config = dict(
+        common=dict(task_category='timeseries_classification'),
+        training=dict(native_amp=False),
+    )
+    with patch('torch.cuda.is_available', return_value=True):
+        params = init_params()
+    params.update(user_config)
+    assert params.training.native_amp is False
+    assert params.training.compile_model == 1  # still auto-enabled
+
+
+def test_init_params_no_change_without_cuda():
+    """Without CUDA, both flags stay at defaults."""
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    user_config = dict(common=dict(task_category='timeseries_classification'))
+    with patch('torch.cuda.is_available', return_value=False):
+        params = init_params(user_config)
+    assert params.training.compile_model == 0
+    assert params.training.native_amp is False
+
+
+def test_init_params_does_not_crash_on_non_dict_first_arg():
+    """ConfigDict documents a YAML path string (or None) as valid first
+    positional input, not just a dict. init_params's own explicitly-set-key
+    detection must not assume args[0] is dict-like and crash instead."""
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    with patch('torch.cuda.is_available', return_value=True):
+        # A bogus path is fine here -- ConfigDict's *args merge loop only
+        # merges dict/ConfigDict values and silently ignores strings, so
+        # this never actually attempts to read the file. The point of this
+        # test is solely that passing a string doesn't crash init_params's
+        # own explicitly-set-key detection before ConfigDict is even built.
+        params = init_params('/nonexistent/path/to/config.yaml')
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True
+
+
+def test_init_params_does_not_crash_on_none_first_arg():
+    from tinyml_modelmaker.ai_modules.timeseries.params import init_params
+    with patch('torch.cuda.is_available', return_value=True):
+        params = init_params(None)
+    assert params.training.compile_model == 1
+    assert params.training.native_amp is True

tinyml-modelmaker/tinyml_modelmaker/ai_modules/timeseries/params.py

@@ -35,6 +35,7 @@
 
 from ... import utils
 from . import constants
+from ...utils.hardware_defaults import apply_hardware_defaults
 
 
 def init_params(*args, **kwargs):
@@ -222,5 +223,11 @@ def init_params(*args, **kwargs):
         ),
     )
 
+    # args[0] is usually a user config dict, but ConfigDict itself also
+    # accepts a YAML path string or None as a first positional argument --
+    # only inspect it as a mapping when it actually is one.
+    user_training_keys = set(args[0].get('training', {}).keys()) \
+        if args and isinstance(args[0], dict) else set()
     params = utils.ConfigDict(default_params, *args, **kwargs)
+    apply_hardware_defaults(params, user_training_keys)
     return params

tinyml-modelmaker/tinyml_modelmaker/utils/hardware_defaults.py

@@ -0,0 +1,19 @@
+import torch
+
+
+def apply_hardware_defaults(params, explicitly_set: set) -> None:
+    """Auto-enable compile_model and native_amp when CUDA is available.
+
+    Skips fields present in explicitly_set — those are deliberate user
+    choices from the YAML config and must not be overridden.
+    hasattr guards keep this safe for params that don't carry these fields
+    yet (vision, audio — Phase 2).
+    """
+    if not torch.cuda.is_available():
+        return
+    if 'compile_model' not in explicitly_set and hasattr(params.training, 'compile_model'):
+        if getattr(params.training, 'compile_model', 0) == 0:
+            params.training.compile_model = 1
+    if 'native_amp' not in explicitly_set and hasattr(params.training, 'native_amp'):
+        if not getattr(params.training, 'native_amp', False):
+            params.training.native_amp = True