df_runner.py

class Runner:
    def __init__(self, X, y):
        self.X = X
        self.y = y

    # called by FedAvg algo.
    def optimise(self, intercept_init, coef_init, hyperparameters):
        _intercept_init = intercept_init.copy()
        _coef_init = coef_init.copy()
        model = train_model(_intercept_init, _coef_init, self.X, self.y, **hyperparameters)
        return model

    def dataset_size(self):
        return self.X.shape[0]