for t in range(num_rounds):
    clients = sample_clients(cohort_size)
    client_features = features(clients)
    configs = linear_model(client_features)

    round_grads, client_evals = [], []
    for config, client in zip(configs, client):
        delta, client_eval_loss = client.train(config, global_model)
        round_grads.append(delta)
        client_evals.append(client_eval_loss)

    global_model = server.step(average(round_grads))
    linear_model = optimize(linear_model, client_evals, configs)