Implementation:Recommenders team Recommenders NNI Utils

Knowledge Sources	Recommenders
Domains	Hyperparameter Tuning, NNI, Experiment Management
Last Updated	2026-02-10 00:00 GMT

Overview

Utility functions for managing NNI (Neural Network Intelligence) hyperparameter tuning experiments, including starting, monitoring, stopping experiments, and retrieving trial results via NNI's REST API.

Description

This module provides a complete Python interface for orchestrating NNI hyperparameter tuning experiments. It interacts with NNI's REST API (default endpoint http://localhost:8080/api/v1/nni) and the nnictl command-line tool through several functions:

• start_nni launches an experiment via the nnictl create subprocess with a YAML configuration file and waits for the experiment to complete.
• check_experiment_status polls the NNI status endpoint with configurable wait time (default 20 seconds) and max retries (default 10), checking for DONE or error states.
• check_stopped verifies that the NNI endpoint is no longer accessible after stopping an experiment.
• check_metrics_written waits until all trials have finalMetricData available in the REST response.
• get_trials retrieves all trial results, parses metric data using ast.literal_eval, sorts by the default metric according to the specified optimize_mode ("minimize" or "maximize"), and reads the best trial's metrics and parameters from its log directory.
• stop_nni calls nnictl stop and verifies shutdown.
• get_experiment_status is a low-level helper that performs a GET request to the status URL.

All polling functions use retry loops with configurable wait times and raise TimeoutError if the maximum number of retries is exceeded.

Usage

Use these utilities to programmatically manage NNI experiments from Python scripts or Jupyter notebooks. They abstract away the complexity of NNI's REST API and command-line interface, enabling automated hyperparameter tuning workflows for recommendation models.

Code Reference

Source Location

• Repository: Recommenders
• File: recommenders/tuning/nni/nni_utils.py
• Lines: 1-154

Signature

def get_experiment_status(status_url=NNI_STATUS_URL)

def check_experiment_status(wait=WAITING_TIME, max_retries=MAX_RETRIES)

def check_stopped(wait=WAITING_TIME, max_retries=MAX_RETRIES)

def check_metrics_written(wait=WAITING_TIME, max_retries=MAX_RETRIES)

def get_trials(optimize_mode)

def stop_nni()

def start_nni(config_path, wait=WAITING_TIME, max_retries=MAX_RETRIES)

Import

from recommenders.tuning.nni.nni_utils import (
    start_nni,
    stop_nni,
    get_trials,
    check_experiment_status,
    check_stopped,
    check_metrics_written,
    get_experiment_status,
)

I/O Contract

Inputs

start_nni

Name	Type	Required	Description
config_path	str	Yes	Path to the NNI experiment YAML configuration file
wait	numeric	No	Seconds to wait between status checks (default: 20)
max_retries	int	No	Maximum number of polling retries (default: 10)

check_experiment_status

Name	Type	Required	Description
wait	numeric	No	Seconds to wait between status checks (default: 20)
max_retries	int	No	Maximum number of polling retries (default: 10)

get_trials

Name	Type	Required	Description
optimize_mode	str	Yes	One of "minimize" or "maximize"; determines sorting order for finding the best trial

get_experiment_status

Name	Type	Required	Description
status_url	str	No	URL for the NNI REST status endpoint (default: NNI_STATUS_URL)

Outputs

get_experiment_status

Name	Type	Description
return	dict	JSON response from the NNI status endpoint containing experiment status

get_trials

Name	Type	Description
trials	list	List of tuples: (metrics_dict, log_path) for each trial
best_metrics	dict	Metrics dictionary for the best trial
best_params	dict	Hyperparameters used by the best trial
best_trial_path	str	File system path to the best trial's log directory

Usage Examples

Basic Usage

from recommenders.tuning.nni.nni_utils import start_nni, stop_nni, get_trials

# Start an NNI experiment
start_nni("nni_config.yml", wait=30, max_retries=50)

# Retrieve trial results (minimize the primary metric)
trials, best_metrics, best_params, best_path = get_trials("minimize")

print("Best metrics:", best_metrics)
print("Best hyperparameters:", best_params)

# Stop the NNI experiment
stop_nni()