Repository: microsoft/nni
File: nni/experiment/config/utils/internal.py (Python)

Code before:

    for cls in TrainingServiceConfig.__subclasses__():
        if cls.platform == platform:
            return cls
    return None

## misc ##

def get_ipv4_address() -> str:
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(('192.0.2.0', 80))
    addr = s.getsockname()[0]
    s.close()
    return addr

def load_experiment_config(config_json) -> typing.Union[ExperimentConfig, RetiariiExeConfig]:
    # avoid circular import
    from nni.nas.experiment.config import RetiariiExeConfig
    from ..experiment_config import ExperimentConfig
    if 'experimentType' in config_json:
        if config_json['experimentType'] == 'hpo':
            return ExperimentConfig(**config_json)
        elif config_json['experimentType'] == 'nas':
            return RetiariiExeConfig(**config_json)
        else:
            raise KeyError(f'Unknown experiment_type: {config_json["experimentType"]}')
    else:
        # for backward compatibility, experiment config <= v2.8 does not have "experiment_type"
        if 'executionEngine' in config_json:
            return RetiariiExeConfig(**config_json)
        else:
            return ExperimentConfig(**config_json)

def get_experiment_cls_using_config(config_json):
    from nni.nas.experiment.pytorch import RetiariiExperiment
    from ...experiment import Experiment
    if 'experimentType' in config_json:
        if config_json['experimentType'] == 'hpo':
            return Experiment
        elif config_json['experimentType'] == 'nas':
            return RetiariiExperiment
        else:
            raise KeyError(f'Unknown experiment_type: {config_json["experimentType"]}')
    else:
        if 'executionEngine' in config_json:
            return RetiariiExperiment
        else:
            return Experiment


Review comment (bug):
`KeyError` is used when key is wrong.
It should raise `ValueError` here.

Code after:
        pkg = importlib.import_module(custom_ts_pkg)
        _config_class = pkg.nni_training_service_info.config_class

    for cls in TrainingServiceConfig.__subclasses__():
        if cls.platform == platform:
            return cls
    return None

## misc ##

def get_ipv4_address() -> str:
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(('192.0.2.0', 80))
    addr = s.getsockname()[0]
    s.close()
    return addr

def load_experiment_config(config_json: dict) -> ExperimentConfig | RetiariiExeConfig:
    _, exp_conf_cls = get_experiment_cls_using_config(config_json)
    return exp_conf_cls(**config_json)

def get_experiment_cls_using_config(config_json: dict) -> Tuple[type[Experiment] | type[RetiariiExperiment],
                                                                type[ExperimentConfig] | type[RetiariiExeConfig]]:
    # avoid circular import and unnecessary dependency on pytorch
    if 'experimentType' in config_json:
        if config_json['experimentType'] == 'hpo':
            from ...experiment import Experiment
            from ..experiment_config import ExperimentConfig
            return Experiment, ExperimentConfig
        elif config_json['experimentType'] == 'nas':
            from nni.nas.experiment.pytorch import RetiariiExperiment
            from nni.nas.experiment.config import RetiariiExeConfig
            return RetiariiExperiment, RetiariiExeConfig
        else:
            raise ValueError(f'Unknown experiment_type: {config_json["experimentType"]}')
    else:
        if 'executionEngine' in config_json:
            from nni.nas.experiment.pytorch import RetiariiExperiment
            from nni.nas.experiment.config import RetiariiExeConfig
            return RetiariiExperiment, RetiariiExeConfig
        else:
            from ...experiment import Experiment
            from ..experiment_config import ExperimentConfig
            return Experiment, ExperimentConfig