Class FedMLRunner

python/fedml/runner.py:19–183

class FedMLRunner:
    def __init__(
        self,
        args,
        device,
        dataset,
        model,
        client_trainer: ClientTrainer = None,
        server_aggregator: ServerAggregator = None,
        algorithm_flow: FedMLAlgorithmFlow = None,
    ):
        if algorithm_flow is not None:
            self.runner = algorithm_flow
            return

        if args.training_type == FEDML_TRAINING_PLATFORM_SIMULATION:
            init_runner_func = self._init_simulation_runner

        elif args.training_type == FEDML_TRAINING_PLATFORM_CROSS_SILO:
            init_runner_func = self._init_cross_silo_runner

        elif args.training_type == FEDML_TRAINING_PLATFORM_CROSS_CLOUD:
            init_runner_func = self._init_cheetah_runner

        elif args.training_type == FEDML_TRAINING_PLATFORM_SERVING:
            init_runner_func = self._init_model_serving_runner

        elif args.training_type == FEDML_TRAINING_PLATFORM_CROSS_DEVICE:
            init_runner_func = self._init_cross_device_runner
        else:
            raise Exception("no such setting")

        self.runner = init_runner_func(
            args, device, dataset, model, client_trainer, server_aggregator
        )

    def _init_simulation_runner(
        self, args, device, dataset, model, client_trainer=None, server_aggregator=None
    ):
        if hasattr(args, "backend") and args.backend == FEDML_SIMULATION_TYPE_SP:
            from .simulation.simulator import SimulatorSingleProcess

            runner = SimulatorSingleProcess(
                args, device, dataset, model, client_trainer, server_aggregator
            )
        elif hasattr(args, "backend") and args.backend == FEDML_SIMULATION_TYPE_MPI:
            from .simulation.simulator import SimulatorMPI

            runner = SimulatorMPI(
                args, device, dataset, model, client_trainer, server_aggregator
            )
        elif hasattr(args, "backend") and args.backend == FEDML_SIMULATION_TYPE_NCCL:
            from .simulation.simulator import SimulatorNCCL

            runner = SimulatorNCCL(
                args, device, dataset, model, client_trainer, server_aggregator
            )
        else: