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

Code before:
    for custom_ts_pkg in custom_ts_config.keys():
        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:
    # avoid circular import
    from nni.nas.experiment.config import RetiariiExeConfig
    from nni.nas.experiment.pytorch import RetiariiExperiment
    from ..experiment_config import ExperimentConfig
    from ...experiment import Experiment
    exp_cls = get_experiment_cls_using_config(config_json)
    if exp_cls is Experiment:
        return ExperimentConfig(**config_json)
    elif exp_cls is RetiariiExperiment:
        return RetiariiExeConfig(**config_json)
    else:
        raise TypeError(f'Unsupported experiment type: {type(exp_cls)}')

def get_experiment_cls_using_config(config_json: dict) -> type[Experiment] | type[RetiariiExperiment]:
    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 ValueError(f'Unknown experiment_type: {config_json["experimentType"]}')
    else:
        if 'executionEngine' in config_json:
            return RetiariiExperiment
        else:
            return Experiment


Review comment (suggestion):
If you raise exception in `get_experiment_cls_using_config`, this part can be `return exp_cls(**config_json)`

Code after:
    for custom_ts_pkg in custom_ts_config.keys():
        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