Repository: microsoft/nni
File: test/ut/retiarii/test_multitrial.py (Python)

Code before:

def test_multi_trial():
    evaluator_kwargs = {
        'max_epochs': 1
    }

    to_test = [
        # (model, evaluator)
        _mnist_net('simple', evaluator_kwargs),
        _mnist_net('simple_value_choice', evaluator_kwargs),
        _mnist_net('value_choice', evaluator_kwargs),
        _mnist_net('repeat', evaluator_kwargs),
        _mnist_net('custom_op', evaluator_kwargs),
    ]

    for base_model, evaluator in to_test:
        search_strategy = strategy.Random()
        exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
        exp_config = RetiariiExeConfig('local')
        exp_config.experiment_name = 'mnist_unittest'
        exp_config.trial_concurrency = 1
        exp_config.max_trial_number = 1
        exp_config.training_service.use_active_gpu = False
        exp.run(exp_config, 8080)
        assert isinstance(exp.export_top_models()[0], dict)
        exp.stop()

python_script = """
from nni.retiarii import strategy
from nni.retiarii.experiment.pytorch import RetiariiExeConfig, RetiariiExperiment
from test_oneshot import _mnist_net

base_model, evaluator = _mnist_net('simple', {'max_epochs': 1})
search_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_unittest'
exp_config.trial_concurrency = 1
exp_config.max_trial_number = 1
exp_config.training_service.use_active_gpu = False
exp.run(exp_config, 8080)
assert isinstance(exp.export_top_models()[0], dict)
"""

def test_exp_exit_without_stop():
    script_name = 'tmp_multi_trial.py'
    with open(script_name, 'w') as f:
        f.write(python_script)
    proc = Popen([sys.executable, script_name])
    proc.wait()
    os.remove(script_name)

Review comment (suggestion):
experiment can quickly release the port after `stop()`? (I don't know, maybe `local` don't have to wait for the port release)

Code after:

def test_multi_trial():
    evaluator_kwargs = {
        'max_epochs': 1
    }

    to_test = [
        # (model, evaluator)
        _mnist_net('simple', evaluator_kwargs),
        _mnist_net('simple_value_choice', evaluator_kwargs),
        _mnist_net('value_choice', evaluator_kwargs),
        _mnist_net('repeat', evaluator_kwargs),
        _mnist_net('custom_op', evaluator_kwargs),
    ]

    for base_model, evaluator in to_test:
        search_strategy = strategy.Random()
        exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
        exp_config = RetiariiExeConfig('local')
        exp_config.experiment_name = 'mnist_unittest'
        exp_config.trial_concurrency = 1
        exp_config.max_trial_number = 1
        exp_config.training_service.use_active_gpu = False
        exp.run(exp_config, 8080)
        assert isinstance(exp.export_top_models()[0], dict)
        exp.stop()

python_script = """
from nni.retiarii import strategy
from nni.retiarii.experiment.pytorch import RetiariiExeConfig, RetiariiExperiment
from test_oneshot import _mnist_net

base_model, evaluator = _mnist_net('simple', {'max_epochs': 1})
search_strategy = strategy.Random()
exp = RetiariiExperiment(base_model, evaluator, strategy=search_strategy)
exp_config = RetiariiExeConfig('local')
exp_config.experiment_name = 'mnist_unittest'
exp_config.trial_concurrency = 1
exp_config.max_trial_number = 1
exp_config.training_service.use_active_gpu = False
exp.run(exp_config, 8080)
assert isinstance(exp.export_top_models()[0], dict)
"""

@pytest.mark.timeout(600)
def test_exp_exit_without_stop():
    script_name = 'tmp_multi_trial.py'
    with open(script_name, 'w') as f:
        f.write(python_script)
    proc = Popen([sys.executable, script_name])
    proc.wait()