bug fix

Author: chenyushuo

tests/explorer/runner_pool_test.py

@@ -6,6 +6,7 @@
 import ray
 import torch
 
+from tests.tools import get_unittest_dataset_config
 from trinity.buffer.reader.queue_reader import QueueReader
 from trinity.common.config import StorageConfig, load_config
 from trinity.common.constants import AlgorithmType, StorageType
@@ -21,7 +22,7 @@
 @WORKFLOWS.register_module("dummy_workflow")
 class DummyWorkflow(Workflow):
     def __init__(self, model, **kwargs):
-        super().__init__(model)
+        super().__init__(model, **kwargs)
         self.error_type = kwargs.get("task_desc")
         self.seconds = None
         if "timeout" in self.error_type:
@@ -81,30 +82,37 @@ def setUp(self):
 
     def test_runner_pool(self):
         pool = RunnerPool(self.config, [DummyModel.remote(), DummyModel.remote()])
+        taskset_config = get_unittest_dataset_config("countdown")
         tasks = [
             Task(
                 task_desc="timeout_100",
                 workflow=DummyWorkflow,
+                taskset_config=taskset_config,
             ),
             Task(
                 task_desc="exception",
                 workflow=DummyWorkflow,
+                taskset_config=taskset_config,
             ),
             Task(
                 task_desc="timeout_2",
                 workflow=DummyWorkflow,
+                taskset_config=taskset_config,
             ),
             Task(
                 task_desc="success",
                 workflow=DummyWorkflow,
+                taskset_config=taskset_config,
             ),
             Task(
                 task_desc="timeout_101",
                 workflow=DummyWorkflow,
+                taskset_config=taskset_config,
             ),
             Task(
                 task_desc="exit",
                 workflow=DummyWorkflow,
+                taskset_config=taskset_config,
             ),
         ]
 

tests/template/config.yaml

@@ -18,7 +18,7 @@ buffer:
     taskset:
       name: taskset
       storage_type: file
-      path: ''
+      path: 'placeholder'
       split: 'train'
     default_workflow_type: ''
     default_reward_fn_type: ''

trinity/explorer/workflow_runner.py

@@ -77,7 +77,7 @@ def run_task(self, task: Task) -> Status:
             if metrics:
                 for k, v in metrics.items():
                     metric[k] = sum(v) / len(v)  # type: ignore
-            if not task.storage_config.task_type == TaskType.EVAL:
+            if not task.taskset_config.task_type == TaskType.EVAL:
                 self.experience_buffer.write(exps)
             return Status(True, metric=metric)
         except Exception as e: