Fix blocking issue at test_tasks.py

Author: KKIEEK

tests/test_mm/test_tasks.py

@@ -5,86 +5,82 @@
 import mmdet  # noqa: F401
 import mmseg  # noqa: F401
 import pytest
+import ray
 import torch
 from mmcv.utils import Config
 from ray import tune
 from ray.air import session
+from ray.tune.result_grid import ResultGrid
 
 from siatune.mm.tasks import (TASKS, BaseTask, BlackBoxTask,
                               ContinuousTestFunction, DiscreteTestFunction,
                               MMClassification, MMDetection, MMSegmentation,
                               MMTrainBasedTask, build_task_processor)
 from siatune.utils.config import dump_cfg
 
-_session = dict()
-
 
-def report_to_session(*args, **kwargs):
-    _session = get_session()
-    _session.update(kwargs)
-    for arg in args:
-        if isinstance(arg, dict):
-            _session.update(arg)
+@pytest.fixture
+def init_ray():
+    if ray.is_initialized():
+        ray.shutdown()
+    return ray.init(num_cpus=1)
 
 
-def get_session():
-    global _session
-    return _session
-
-
-@patch('ray.tune.report', side_effect=report_to_session)
-def test_base_task(mock_report):
+def test_base_task(init_ray):
     with pytest.raises(TypeError):
         BaseTask()
 
     class TestRewriter:
 
         def __call__(self, context):
-            context.get('args').test = -1
-            return context
+            args = context.pop('args')
+            args.test = 'success'
+            return dict(args=args)
 
     class TestTask(BaseTask):
 
-        def parse_args(self, *args, **kwargs):
-            return argparse.Namespace(test=1)
+        def parse_args(self, args):
+            return argparse.Namespace(test='init')
 
-        def run(self, *, args, **kwargs):
+        def run(self, args):
             tune.report(test=args.test)
-            return args.test
 
         def create_trainable(self):
             return self.context_aware_run
 
     task = TestTask([TestRewriter()])
-    task.set_args('')
-    assert task.args == argparse.Namespace(test=1)
-    assert isinstance(task.rewriters, list)
-    task.context_aware_run({})
-    assert get_session().get('test') == -1
+    task.set_args([])
+    assert task.args == argparse.Namespace(test='init')
 
-    tune.run(task.create_trainable(), config={})
+    trainable = task.create_trainable()
+    result = ResultGrid(tune.run(trainable, config={}))[0]
+    assert result.metrics['test'] == 'success'
 
 
-def test_black_box_task():
+def test_black_box_task(init_ray):
     with pytest.raises(TypeError):
         BlackBoxTask()
 
     class TestTask(BlackBoxTask):
 
         def run(self, *args, **kwargs):
-            tune.report(test=1)
+            tune.report(test='success')
 
     task = TestTask()
-    task.set_args('')
+    task.set_args([])
     assert task.args == argparse.Namespace()
-    tune.run(task.create_trainable(), config={})
+
+    trainable = task.create_trainable()
+    result = ResultGrid(tune.run(trainable, config={}))[0]
+    assert result.metrics['test'] == 'success'
 
 
 def test_build_task_processor():
 
+    @TASKS.register_module()
     class TestTaks(BaseTask):
 
-        def parse_args(self, *args, **kwargs):
+        def parse_args(self, args):
             pass
 
         def run(self, *args, **kwargs):
@@ -93,12 +89,22 @@ def run(self, *args, **kwargs):
         def create_trainable(self, *args, **kwargs):
             pass
 
-    TASKS.register_module(TestTaks)
-    assert isinstance(build_task_processor(dict(type='TestTaks')), TestTaks)
+    task = build_task_processor(dict(type='TestTaks'))
+    assert isinstance(task, TestTaks)
+
+
+_session = dict()
+
+
+def report_to_session(*args, **kwargs):
+    _session.update(kwargs)
+    for arg in args:
+        if isinstance(arg, dict):
+            _session.update(arg)
 
 
 @patch('ray.tune.report', side_effect=report_to_session)
-def test_continuous_test_function(mock_report):
+def test_continuous_test_function(init_ray):
     func = ContinuousTestFunction()
     predefined_cont_funcs = [
         'delayedsphere',
@@ -147,11 +153,11 @@ def test_continuous_test_function(mock_report):
             'test.py')
         args = argparse.Namespace(config='test.py')
         func.run(args=args)
-        assert isinstance(get_session().get('result'), float)
+        assert isinstance(_session['result'], float)
 
 
 @patch('ray.tune.report', side_effect=report_to_session)
-def test_discrete_test_function(mock_report):
+def test_discrete_test_function(init_ray):
     func = DiscreteTestFunction()
 
     predefined_discrete_funcs = ['onemax', 'leadingones', 'jump']
@@ -161,7 +167,7 @@ def test_discrete_test_function(mock_report):
             'test.py')
         args = argparse.Namespace(config='test.py')
         func.run(args=args)
-        assert isinstance(get_session().get('result'), float)
+        assert isinstance(_session['result'], float)
 
 
 @patch('mmcls.apis.train_model')
@@ -194,8 +200,7 @@ def test_mmseg(*not_used):
     task.run(args=task.args)
 
 
-@patch('ray.air.session.report', side_effect=report_to_session)
-def test_mm_train_based_task(mock_report):
+def test_mm_train_based_task(init_ray):
     with pytest.raises(TypeError):
         MMTrainBasedTask()
 
@@ -250,7 +255,7 @@ def train_model(self, model, dataset, cfg):
                     loss.backward()
                     optimizer.step()
                     total_loss += loss.item()
-                session.report(loss=total_loss / (batch_idx + 1))
+                session.report(dict(loss=total_loss / (batch_idx + 1)))
 
         def run(self, *, searched_cfg, **kwargs):
             cfg = searched_cfg.get('cfg')
@@ -274,7 +279,7 @@ def run(self, *, searched_cfg, **kwargs):
     task = TestTask()
     task.set_resource()
     task.context_aware_run(searched_cfg=dict(cfg=cfg))
-    assert 'loss' in get_session()
 
     trainable = task.create_trainable()
-    tune.Tuner(trainable).fit()
+    result = tune.Tuner(trainable).fit()[0]
+    assert 'loss' in result.metrics