# Buggy code
def check_experiment_status(wait=WAITING_TIME, max_retries=MAX_RETRIES):
    """Checks the status of the current experiment on the NNI REST endpoint
    Waits until the tuning has completed

    Args:
        wait (numeric) : time to wait in seconds
        max_retries (int): max number of retries
    """
    i = 0
    while i < max_retries:
        status = get_experiment_status(NNI_STATUS_URL)
        if status in ["DONE", "TUNER_NO_MORE_TRIAL"]:
            break
        elif status not in ["RUNNING", "NO_MORE_TRIAL"]:
            raise RuntimeError(
                "NNI experiment failed to complete with status {}".format(status)
            )
        time.sleep(wait)
        i += 1
    if i == max_retries:
        raise TimeoutError("check_experiment_status() timed out")


# Traceback (PapermillExecutionError)
tests/integration/test_movielens.py .........                            [ 62%]
tests/integration/test_notebooks_python.py .....F                        [100%]

=================================== FAILURES ===================================
_____________________________ test_nni_tuning_svd ______________________________

notebooks = {'als_deep_dive': '/data/home/recocat/cicd/7/s/notebooks/02_model/als_deep_dive.ipynb', 'als_pyspark': '/data/home/rec...baseline_deep_dive.ipynb', 'data_split': '/data/home/recocat/cicd/7/s/notebooks/01_prepare_data/data_split.ipynb', ...}
tmp = '/tmp/pytest-of-recocat/pytest-200/tmpuj9_hwzq'

@pytest.mark.integration
def test_nni_tuning_svd(notebooks, tmp):
notebook_path = notebooks["nni_tuning_svd"]
# First stop NNI in case it is running
subprocess.run([sys.prefix + '/bin/nnictl', 'stop'])
check_stopped()
pm.execute_notebook(notebook_path, OUTPUT_NOTEBOOK, kernel_name=KERNEL_NAME,
parameters=dict(MOVIELENS_DATA_SIZE="100k",
SURPRISE_READER="ml-100k",
TMP_DIR=tmp,
MAX_TRIAL_NUM=1,
NUM_EPOCHS=1))

tests/integration/test_notebooks_python.py:164:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
/anaconda/envs/nightly_reco_base/lib/python3.6/site-packages/papermill/execute.py:94: in execute_notebook
raise_for_execution_errors(nb, output_path)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

nb = {'cells': [{'cell_type': 'code', 'metadata': {'inputHidden': True, 'hide_input': True}, 'execution_count': None, 'sour...d_time': '2019-04-18T12:09:01.272263', 'duration': 345.187067, 'exception': True}}, 'nbformat': 4, 'nbformat_minor': 2}
output_path = 'output.ipynb'

def raise_for_execution_errors(nb, output_path):
"""Assigned parameters into the appropriate place in the input notebook

Parameters
----------
nb : NotebookNode
Executable notebook object
output_path : str
Path to write executed notebook
"""
error = None
for cell in nb.cells:
if cell.get("outputs") is None:
continue

for output in cell.outputs:
if output.output_type == "error":
error = PapermillExecutionError(
exec_count=cell.execution_count,
source=cell.source,
ename=output.ename,
evalue=output.evalue,
traceback=output.traceback,
)
break

if error:
# Write notebook back out with the Error Message at the top of the Notebook.
error_msg = ERROR_MESSAGE_TEMPLATE % str(error.exec_count)
error_msg_cell = nbformat.v4.new_code_cell(
source="%%html\n" + error_msg,
outputs=[
nbformat.v4.new_output(output_type="display_data", data={"text/html": error_msg})
],
metadata={"inputHidden": True, "hide_input": True},
)
nb.cells = [error_msg_cell] + nb.cells
write_ipynb(nb, output_path)
raise error
E           papermill.exceptions.PapermillExecutionError:
E           ---------------------------------------------------------------------------
E           Exception encountered at "In [32]":
E           ---------------------------------------------------------------------------
E           RuntimeError                              Traceback (most recent call last)
E           <ipython-input-32-120fb8357fc9> in <module>
E                 1 t = time.time()
E           ----> 2 stop_and_restart()
E                 3 time_smac = time.time() - t
E
E           <ipython-input-20-609ee7726ae6> in stop_and_restart()
E                 6     proc = subprocess.run([sys.prefix + '/bin/nnictl', 'create', '--config', config_path])
E                 7     if proc.returncode != 0:
E           ----> 8         raise RuntimeError("'nnictl create' failed with code %d" % proc.returncode)
E                 9     check_experiment_status(wait=WAITING_TIME, max_retries=MAX_RETRIES)
E
E           RuntimeError: 'nnictl create' failed with code 1

/anaconda/envs/nightly_reco_base/lib/python3.6/site-packages/papermill/execute.py:241: PapermillExecutionError

# Fixed code
def check_experiment_status(wait=WAITING_TIME, max_retries=MAX_RETRIES):
    """Checks the status of the current experiment on the NNI REST endpoint
    Waits until the tuning has completed

    Args:
        wait (numeric) : time to wait in seconds
        max_retries (int): max number of retries
    """
    i = 0
    while i < max_retries:
        nni_status = get_experiment_status(NNI_STATUS_URL)
        if nni_status["status"] in ["DONE", "TUNER_NO_MORE_TRIAL"]:
            break
        elif nni_status["status"] not in ["RUNNING", "NO_MORE_TRIAL"]:
            raise RuntimeError(
                "NNI experiment failed to complete with status {} - {}".format(
                    nni_status["status"], nni_status["errors"][0]
                )
            )
        time.sleep(wait)
        i += 1
    if i == max_retries:
        raise TimeoutError("check_experiment_status() timed out")