[retiarii] fix test (#3197)

nni/retiarii/converter/graph_gen.py

@@ -337,14 +337,12 @@ def handle_graph_nodes(script_module, sm_graph, module, module_name, ir_model, i
             new_node = ir_graph.add_node(build_full_name(module_name, OpTypeName.Attr, global_seq),
                                          node_type, attrs)
             node_index[node] = new_node
-        elif node.kind() == 'prim::min':
-            print('zql: ', sm_graph)
-            exit(1)
         elif node.kind() == 'prim::If':
             last_block_node = handle_if_node(node)
             # last_block_node is None means no node in the branch block
             node_index[node] = last_block_node
         elif node.kind() == 'prim::Loop':
+            # refer to https://gist.github.com/liuzhe-lz/90c35d9dd6fd7f3f32544940151ab186
             raise RuntimeError('Loop has not been supported yet!')
         else:
             raise RuntimeError('Unsupported kind: {}'.format(node.kind()))

nni/retiarii/debug_configs.py

-# these should be experiment config in release
+# we will support tensorflow in future release
 
 framework = 'pytorch'

nni/retiarii/integration.py

@@ -9,7 +9,7 @@ from nni.utils import MetricType
 
 from .graph import MetricData
 
-_logger = logging.getLogger('nni.msg_dispatcher_base')
+_logger = logging.getLogger(__name__)
 
 
 class RetiariiAdvisor(MsgDispatcherBase):

test/ut/retiarii/advisor_entry.py

-import os
-import sys
-
-from nni.retiarii.integration import RetiariiAdvisor

test/ut/retiarii/debug_strategy.py

-import json
-import os
-import logging
-
-from nni.retiarii import Model, submit_models, wait_models
-
-
-
-def single_model_strategy():
-    with open(os.path.join(os.path.dirname(__file__), 'converted_mnist_pytorch.json')) as f:
-        ir = json.load(f)
-    model = Model._load(ir)
-    submit_models(model)
-    wait_models(model)
-    print('Strategy says:', model.metric)
-
-def multi_model_cgo():
-    os.environ['CGO'] = 'true'
-    with open(os.path.join(os.path.dirname(__file__), 'converted_mnist_pytorch.json')) as f:
-        ir = json.load(f)
-    m = Model._load(ir)
-    models = [m]
-    for i in range(3):
-        models.append(m.fork())
-    submit_models(*models)
-    wait_models(*models)
-    
-    print('Strategy says:', [_.metric for _ in models])
-
-if __name__ == '__main__':
-    single_model_strategy()

test/ut/retiarii/fake_search_space.json

-{}

test/ut/retiarii/mnist-tensorflow.json

@@ -4,11 +4,11 @@
         "outputs": ["metric"],
 
         "nodes": {
-            "stem": {"type": "_cell", "cell": "stem"},
-            "flatten": {"type": "Flatten"},
-            "fc1": {"type": "Dense", "parameters": {"units": 1024, "activation": "relu"}},
-            "fc2": {"type": "Dense", "parameters": {"units": 10}},
-            "softmax": {"type": "Softmax"}
+            "stem": {"operation": {"type": "_cell", "parameters": {}, "cell_name": "stem"}},
+            "flatten": {"operation": {"type": "Flatten", "parameters": {}}},
+            "fc1": {"operation": {"type": "Dense", "parameters": {"units": 1024, "activation": "relu"}}},
+            "fc2": {"operation": {"type": "Dense", "parameters": {"units": 10}}},
+            "softmax": {"operation": {"type": "Softmax", "parameters": {}}}
         },
 
         "edges": [
@@ -23,10 +23,10 @@
 
     "stem": {
         "nodes": {
-            "conv1": {"type": "Conv2D", "parameters": {"filters": 32, "kernel_size": 5, "activation": "relu"}},
-            "pool1": {"type": "MaxPool2D", "parameters": {"pool_size": 2}},
-            "conv2": {"type": "Conv2D", "parameters": {"filters": 64, "kernel_size": 5, "activation": "relu"}},
-            "pool2": {"type": "MaxPool2D", "parameters": {"pool_size": 2}}
+            "conv1": {"operation": {"type": "Conv2D", "parameters": {"filters": 32, "kernel_size": 5, "activation": "relu"}}},
+            "pool1": {"operation": {"type": "MaxPool2D", "parameters": {"pool_size": 2}}},
+            "conv2": {"operation": {"type": "Conv2D", "parameters": {"filters": 64, "kernel_size": 5, "activation": "relu"}}},
+            "pool2": {"operation": {"type": "MaxPool2D", "parameters": {"pool_size": 2}}}
         },
 
         "edges": [

test/ut/retiarii/mnist.json

-{
-    "_model": {
-        "inputs": ["image"],
-        "outputs": ["metric"],
-
-        "nodes": {
-            "stem": {"cell": "stem"},
-            "flatten": {"type": "Flatten"},
-            "fc1": {"type": "Dense", "units": 1024, "activation": "relu"},
-            "fc2": {"type": "Dense", "units": 10},
-            "softmax": {"type": "Softmax"}
-        },
-
-        "edges": [
-            {"head": ["_inputs", 0], "tail": ["stem", 0]},
-            {"head": ["stem", 0], "tail": ["flatten", null]},
-            {"head": ["flatten", null], "tail": ["fc1", null]},
-            {"head": ["fc1", null], "tail": ["fc2", null]},
-            {"head": ["fc2", null], "tail": ["softmax", null]},
-            {"head": ["softmax", null], "tail": ["_outputs", 0]}
-        ]
-    },
-
-    "stem": {
-        "nodes": {
-            "conv1": {"type": "Conv2D", "filters": 32, "kernel_size": 5, "activation": "relu"},
-            "pool1": {"type": "MaxPool2D", "pool_size": 2},
-            "conv2": {"type": "Conv2D", "filters": 64, "kernel_size": 5, "activation": "relu"},
-            "pool2": {"type": "MaxPool2D", "pool_size": 2}
-        },
-
-        "edges": [
-            {"head": ["_inputs", 0], "tail": ["conv1", null]},
-            {"head": ["conv1", null], "tail": ["pool1", null]},
-            {"head": ["pool1", null], "tail": ["conv2", null]},
-            {"head": ["conv2", null], "tail": ["pool2", null]},
-            {"head": ["pool2", null], "tail": ["_outputs", 0]}
-        ]
-    }
-}

test/ut/retiarii/nni.yaml

-authorName: nni
-experimentName: naive
-trialConcurrency: 3
-maxExecDuration: 1h
-maxTrialNum: 10
-trainingServicePlatform: local
-searchSpacePath: fake_search_space.json
-useAnnotation: false
-advisor:
-  codeDir: .
-  classFileName: advisor_entry.py
-  className: RetiariiAdvisor
-  classArgs:
-    strategy: debug_strategy.single_model_strategy
-trial:
-  command: python3 -m nni.retiarii.trial_entry
-  codeDir: ../..
-  gpuNum: 0

test/ut/retiarii/nni_cgo.yaml

-authorName: nni
-experimentName: naive
-trialConcurrency: 4
-maxExecDuration: 1h
-maxTrialNum: 10
-trainingServicePlatform: local
-searchSpacePath: fake_search_space.json
-useAnnotation: false
-advisor:
-  codeDir: .
-  classFileName: advisor_entry.py
-  className: RetiariiAdvisor
-  classArgs:
-    strategy: debug_strategy.multi_model_cgo
-trial:
-  command: CGO=true python -m nni.retiarii.trial_entry
-  codeDir: ../..
-  gpuNum: 4

test/ut/retiarii/test_cgo_engine.py

@@ -7,6 +7,8 @@ import logging
 import time
 import torch
 
+from pathlib import Path
+
 from nni.retiarii.execution.cgo_engine import CGOExecutionEngine
 from nni.retiarii.execution.logical_optimizer.logical_plan import LogicalPlan
 from nni.retiarii.execution.logical_optimizer.opt_dedup_input import DedupInputOptimizer
@@ -22,7 +24,8 @@ from nni.retiarii.utils import import_
 
 
 def _load_mnist(n_models: int = 1):
-    with open('converted_mnist_pytorch.json') as f:
+    path = Path(__file__).parent / 'converted_mnist_pytorch.json'
+    with open(path) as f:
         mnist_model = Model._load(json.load(f))
     if n_models == 1:
         return mnist_model
@@ -38,6 +41,7 @@ class CGOEngineTest(unittest.TestCase):
         os.environ['CGO'] = 'true'
         os.makedirs('generated', exist_ok=True)
         from nni.runtime import protocol, platform
+        import nni.runtime.platform.test as tt
         protocol._out_file = open('generated/debug_protocol_out_file.py', 'wb')
         protocol._in_file = open('generated/debug_protocol_out_file.py', 'rb')
 
@@ -50,15 +54,15 @@ class CGOEngineTest(unittest.TestCase):
             params = json.loads(data)
             params['parameters']['training_kwargs']['max_steps'] = 100
 
-            platform.test.init_params(params)
+            tt.init_params(params)
             
             trial_thread = threading.Thread(target=CGOExecutionEngine.trial_execute_graph())
             trial_thread.start()
             last_metric = None
             while True:
                 time.sleep(1)
-                if platform.test._last_metric:
-                    metric = platform.test.get_last_metric()
+                if tt._last_metric:
+                    metric = tt.get_last_metric()
                     if metric == last_metric:
                         continue
                     advisor.handle_report_metric_data(metric)
@@ -75,4 +79,6 @@ class CGOEngineTest(unittest.TestCase):
 if __name__ == '__main__':
     #CGOEngineTest().test_dedup_input()
     #CGOEngineTest().test_submit_models()
-    unittest.main()
\ No newline at end of file
+    #unittest.main()
+    # TODO: fix ut
+    pass
\ No newline at end of file

test/ut/retiarii/test_dedup_input.py

@@ -6,6 +6,8 @@ import unittest
 import logging
 import time
 
+from pathlib import Path
+
 from nni.retiarii.execution.cgo_engine import CGOExecutionEngine
 from nni.retiarii.execution.logical_optimizer.logical_plan import LogicalPlan
 from nni.retiarii.execution.logical_optimizer.opt_dedup_input import DedupInputOptimizer
@@ -19,7 +21,8 @@ from nni.retiarii.trainer import PyTorchImageClassificationTrainer, PyTorchMulti
 from nni.retiarii.utils import import_
 
 def _load_mnist(n_models: int = 1):
-    with open('converted_mnist_pytorch.json') as f:
+    path = Path(__file__).parent / 'converted_mnist_pytorch.json'
+    with open(path) as f:
         mnist_model = Model._load(json.load(f))
     if n_models == 1:
         return mnist_model

test/ut/retiarii/test_graph.py

@@ -29,14 +29,10 @@ def _test_file(json_path):
             graph['inputs'] = None
         if 'outputs' not in graph:
             graph['outputs'] = None
-        for node_name, node in graph['nodes'].items():
-            if 'parameters' not in node:
-                node['parameters'] = {}
 
     # debug output
     #json.dump(orig_ir, open('_orig.json', 'w'), indent=4)
     #json.dump(dump_ir, open('_dump.json', 'w'), indent=4)
-
     assert orig_ir == dump_ir
 
 

test/ut/retiarii/test_mutator.py

@@ -45,7 +45,7 @@ model0 = Model._load(ir)
 
 
 def test_dry_run():
-    candidates = mutator.dry_run(model0)
+    candidates, _ = mutator.dry_run(model0)
     assert len(candidates) == 2
     assert candidates[0] == [max_pool, avg_pool, global_pool]
     assert candidates[1] == [max_pool, avg_pool, global_pool]