Repository: microsoft/nni
File: nni/retiarii/execution/base.py (Python)

Code before:

import logging
import os
import random
import string
from typing import Any, Dict, Iterable, List

from .interface import AbstractExecutionEngine, AbstractGraphListener
from .. import codegen, utils
from ..graph import Model, ModelStatus, MetricData, Evaluator
from ..integration_api import send_trial, receive_trial_parameters, get_advisor, report_search_space

_logger = logging.getLogger(__name__)

def extract_visual_hp(model: Model) -> dict:
    hp_config = {}
    for mut in model.history:
        if len(mut.samples) == 1:
            hp_config[mut.mutator.label] = mut.samples[0]
        else:
            for i, sample in enumerate(mut.samples):
                hp_config[f'{mut.mutator.label}_{i}'] = sample
    return hp_config

class BaseGraphData:
    def __init__(self, model_script: str, evaluator: Evaluator, visual_hp: dict = None) -> None:
        self.model_script = model_script
        self.evaluator = evaluator
        self.visual_hp = visual_hp

    def dump(self) -> dict:
        return {
            'model_script': self.model_script,
            'evaluator': self.evaluator,
            '_visual_hyper_params_': self.visual_hp
        }

    @staticmethod
    def load(data) -> 'BaseGraphData':
        return BaseGraphData(data['model_script'], data['evaluator'])


class BaseExecutionEngine(AbstractExecutionEngine):
    """
    The execution engine with no optimization at all.
    Resource management is implemented in this class.
    """

    def __init__(self) -> None:
        """
        Upon initialization, advisor callbacks need to be registered.

Review comment (suggestion):
Suggest renaming it into something like `mutation_summary`. It's not for visualization purpose only.

Code after:

import logging
import os
import random
import string
from typing import Any, Dict, Iterable, List

from .interface import AbstractExecutionEngine, AbstractGraphListener
from .utils import get_mutation_summary
from .. import codegen, utils
from ..graph import Model, ModelStatus, MetricData, Evaluator
from ..integration_api import send_trial, receive_trial_parameters, get_advisor

_logger = logging.getLogger(__name__)

class BaseGraphData:
    """
    Attributes
    ----------
    model_script
        code of an instantiated PyTorch model
    evaluator
        training approach for model_script
    mutation_summary
        a dict of all the choices during mutations in the HPO search space format
    """
    def __init__(self, model_script: str, evaluator: Evaluator, mutation_summary: dict) -> None:
        self.model_script = model_script
        self.evaluator = evaluator
        self.mutation_summary = mutation_summary

    def dump(self) -> dict:
        return {
            'model_script': self.model_script,
            'evaluator': self.evaluator,
            'mutation_summary': self.mutation_summary
        }

    @staticmethod
    def load(data) -> 'BaseGraphData':
        return BaseGraphData(data['model_script'], data['evaluator'], data['mutation_summary'])


class BaseExecutionEngine(AbstractExecutionEngine):
    """
    The execution engine with no optimization at all.
    Resource management is implemented in this class.
    """

    def __init__(self) -> None:
        """