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

Code before:
_logger = logging.getLogger(__name__)


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

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

    @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, devices=None) -> None:
        """
        Upon initialization, advisor callbacks need to be registered.
        Advisor will call the callbacks when the corresponding event has been triggered.
        Base execution engine will get those callbacks and broadcast them to graph listener.
        """
        self._listeners: List[AbstractGraphListener] = []

        # register advisor callbacks
        advisor = get_advisor()
        advisor.send_trial_callback = self._send_trial_callback
        advisor.request_trial_jobs_callback = self._request_trial_jobs_callback
        advisor.trial_end_callback = self._trial_end_callback
        advisor.intermediate_metric_callback = self._intermediate_metric_callback
        advisor.final_metric_callback = self._final_metric_callback

        self._running_models: Dict[int, Model] = dict()
        self._history: List[Model] = []

        self.resources = 0
        self.devices = [] if devices is None else devices

    def submit_models(self, *models: Model) -> None:
        for model in models:
            data = self.pack_model_data(model)
            self._running_models[send_trial(data.dump())] = model

Review comment (suggestion):
Who will use this `devices`? If it's only useful for `CGOExecutionEngine`, suggest putting it there.

Code after:
_logger = logging.getLogger(__name__)


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

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

    @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.
        Advisor will call the callbacks when the corresponding event has been triggered.
        Base execution engine will get those callbacks and broadcast them to graph listener.
        """
        self._listeners: List[AbstractGraphListener] = []

        # register advisor callbacks
        advisor = get_advisor()
        advisor.send_trial_callback = self._send_trial_callback
        advisor.request_trial_jobs_callback = self._request_trial_jobs_callback
        advisor.trial_end_callback = self._trial_end_callback
        advisor.intermediate_metric_callback = self._intermediate_metric_callback
        advisor.final_metric_callback = self._final_metric_callback

        self._running_models: Dict[int, Model] = dict()
        self._history: List[Model] = []

        self.resources = 0

    def submit_models(self, *models: Model) -> None:
        for model in models:
            data = self.pack_model_data(model)
            self._running_models[send_trial(data.dump())] = model
            self._history.append(model)