[Feature] Support trial name template (#146)

Author: KKIEEK

siatune/tune/tuner.py

@@ -13,6 +13,7 @@
 from siatune.codebase import build_task
 from siatune.tune import (build_callback, build_scheduler, build_searcher,
                           build_space, build_stopper)
+from .utils import NAME_CREATOR
 
 
 class Tuner:
@@ -106,8 +107,10 @@ def __init__(
             tune_config=TuneConfig(
                 search_alg=searcher,
                 scheduler=trial_scheduler,
-                trial_name_creator=lambda trial: trial.trial_id,
-                trial_dirname_creator=lambda trial: trial.experiment_tag,
+                trial_name_creator=NAME_CREATOR.get(
+                    tune_cfg.pop('trial_name_creator', 'trial_id')),
+                trial_dirname_creator=NAME_CREATOR.get(
+                    tune_cfg.pop('trial_dirname_creator', 'experiment_tag')),
                 **tune_cfg),
             run_config=RunConfig(
                 name=self.experiment_name,

siatune/tune/utils/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) SI-Analytics. All rights reserved.
+from .name_creator import NAME_CREATOR, experiment_tag, trial_id
+
+__all__ = ['NAME_CREATOR', 'experiment_tag', 'trial_id']

siatune/tune/utils/name_creator.py

@@ -0,0 +1,15 @@
+# Copyright (c) SI-Analytics. All rights reserved.
+from mmengine.registry import Registry
+from ray.tune.experiment import Trial
+
+NAME_CREATOR = Registry('name creator')
+
+
+@NAME_CREATOR.register_module()
+def trial_id(trial: Trial) -> str:
+    return trial.trial_id
+
+
+@NAME_CREATOR.register_module()
+def experiment_tag(trial: Trial) -> str:
+    return trial.experiment_tag

tests/test_tune/test_utils.py

@@ -0,0 +1,28 @@
+import inspect
+from unittest.mock import MagicMock
+
+import pytest
+
+from siatune.tune.utils import NAME_CREATOR
+
+
+@pytest.fixture
+def trial():
+    return MagicMock(
+        trainable_name='test',
+        trial_id='trial_id',
+        experiment_tag='experiment_tag')
+
+
+def test_trial_id(trial):
+    tmpl = NAME_CREATOR.get('trial_id')
+    assert inspect.isfunction(tmpl)
+    assert tmpl.__name__ == 'trial_id'
+    assert tmpl(trial) == trial.trial_id
+
+
+def test_experiment_tag(trial):
+    tmpl = NAME_CREATOR.get('experiment_tag')
+    assert inspect.isfunction(tmpl)
+    assert tmpl.__name__ == 'experiment_tag'
+    assert tmpl(trial) == trial.experiment_tag