Add Task.parameter-sets

Fixes #134

examples/task-example.yaml

@@ -77,3 +77,21 @@
           base: 10
           distribution: uniform
           integerify: False
+- task:
+    step: run training
+    name: parameter sets
+    type: manual_search
+    parameters:
+      - name: A
+        style: multiple
+        rules: {}
+      - name: B
+        style: multiple
+        rules: {}
+    parameter-sets:
+      - A: 5
+        B: 6
+      - A: 8
+        B: 9
+      - A: 72
+        B: 42

tests/test_task.py

@@ -23,3 +23,12 @@ def test_task_additional_fields(task_config: Config):
     assert task.optimization_target_metric == "goodness"
     assert task.optimization_target_value == 7.2
     assert task.type == TaskType.RANDOM_SEARCH
+
+
+def test_task_parameter_sets(task_config: Config):
+    task = task_config.tasks["parameter sets"]
+    assert task.parameter_sets == [
+        {"A": 5, "B": 6},
+        {"A": 8, "B": 9},
+        {"A": 72, "B": 42},
+    ]

valohai_yaml/objs/task.py

@@ -56,6 +56,7 @@ class Task(Item):
     step: str
     type: TaskType
     parameters: list[VariantParameter]
+    parameter_sets: list[dict[str, Any]]
     name: str
     execution_count: int | None
     execution_batch_size: int | None
@@ -72,6 +73,7 @@ def __init__(
         step: str,
         type: TaskType | str | None = None,
         parameters: list[VariantParameter] | None = None,
+        parameter_sets: list[dict[str, Any]] | None = None,
         execution_count: int | None = None,
         execution_batch_size: int | None = None,
         maximum_queued_executions: int | None = None,
@@ -85,6 +87,9 @@ def __init__(
         self.step = step
         self.type = TaskType.cast(type)
         self.parameters = check_type_and_listify(parameters, VariantParameter)
+        self.parameter_sets = [
+            ps for ps in check_type_and_listify(parameter_sets, dict) if ps
+        ]
         self.execution_count = execution_count
         self.execution_batch_size = execution_batch_size
         self.maximum_queued_executions = maximum_queued_executions
@@ -105,3 +110,5 @@ def parse(cls, data: Any) -> Task:
     def lint(self, lint_result: LintResult, context: LintContext) -> None:
         context = dict(context, task=self, object_type="task")
         lint_expression(lint_result, context, "stop-condition", self.stop_condition)
+        if self.parameter_sets and self.type != TaskType.MANUAL_SEARCH:
+            lint_result.add_warning("Parameter sets only make sense with manual search")

valohai_yaml/schema/task.yaml

@@ -37,3 +37,8 @@ properties:
       "$ref": "./variant-param.json"
   stop-condition:
     type: string
+  parameter-sets:
+    type: array
+    description: Parameter sets for manual search mode.
+    items:
+      type: object