diff --git a/cezo_fl/run_client_jobs.py b/cezo_fl/run_client_jobs.py index 967e92c..dde19ae 100644 --- a/cezo_fl/run_client_jobs.py +++ b/cezo_fl/run_client_jobs.py @@ -2,9 +2,11 @@ from typing import Sequence, TypeAlias import torch +from torch.distributed import rpc from cezo_fl.client import AbstractClient, LocalUpdateResult from cezo_fl.util.metrics import Metric +from cezo_fl.fl_helpers import get_client_name def parallalizable_client_job( @@ -38,7 +40,19 @@ def parallalizable_client_job( client_local_update_result = client.local_update(seeds=local_update_seeds) # move result to server device and return - return client_local_update_result.to(server_device) + return client_local_update_result.to("cpu") + + +def parallalizable_client_job_rref( + clientRRef: AbstractClient, + pull_seeds_list: Sequence[Sequence[int]], + pull_grad_list: Sequence[Sequence[torch.Tensor]], + local_update_seeds: Sequence[int], + server_device: torch.device, +) -> LocalUpdateResult: + return parallalizable_client_job( + clientRRef.to_here(), pull_seeds_list, pull_grad_list, local_update_seeds, server_device + ) # Outer list is for clients and inner list for local update data @@ -100,3 +114,33 @@ def execute_sampled_clients( local_grad_scalar_list.append(client_local_update_result.grad_tensors) return step_train_loss, step_train_accuracy, local_grad_scalar_list + + +def rpc_execute_sampled_clients(server, sampled_client_index: Sequence[int], seeds: Sequence[int]): + local_grad_scalar_list: LOCAL_GRAD_SCALAR_LIST = [] # Clients X Local_update + step_train_loss = Metric("Step train loss") + step_train_accuracy = Metric("Step train accuracy") + + futures = [] + for index in sampled_client_index: + client = server.clients[index] + last_update_iter = server.client_last_updates[index] + # The seed and grad in last_update_iter is fetched as well + # Note at that iteration, we just reset the client model so that iteration + # information is needed as well. + seeds_list = server.seed_grad_records.fetch_seed_records(last_update_iter) + grad_list = server.seed_grad_records.fetch_grad_records(last_update_iter) + grad_list = [[tensor.to("cpu") for tensor in tensors] for tensors in grad_list] + args = (client, seeds_list, grad_list, seeds, server.device) + + futures.append( + rpc.rpc_async(to=get_client_name(index), func=parallalizable_client_job_rref, args=args) + ) + + client_results = [f.wait().to(server.device) for f in futures] + for client_local_update_result in client_results: + step_train_loss.update(client_local_update_result.step_loss) + step_train_accuracy.update(client_local_update_result.step_accuracy) + local_grad_scalar_list.append(client_local_update_result.grad_tensors) + + return step_train_loss, step_train_accuracy, local_grad_scalar_list diff --git a/cezo_fl/server.py b/cezo_fl/server.py index 9a753f8..ad39fdd 100644 --- a/cezo_fl/server.py +++ b/cezo_fl/server.py @@ -8,7 +8,7 @@ from cezo_fl.client import AbstractClient from cezo_fl.random_gradient_estimator import RandomGradientEstimator -from cezo_fl.run_client_jobs import execute_sampled_clients +from cezo_fl.run_client_jobs import execute_sampled_clients, rpc_execute_sampled_clients from cezo_fl.shared import CriterionType from cezo_fl.util.metrics import Metric @@ -163,8 +163,11 @@ def train_one_step(self, iteration: int) -> tuple[float, float]: seeds = [random.randint(0, 1000000) for _ in range(self.local_update_steps)] # Step 1 & 2: pull model and local update - step_train_loss, step_train_accuracy, local_grad_scalar_list = execute_sampled_clients( - self, sampled_client_index, seeds, parallel=False + # step_train_loss, step_train_accuracy, local_grad_scalar_list = execute_sampled_clients( + # self, sampled_client_index, seeds, parallel=False + # ) + step_train_loss, step_train_accuracy, local_grad_scalar_list = rpc_execute_sampled_clients( + self, sampled_client_index, seeds ) for index in sampled_client_index: diff --git a/rpc_decomfl.py b/rpc_decomfl.py new file mode 100644 index 0000000..b3806db --- /dev/null +++ b/rpc_decomfl.py @@ -0,0 +1,208 @@ +import functools +from os import path +import torch + + +from tensorboardX import SummaryWriter +from tqdm import tqdm + +from byzantine import aggregation as byz_agg + + +from byzantine import attack as byz_attack +from cezo_fl.client import ResetClient +from cezo_fl.fl_helpers import get_client_name, get_server_name +from cezo_fl.server import CeZO_Server +from cezo_fl.util import model_helpers +from experiment_helper import prepare_settings +from experiment_helper.cli_parser import ( + GeneralSetting, + DeviceSetting, + DataSetting, + ModelSetting, + OptimizerSetting, + FederatedLearningSetting, + RGESetting, + ByzantineSetting, +) +from experiment_helper.device import use_device +from experiment_helper.data import get_dataloaders +import os +import torch +import torch.distributed.rpc as rpc +import torch.multiprocessing as mp +import time +import os + + +class CliSetting( + GeneralSetting, + DeviceSetting, + DataSetting, + ModelSetting, + OptimizerSetting, + FederatedLearningSetting, + RGESetting, + ByzantineSetting, +): + """ + This is a replacement for regular argparse module. + We used a third party library pydantic_setting to make command line interface easier to manage. + Example: + if __name__ == "__main__": + args = CliSetting() + + args will have all parameters defined by all components. + """ + + pass + + +def setup_server_and_clients( + args: CliSetting, device_map: dict[str, torch.device], train_loaders +) -> CeZO_Server: + model_inferences, metrics = prepare_settings.get_model_inferences_and_metrics( + args.dataset, args.model_setting + ) + client_rrefs = [] + + for i in range(args.num_clients): + client_ref = rpc.remote( + to=get_client_name(i), func=rpc_setup_client_from_server, args=(i, args) + ) + client_rrefs += [client_ref] + + server_device = device_map[get_server_name()] + server = CeZO_Server( + client_rrefs, + server_device, + num_sample_clients=args.num_sample_clients, + local_update_steps=args.local_update_steps, + ) + + # set server tools + server_model = prepare_settings.get_model( + dataset=args.dataset, model_setting=args.model_setting, seed=args.seed + ).to(server_device) + server_optimizer = prepare_settings.get_optimizer( + model=server_model, dataset=args.dataset, optimizer_setting=args.optimizer_setting + ) + server_grad_estimator = prepare_settings.get_random_gradient_estimator( + model=server_model, + device=server_device, + rge_setting=args.rge_setting, + model_setting=args.model_setting, + ) + + server.set_server_model_and_criterion( + server_model, + model_inferences.test_inference, + metrics.test_loss, + metrics.test_acc, + server_optimizer, + server_grad_estimator, + ) + + return server + + +def rpc_setup_client_from_server(rank, args): + device_map = use_device(args.device_setting, args.num_clients) + train_loaders, _ = get_dataloaders( + args.data_setting, args.num_clients, args.seed, args.get_hf_model_name() + ) + + model_inferences, metrics = prepare_settings.get_model_inferences_and_metrics( + args.dataset, args.model_setting + ) + + client_name = get_client_name(rank) + client_device = device_map[client_name] + client_model = prepare_settings.get_model( + dataset=args.dataset, model_setting=args.model_setting, seed=args.seed + ).to(client_device) + client_optimizer = prepare_settings.get_optimizer( + model=client_model, dataset=args.dataset, optimizer_setting=args.optimizer_setting + ) + client_grad_estimator = prepare_settings.get_random_gradient_estimator( + model=client_model, + device=client_device, + rge_setting=args.rge_setting, + model_setting=args.model_setting, + ) + + client = ResetClient( + client_model, + model_inferences.train_inference, + train_loaders[rank], + client_grad_estimator, + client_optimizer, + metrics.train_loss, + metrics.train_acc, + client_device, + ) + return client + + +def rpc_worker(rank, world_size, args: CliSetting): + # rank >= 1 for worker, thus need to minus to keep it the same as device_map's name + rpc.init_rpc(get_client_name(rank - 1), rank=rank, world_size=world_size) + rpc.shutdown() + + +def print_work_info(client_ref): + print("worker", rpc.get_worker_info()) + print("client", client_ref.owner_name()) + print(client_ref.to_here().device) + return client_ref.to_here().device + + +def rpc_server(rank, world_size, args: CliSetting): + """Caller function that sends an RPC to the worker.""" + rpc.init_rpc(get_server_name(), rank=rank, world_size=world_size) + print(args) + # Name of the worker to send the task to + device_map = use_device(args.device_setting, args.num_clients) + _, test_loader = get_dataloaders( + args.data_setting, args.num_clients, args.seed, args.get_hf_model_name() + ) + server = setup_server_and_clients(args, device_map, test_loader) + # add 5s sleep here, need to fix this + # time.sleep(5) + + # ret = rpc.rpc_async( + # get_client_name(0), + # print_work_info, + # (server.clients[0],), + # ) + # print(ret.wait()) + # server.train_one_step(0) + with tqdm(total=args.iterations, desc="Training:") as t, torch.no_grad(): + for ite in range(args.iterations): + step_loss, step_accuracy = server.train_one_step(ite) + t.set_postfix({"Loss": step_loss, "Accuracy": step_accuracy}) + t.update(1) + + # eval + if args.eval_iterations != 0 and (ite + 1) % args.eval_iterations == 0: + eval_loss, eval_accuracy = server.eval_model(test_loader) + + rpc.shutdown() + + +def fn(rank, world_size, args): + if rank == 0: + return rpc_server(rank, world_size, args) + else: + return rpc_worker(rank, world_size, args) + + +if __name__ == "__main__": + args = CliSetting() + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29500" + world_size = args.num_clients + 1 # We will have 1 server + num_clients worker + + # Spawn two processes, one for the caller and one for the worker + mp.spawn(fn, args=(world_size, args), nprocs=world_size, join=True)