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

Code before:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import time
from dataclasses import dataclass
from typing import Iterable

from ..graph import Model, ModelStatus
from .interface import AbstractExecutionEngine
from .listener import DefaultListener

_execution_engine = None
_default_listener = None

__all__ = ['get_execution_engine', 'get_and_register_default_listener',
           'list_models', 'submit_models', 'wait_models', 'query_available_resources',
           'set_execution_engine', 'is_stopped_exec', 'budget_exhausted', 'GPUDevice']


@dataclass(init=False)
class GPUDevice:
    def __init__(self, node_id: int, gpu_id: int):
        self.node_id = node_id
        self.gpu_id = gpu_id
        self.status = "free"

    def __eq__(self, o) -> bool:
        return self.node_id == o.node_id and self.gpu_id == o.gpu_id

    def __lt__(self, o) -> bool:
        if self.node_id < o.node_id:
            return True
        elif self.node_id > o.node_id:
            return False
        else:
            return self.gpu_id < o.gpu_id

    def __repr__(self) -> str:
        return "{Server-%d, GPU-%d, Status: %s}" % (self.node_id, self.gpu_id, self.status)

    def __hash__(self) -> int:
        return hash(self.node_id + '_' + self.gpu_id)

    def set_status(self, status):
        self.status = status


def set_execution_engine(engine: AbstractExecutionEngine) -> None:

Review comment (suggestion):
When using dataclass, it should be:
```
@dataclass
class GPUDevice:
    node_id: int
    gpu_id: int
    status: str = 'free'

    def __eq__():
        …
```

Code after:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import time
from typing import Iterable

from ..graph import Model, ModelStatus
from .interface import AbstractExecutionEngine
from .listener import DefaultListener

_execution_engine = None
_default_listener = None

__all__ = ['get_execution_engine', 'get_and_register_default_listener',
           'list_models', 'submit_models', 'wait_models', 'query_available_resources',
           'set_execution_engine', 'is_stopped_exec', 'budget_exhausted']


def set_execution_engine(engine: AbstractExecutionEngine) -> None:
    global _execution_engine
    if _execution_engine is None:
        _execution_engine = engine
    else:
        raise RuntimeError('Execution engine is already set.')


def get_execution_engine() -> AbstractExecutionEngine:
    global _execution_engine
    assert _execution_engine is not None, 'You need to set execution engine, before using it.'
    return _execution_engine


def get_and_register_default_listener(engine: AbstractExecutionEngine) -> DefaultListener:
    global _default_listener
    if _default_listener is None:
        _default_listener = DefaultListener()
        engine.register_graph_listener(_default_listener)
    return _default_listener


def submit_models(*models: Model) -> None:
    engine = get_execution_engine()
    get_and_register_default_listener(engine)
    engine.submit_models(*models)


def list_models(*models: Model) -> Iterable[Model]:
    engine = get_execution_engine()