Module run
Expand source code
#!/usr/bin/env python3
import os
import sys
import json
import socket
import math
import signal
# Import PyTorch root package import torch
import torch
from torch.utils.collect_env import get_pretty_env_info
import numpy as np
import time
import copy
import threading
import pickle
from utils import comm_socket
from utils import gpu_utils
from utils import execution_context
from utils import thread_pool
from utils import wandb_wrapper
from opts import parse_args
from utils.logger import Logger
from data_preprocess.data_loader import load_data, get_test_batch_size
from utils.model_funcs import get_training_elements, evaluate_model, get_lr_scheduler, get_optimiser, \
run_one_communication_round, local_training, initialise_model
from utils.checkpointing import save_checkpoint
from utils.checkpointing import defered_eval_and_save_checkpoint
from utils.utils import create_model_dir, create_metrics_dict
from utils.fl_funcs import get_sampled_clients
from models import RNN_MODELS
from models import mutils
from utils import algorithms
from utils.buffer import Buffer
def main(args, raw_cmdline, extra_):
# Init project and save in args url for the plots
projectWB = None
if len(args.wandb_key) > 0:
projectWB = wandb_wrapper.initWandbProject(
args.wandb_key,
args.wandb_project_name,
f"{args.algorithm}-rounds-{args.rounds}-{args.model}@{args.dataset}-runid-{args.run_id}",
args
)
if projectWB is not None:
args.wandb_url = projectWB.url
# In case of DataParallel for .to() to
# This a device used by master for FL algorithms
args.device = gpu_utils.get_target_device_str(args.gpu[0]) if type(args.gpu) == list else args.gpu
# Instantiate execution context
exec_ctx = execution_context.initExecutionContext()
exec_ctx.extra_ = extra_
# Set all seeds
if args.deterministic:
exec_ctx.random.seed(args.manual_init_seed)
exec_ctx.np_random.seed(args.manual_init_seed)
# Setup extra options for experiments
for option in args.algorithm_options.split(","):
kv = option.split(":")
if len(kv) == 1:
exec_ctx.experimental_options[kv[0]] = True
else:
exec_ctx.experimental_options[kv[0]] = kv[1]
# Load validation set
trainset, testset = load_data(exec_ctx, args.data_path, args.dataset, args, load_trainset=True, download=True)
test_batch_size = get_test_batch_size(args.dataset, args.batch_size)
# Test set dataloader
testloader = torch.utils.data.DataLoader(testset,
batch_size=test_batch_size,
num_workers=args.num_workers_test,
shuffle=False,
pin_memory=False,
drop_last=False
)
# Reset all seeds
if args.deterministic:
exec_ctx.random.seed(args.manual_init_seed)
exec_ctx.np_random.seed(args.manual_init_seed)
if not args.evaluate: # Training mode
trainloader = torch.utils.data.DataLoader(
trainset,
batch_size=args.batch_size,
num_workers=args.num_workers_train,
shuffle=False,
pin_memory=False,
drop_last=False)
# Is target device for master is GPU
is_target_dev_gpu = gpu_utils.is_target_dev_gpu(args.device)
# Initialize local training threads in case of using multithreading implementation
exec_ctx.local_training_threads.adjust_num_workers(w=args.threadpool_for_local_opt, own_cuda_stream=True,
device_list=args.gpu)
for th in exec_ctx.local_training_threads.threads:
th.trainset_copy = copy.deepcopy(trainset)
th.train_loader_copy = torch.utils.data.DataLoader(th.trainset_copy,
batch_size=args.batch_size,
num_workers=args.num_workers_train,
shuffle=False,
pin_memory=False,
drop_last=False)
if hasattr(th.trainset_copy, "load_data"):
th.trainset_copy.load_data()
# Initialize remote clients - one remote connection per one remote computation device
external_devices = [d.strip() for d in args.external_devices.split(",") if len(d.strip()) > 0]
exec_ctx.remote_training_threads.adjust_num_workers(w=len(external_devices), own_cuda_stream=False,
device_list=["cpu"])
for th_i in range(len(exec_ctx.remote_training_threads.threads)):
th = exec_ctx.remote_training_threads.threads[th_i]
dev = external_devices[th_i]
th.external_dev_long = dev
th.external_ip, th.external_port, th.external_dev, th.external_dev_dumber = dev.split(":")
try:
th.external_socket = comm_socket.CommSocket()
th.external_socket.sock.connect((th.external_ip, int(th.external_port))) # connect with remote side
th.external_socket.rawSendString("execute_work") # initiate work execution COMMAND
th.external_socket.rawSendString(th.external_dev_long) # provide device specification
th.external_socket.rawSend(pickle.dumps(
raw_cmdline)) # provide original command line (it will be changed by the client) to be consistent with used devices
th.external_socket_online = True
except socket.error:
th.external_socket_online = False
# Initialize saver treads for information serialization (0 is acceptable)
exec_ctx.saver_thread.adjust_num_workers(w=args.save_async_threads, own_cuda_stream=True, device_list=args.gpu)
# Initialize local training threads in case of using multithreading implementation (0 is acceptable)
exec_ctx.eval_thread_pool.adjust_num_workers(w=args.eval_async_threads, own_cuda_stream=True,
device_list=args.gpu)
for th in exec_ctx.eval_thread_pool.threads:
th.testset_copy = copy.deepcopy(testset)
th.testloader_copy = torch.utils.data.DataLoader(th.testset_copy,
batch_size=test_batch_size,
num_workers=args.num_workers_test,
shuffle=False,
pin_memory=False,
drop_last=False
)
if hasattr(th.testset_copy, "load_data"):
th.testset_copy.load_data()
if args.worker_listen_mode <= 0:
# Path of execution for local simulation
init_and_train_model(args, raw_cmdline, trainloader, testloader, exec_ctx)
else:
# Path of execution for assist simulation with a remote means
init_and_help_with_compute(args, raw_cmdline, trainloader, testloader, exec_ctx)
else: # Evaluation mode
# TODO: figure out how to exploit DataParallel. Currently - parallelism across workers
model, criterion, round = get_training_elements(args.model, args.dataset, args, args.resume_from,
args.load_best, args.device, args.loss,
args.turn_off_batch_normalization_and_dropout)
metrics = evaluate_model(model, testloader, criterion, args.device, round,
print_freq=10, metric_to_optim=args.metric,
is_rnn=args.model in RNN_MODELS)
metrics_dict = create_metrics_dict(metrics)
logger.info(f'Validation metrics: {metrics_dict}')
wandb_wrapper.finishProject(projectWB)
def init_and_help_with_compute(args, raw_cmdline, trainloader, testloader, exec_ctx):
logger = Logger.get(args.run_id)
model_dir = create_model_dir(args)
# don't train if setup already exists
if os.path.isdir(model_dir):
logger.info(f"{model_dir} already exists.")
logger.info("Skipping this setup.")
return
# create model directory
os.makedirs(model_dir, exist_ok=True)
# save used args as json to experiment directory
with open(os.path.join(create_model_dir(args), 'args.json'), 'w') as f:
json.dump(vars(args), f, indent=4)
is_rnn = args.model in RNN_MODELS
# TODO: figure out how to exploit DataParallel. Currently - parallelism across workers
model, criterion, current_round = get_training_elements(args.model, args.dataset, trainloader.dataset, args,
args.resume_from, args.load_best, args.device, args.loss,
args.turn_off_batch_normalization_and_dropout)
# ==================================================================================================================
# Reset execution seeds for tunable runtime behaviour
if args.deterministic:
exec_ctx.random.seed(args.manual_runtime_seed)
exec_ctx.np_random.seed(args.manual_runtime_seed)
# ==================================================================================================================
mutils.print_models_info(model, args)
logger.info(f"Number of parameters in the model: {mutils.number_of_params(model):,d}\n")
gpu_utils.print_info_about_used_gpu(args.device, args.run_id)
# ==================================================================================================================
# Initialize server state
# ==================================================================================================================
trainloader.dataset.set_client(None)
algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False,
evaluate_function=False, device=args.device, args=args)
gradient_at_start = mutils.get_gradient(model)
for p in model.parameters():
p.grad = None
H = algorithms.initializeServerState(args, model, trainloader.dataset.num_clients, gradient_at_start, exec_ctx)
logger.info(f'D: {H["D"]} / D_include_frozen : {H["D_include_frozen"]}')
# Append all launch arguments (setup and default)
H["args"] = args
H["raw_cmdline"] = raw_cmdline
H["execution_context"] = exec_ctx
H["total_clients"] = trainloader.dataset.num_clients
H["comment"] = args.comment
H["group-name"] = args.group_name
# ==================================================================================================================
local_optimiser = get_optimiser(model.parameters(), args.local_optimiser, args.local_lr, args.local_momentum,
args.local_weight_decay)
socket = exec_ctx.extra_
state_dicts_thread_safe = Buffer()
while True:
# Main loop for obtaining commands from master
cmd = socket.rawRecvString()
if cmd != "finish_work" and cmd != "non_local_training":
print("Unknown command: ", cmd)
break
if cmd == "finish_work":
# Work termination
break
if cmd == "non_local_training":
# Remote training. Param# 1- all params with local training
msg = socket.rawRecv()
serialized_args = pickle.loads(msg)
(client_state, client_id, msg, model_dict_original, optimiser_dict_original,
model_, train_loader_, criterion, local_optimiser_, device_, round, run_local_iters,
number_of_local_steps, is_rnn, print_freq) = serialized_args
# 1. Update to client specific parameters - compute device, execution context
# 2. Move all tensors into device
device = args.device
client_state['device'] = args.device
client_state['H']["execution_context"] = H["execution_context"]
for k, v in client_state['H'].items():
if torch.is_tensor(v):
client_state['H'][k] = v.to(device=device)
model_dict_original = model_dict_original.to(device=device)
res = local_training(None, client_state, client_id, msg, model_dict_original, optimiser_dict_original,
model, trainloader, criterion, local_optimiser, device, round, run_local_iters,
number_of_local_steps, is_rnn, print_freq, state_dicts_thread_safe,
client_state['H']["run_id"])
# ==========================================================================================================
# Training is finished - response
socket.rawSendString("result_of_local_training")
msg = state_dicts_thread_safe.popFront()
# Temporary remove reference to server state from client
h_ctx = msg['client_state']['H']
# Remove server state from sending
del msg['client_state']['H']
for k, v in msg.items():
if torch.is_tensor(v):
msg[k] = v.cpu()
socket.rawSend(pickle.dumps(msg))
# Revert reference to server state from client
msg['client_state']['H'] = h_ctx
# ==========================================================================================================
def makeBackupOfServerState(H, round):
args = H["args"]
job_id = args.run_id
fname = args.out
if fname == "":
fname = os.path.join(args.checkpoint_dir, f"{args.run_id}_{args.algorithm}_{args.global_lr}_backup.bin")
else:
fname = fname.replace(".bin", f"{args.run_id}_{args.algorithm}_{args.global_lr}_backup.bin")
execution_context = H["execution_context"]
H["execution_context"] = None
exec_ctx_in_arg = ('exec_ctx' in H["args"])
exec_ctx = None
if exec_ctx_in_arg:
exec_ctx = H["args"]['exec_ctx']
H["execution_context"] = None
with open(fname, "wb") as f:
pickle.dump([(job_id, H)], f)
H["execution_context"] = execution_context
if exec_ctx_in_arg:
H["args"]['exec_ctx'] = exec_ctx
def init_and_train_model(args, raw_cmdline, trainloader, testloader, exec_ctx):
logger = Logger.get(args.run_id)
model_dir = create_model_dir(args)
# don't train if setup already exists
if os.path.isdir(model_dir):
logger.info(f"{model_dir} already exists.")
logger.info("Skipping this setup.")
return
# create model directory
os.makedirs(model_dir, exist_ok=True)
# save used args as json to experiment directory
with open(os.path.join(create_model_dir(args), 'args.json'), 'w') as f:
json.dump(vars(args), f, indent=4)
is_rnn = args.model in RNN_MODELS
# TODO: figure out how to exploit DataParallel. Currently - parallelism across workers
model, criterion, current_round = get_training_elements(args.model, args.dataset, trainloader.dataset, args,
args.resume_from, args.load_best, args.device, args.loss,
args.turn_off_batch_normalization_and_dropout)
# ==================================================================================================================
# Reset execution seeds for tunable runtime behaviour
if args.deterministic:
exec_ctx.random.seed(args.manual_runtime_seed)
exec_ctx.np_random.seed(args.manual_runtime_seed)
# ==================================================================================================================
local_optimiser = get_optimiser(model.parameters(), args.local_optimiser, args.local_lr,
args.local_momentum, args.local_weight_decay)
local_scheduler = get_lr_scheduler(local_optimiser, args.rounds, args.local_lr_type)
global_optimiser = get_optimiser(model.parameters(), args.global_optimiser, args.global_lr,
args.global_momentum, args.global_weight_decay)
global_scheduler = get_lr_scheduler(global_optimiser, args.rounds, args.global_lr_type)
metric_to_optim = args.metric
train_time_meter = 0
best_metric = -np.inf
eval_metrics = {}
mutils.print_models_info(model, args)
logger.info(f"Number of parameters in the model: {mutils.number_of_params(model):,d}\n")
gpu_utils.print_info_about_used_gpu(args.device, args.run_id)
sampled_clients = get_sampled_clients(trainloader.dataset.num_clients, args, exec_ctx)
# ============= Initialize worker threads===========================================================================
for th in exec_ctx.local_training_threads.threads:
th.model_copy = copy.deepcopy(model)
# th.model_copy, _ = initialise_model(args.model, args.dataset, trainloader.dataset, args, args.resume_from,
# args.load_best)
# th.model_copy.load_state_dict(copy.deepcopy(model.state_dict()))
th.model_copy = th.model_copy.to(th.device)
th.local_optimiser_copy = get_optimiser(th.model_copy.parameters(), args.local_optimiser, args.local_lr,
args.local_momentum, args.local_weight_decay)
th.local_scheduler = get_lr_scheduler(th.local_optimiser_copy, args.rounds, args.local_lr_type)
for th in exec_ctx.eval_thread_pool.threads:
th.model_copy = copy.deepcopy(model)
# th.model_copy, _ = initialise_model(args.model, args.dataset, trainloader.dataset, args, args.resume_from,
# args.load_best)
# th.model_copy.load_state_dict(copy.deepcopy(model.state_dict()))
th.model_copy = th.model_copy.to(th.device)
for th in exec_ctx.saver_thread.threads:
th.model_copy = copy.deepcopy(model)
# th.model_copy, _ = initialise_model(args.model, args.dataset, trainloader.dataset, args, args.resume_from,
# args.load_best)
# th.model_copy.load_state_dict(copy.deepcopy(model.state_dict()))
th.model_copy = th.model_copy.to(th.device)
exec_ctx.saver_thread.best_metric = -np.inf
exec_ctx.saver_thread.eval_metrics = {}
exec_ctx.saver_thread.best_metric_lock = threading.Lock()
# ==================================================================================================================
optimiser_dict_original = copy.deepcopy(local_optimiser.state_dict())
# ==================================================================================================================
logger.info(f"Use separate {exec_ctx.local_training_threads.workers()} worker threads for make local optimization")
# Initialize server state
if args.initialize_shifts_policy == "full_gradient_at_start":
trainloader.dataset.set_client(None)
algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False,
evaluate_function=False, device=args.device, args=args)
gradient_at_start = mutils.get_gradient(model)
for p in model.parameters():
p.grad = None
else:
D = mutils.number_of_params(model)
gradient_at_start = torch.zeros(D).to(args.device)
H = algorithms.initializeServerState(args, model, trainloader.dataset.num_clients, gradient_at_start, exec_ctx)
logger.info(f'D: {H["D"]} / D_include_frozen : {H["D_include_frozen"]}')
# ==================================================================================================================
# Append all launch arguments (setup and default)
H["args"] = args
H["raw_cmdline"] = raw_cmdline
H["execution_context"] = exec_ctx
H["total_clients"] = trainloader.dataset.num_clients
H["comment"] = args.comment
H["group-name"] = args.group_name
# H["sampled_clients"] = sampled_clients
# Probable pre-scale x0
if algorithms.has_experiment_option(H, "x0_norm"):
H["x0"] = (H["x0"] / H["x0"].norm()) * algorithms.get_experiment_option_f(H, "x0_norm")
# For task where we can obtain good estimation of L
if "L" in dir(trainloader.dataset):
H["L_compute"] = trainloader.dataset.L
if "Li_all_clients" in dir(trainloader.dataset):
H["Li_all_clients"] = trainloader.dataset.Li_all_clients
H["Li_max"] = max(H["Li_all_clients"])
# Initialize metrics
H['best_metric'] = best_metric
H['eval_metrics'] = copy.deepcopy(eval_metrics)
# Initialize starting point of optimization
mutils.set_params(model, H['x0'])
is_target_gpu = gpu_utils.is_target_dev_gpu(args.device)
if execution_context.simulation_start_fn:
execution_context.simulation_start_fn(H)
extra_opts = [a.strip() for a in args.extra_track.split(",")]
fed_dataset_all_clients = np.arange(H['total_clients'])
# Main Algorithm Optimization loop
for i in range(args.rounds):
# Check requirement to force stop simulation
if execution_context.is_simulation_need_earlystop_fn is not None:
if execution_context.is_simulation_need_earlystop_fn(H):
break
# Check that all is ok with current run
if i > 0:
prev_history_stats = H['history'][i - 1]
forceTermination = False
for elem in ["x_before_round", "grad_sgd_server_l2"]:
if elem in prev_history_stats:
if math.isnan(prev_history_stats[elem]) or math.isinf(prev_history_stats[elem]):
logger.error(f"Force early stop due to numerical problems with {elem} at round {i}.")
forceTermination = True
break
if forceTermination:
break
# Update information about current round
H['current_round'] = i
start = time.time()
# Generate client state
# TODO: metrics meter to be used for Tensorboard/WandB
metrics_meter, fed_dataset_clients = run_one_communication_round(H, model, trainloader, criterion,
local_optimiser,
optimiser_dict_original, global_optimiser,
args.device,
current_round, args.run_local_steps,
args.number_of_local_iters,
is_rnn, sampled_clients)
# train_time = time.time() - start
if H['algorithm'] == 'gradskip':
train_time = H['time']
train_time_meter += train_time # Track timings for across epochs average
logger.debug(f'Epoch train time: {train_time}')
# H['history'][current_round]["xi_after"] = mutils.get_params(model)
H['history'][current_round]["train_time"] = train_time
H['history'][current_round]["time"] = time.time() - exec_ctx.context_init_time
if H['algorithm'] == 'gradskip':
if current_round == 0:
H['history'][current_round]["time"] = train_time
else:
H['history'][current_round]["time"] = H['history'][current_round - 1]["time"] + train_time
H['last_round_elapsed_sec'] = train_time
if (i % args.eval_every == 0 or i == (args.rounds - 1)) and metric_to_optim != "none":
# Save results obtained so far into backup file
# ==========================================================================================================
# Serialize server state into filesystem"
makeBackupOfServerState(H, i)
# ==========================================================================================================
# Evaluate model
if args.eval_async_threads > 0:
defered_eval_and_save_checkpoint(model, criterion, args, current_round, is_rnn=is_rnn,
metric_to_optim=metric_to_optim, exec_ctx=exec_ctx)
# Update recent information about eval metrics
exec_ctx.saver_thread.best_metric_lock.acquire()
H['best_metric'] = max(H['best_metric'], exec_ctx.saver_thread.best_metric)
H['eval_metrics'] = copy.deepcopy(exec_ctx.saver_thread.eval_metrics)
exec_ctx.saver_thread.best_metric_lock.release()
else:
metrics = evaluate_model(model, testloader, criterion, args.device, current_round, print_freq=10,
is_rnn=is_rnn, metric_to_optim=metric_to_optim)
avg_metric = metrics[metric_to_optim].get_avg()
cur_metrics = {"loss": metrics["loss"].get_avg(),
"top_1_acc": metrics["top_1_acc"].get_avg(),
"top_5_acc": metrics["top_5_acc"].get_avg(),
"neq_perplexity": metrics["neq_perplexity"].get_avg()
}
eval_metrics.update({metrics['round']: cur_metrics})
# Save model checkpoint
model_filename = '{model}_{run_id}_checkpoint_{round:0>2d}.pth.tar'.format(model=args.model,
run_id=args.run_id,
round=current_round)
is_best = avg_metric > best_metric
save_checkpoint(model, model_filename, is_best=is_best, args=args, metrics=metrics,
metric_to_optim=metric_to_optim)
if is_best:
best_metric = avg_metric
# Update recent information about eval metrics
H['best_metric'] = max(H['best_metric'], best_metric)
H['eval_metrics'] = copy.deepcopy(eval_metrics)
if np.isnan(metrics['loss'].get_avg()):
logger.critical('NaN loss detected, aborting training procedure.')
# return
logger.info(f'Current lrs global:{global_scheduler.get_last_lr()}')
# ==============================================================================================================
# Add number of clients in that round
H['history'][current_round]["number_of_client_in_round"] = len(fed_dataset_clients)
if args.log_gpu_usage:
H['history'][current_round]["memory_gpu_used"] = 0
for dev in args.gpu:
if gpu_utils.is_target_dev_gpu(dev):
memory_gpu_used = torch.cuda.memory_stats(args.device)['reserved_bytes.all.current']
logger.info(
f"GPU: Before round {i} we have used {memory_gpu_used / (1024.0 ** 2):.2f} MB from device {dev}")
H['history'][current_round]["memory_gpu_used"] += memory_gpu_used / (1024.0 ** 2)
if current_round % args.eval_every == 0 or i == (args.rounds - 1):
if "full_gradient_norm_train" in extra_opts and "full_objective_value_train" in extra_opts:
fed_dataset = trainloader.dataset
gradient_avg = mutils.get_zero_gradient_compatible_with_model(model)
fvalue_avg = 0.0
for c in fed_dataset_all_clients:
fed_dataset.set_client(c)
fvalue = algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn,
update_statistics=False, evaluate_function=True,
device=args.device, args=args)
g = mutils.get_gradient(model)
gradient_avg = (gradient_avg * c + g) / (c + 1)
fvalue_avg = (fvalue_avg * c + fvalue) / (c + 1)
H['history'][current_round]["full_gradient_norm_train"] = mutils.l2_norm_of_vec(gradient_avg)
H['history'][current_round]["full_objective_value_train"] = fvalue_avg
if "full_gradient_norm_train" in extra_opts and "full_objective_value_train" not in extra_opts:
fed_dataset = trainloader.dataset
gradient_avg = mutils.get_zero_gradient_compatible_with_model(model)
for c in fed_dataset_all_clients:
fed_dataset.set_client(c)
algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False,
evaluate_function=False, device=args.device, args=args)
g = mutils.get_gradient(model)
gradient_avg = (gradient_avg * c + g) / (c + 1)
H['history'][current_round]["full_gradient_norm_train"] = mutils.l2_norm_of_vec(gradient_avg)
if "full_objective_value_train" in extra_opts and "full_gradient_norm_train" not in extra_opts:
fed_dataset = trainloader.dataset
fvalue_avg = 0.0
for c in fed_dataset_all_clients:
fed_dataset.set_client(c)
fvalue = algorithms.evaluateFunction(None, model, trainloader, criterion, is_rnn,
update_statistics=False, device=args.device, args=args)
fvalue_avg = (fvalue_avg * c + fvalue) / (c + 1)
H['history'][current_round]["full_objective_value_train"] = fvalue_avg
# ==========================================================================================================
if "full_gradient_norm_val" in extra_opts and "full_objective_value_val" in extra_opts:
fvalue = algorithms.evaluateGradient(None, model, testloader, criterion, is_rnn,
update_statistics=False, evaluate_function=True,
device=args.device, args=args)
H['history'][current_round]["full_gradient_norm_val"] = mutils.l2_norm_of_gradient_m(model)
H['history'][current_round]["full_objective_value_val"] = fvalue
if "full_gradient_norm_val" in extra_opts and "full_objective_value_val" not in extra_opts:
algorithms.evaluateGradient(None, model, testloader, criterion, is_rnn, update_statistics=False,
evaluate_function=False, device=args.device, args=args)
H['history'][current_round]["full_gradient_norm_val"] = mutils.l2_norm_of_gradient_m(model)
if "full_objective_value_val" in extra_opts and "full_gradient_norm_val" not in extra_opts:
fvalue = algorithms.evaluateFunction(None, model, testloader, criterion, is_rnn,
update_statistics=False, device=args.device, args=args)
H['history'][current_round]["full_objective_value_val"] = fvalue
# ==========================================================================================================
# Interpolate
# ==========================================================================================================
if current_round > 0:
look_behind = args.eval_every
# ======================================================================================================
if i % args.eval_every == 0:
look_behind = args.eval_every # Exactly look_behind previous evaluation had place to be
elif i == args.rounds - 1:
look_behind = i % args.eval_every # We are in case that we in last round "R-1".
pass # Exactly i%args.eval_every round before previous eval had place to be
else:
assert (not "Impossible case")
# ======================================================================================================
for s in range(1, look_behind):
alpha = s / float(look_behind)
prev_eval_round = current_round - look_behind
inter_round = prev_eval_round + s
# Linearly interpolate between a (alpha=0.0) and b(alpha=1.0)
def lerp(a, b, alpha):
return (1.0 - alpha) * a + alpha * b
if "full_gradient_norm_train" in extra_opts:
H['history'][inter_round]["full_gradient_norm_train"] = lerp(
H['history'][prev_eval_round]["full_gradient_norm_train"],
H['history'][current_round]["full_gradient_norm_train"], alpha)
if "full_objective_value_train" in extra_opts:
H['history'][inter_round]["full_objective_value_train"] = lerp(
H['history'][prev_eval_round]["full_objective_value_train"],
H['history'][current_round]["full_objective_value_train"], alpha)
if "full_gradient_norm_val" in extra_opts:
H['history'][inter_round]["full_gradient_norm_val"] = lerp(
H['history'][prev_eval_round]["full_gradient_norm_val"],
H['history'][current_round]["full_gradient_norm_val"], alpha)
if "full_objective_value_val" in extra_opts:
H['history'][inter_round]["full_objective_value_val"] = lerp(
H['history'][prev_eval_round]["full_objective_value_val"],
H['history'][current_round]["full_objective_value_val"], alpha)
if current_round % args.eval_every == 0 or i == (args.rounds - 1):
if execution_context.simulation_progress_steps_fn is not None:
execution_context.simulation_progress_steps_fn(i / float(args.rounds) * 0.75, H)
wandb_wrapper.logStatistics(H, current_round)
# ==============================================================================================================
# Save parameters of the last model
# if i == args.rounds - 1:
# xfinal = mutils.get_params(model)
# H['xfinal'] = xfinal
# ==============================================================================================================
# Update schedulers
if exec_ctx.local_training_threads.workers() > 0:
# In this mode local optimizers for worker threads are used
for th in exec_ctx.local_training_threads.threads:
th.local_scheduler.step()
else:
# In this mode local optimizer is used
local_scheduler.step()
global_scheduler.step()
# Increment rounder counter
current_round += 1
# Empty caches. Force cleanup the cache after round (good to fix fragmentation issues)
if args.per_round_clean_torch_cache:
execution_context.torch_global_lock.acquire()
torch.cuda.empty_cache()
execution_context.torch_global_lock.release()
xfinal = mutils.get_params(model)
H['xfinal'] = xfinal
logger.debug(f'Average epoch train time: {train_time_meter / args.rounds}')
# ==================================================================================================================
logger.info("Wait for threads from a training threadpool")
exec_ctx.local_training_threads.stop()
if execution_context.simulation_progress_steps_fn is not None:
execution_context.simulation_progress_steps_fn(0.8, H)
logger.info("Wait for threads from a evaluate threadpool")
exec_ctx.eval_thread_pool.stop()
if execution_context.simulation_progress_steps_fn is not None:
execution_context.simulation_progress_steps_fn(0.95, H)
logger.info("Wait for threads from a serialization threadpool")
exec_ctx.saver_thread.stop()
# Not need for use in fact, but we do it for purpose of excluding any bugs inside PyTorch
logger.info("Synchronize all GPU streams")
if is_target_gpu:
torch.cuda.synchronize(args.device)
# Update metrics based on eval and save results
H['best_metric'] = max(H['best_metric'], exec_ctx.saver_thread.best_metric)
H['best_metric'] = max(H['best_metric'], best_metric)
if len(eval_metrics) > 0:
H['eval_metrics'] = copy.deepcopy(eval_metrics)
else:
H['eval_metrics'] = copy.deepcopy(exec_ctx.saver_thread.eval_metrics)
if execution_context.simulation_progress_steps_fn is not None:
execution_context.simulation_progress_steps_fn(1.0, H)
# ==================================================================================================================
# Final prune
HKeys = list(H.keys())
# Remove any tensor field huger then 1 MBytes from final solution
tensor_prune_mb_threshold = 0.0
for field in HKeys:
if torch.is_tensor(H[field]):
size_in_mbytes = H[field].numel() * H[field].element_size() / (1024.0 ** 2)
if size_in_mbytes > tensor_prune_mb_threshold:
del H[field]
continue
else:
# Convert any not removed tensor to CPU
H[field] = H[field].cpu()
if type(H[field]) == list:
sz = len(H[field])
i = 0
while i < sz:
if torch.is_tensor(H[field][i]):
size_in_mbytes = H[field][i].numel() * H[field][i].element_size() / (1024.0 ** 2)
if size_in_mbytes >= tensor_prune_mb_threshold:
del H[field][i]
sz = len(H[field])
else:
# Convert any not removed tensor to CPU
H[field][i] = H[field][i].cpu()
i += 1
else:
i += 1
# Remove local optimiser state, client_compressors, reference to server state
not_to_remove_from_client_state = ["error"]
for round, history_state in H['history'].items():
clients_history = history_state['client_states']
for client_id, client_state in clients_history.items():
client_state = client_state['client_state']
if 'H' in client_state:
del client_state['H']
if 'optimiser' in client_state:
del client_state['optimiser']
if 'buffers' in client_state:
del client_state['buffers']
if 'client_compressor' in client_state:
del client_state['client_compressor']
client_state_keys = list(client_state.keys())
for field in client_state_keys:
if field in not_to_remove_from_client_state:
continue
if torch.is_tensor(client_state[field]):
size_in_mbytes = client_state[field].numel() * client_state[field].element_size() / (1024.0 ** 2)
if size_in_mbytes > tensor_prune_mb_threshold:
del client_state[field]
continue
else:
# Convert any not removed tensor to CPU
client_state[field] = client_state[field].cpu()
continue
if type(client_state[field]) == list:
sz = len(client_state[field])
i = 0
while i < sz:
if torch.is_tensor(client_state[field][i]):
size_in_mbytes = client_state[field][i].numel() * client_state[field][i].element_size() / (
1024.0 ** 2)
if size_in_mbytes >= tensor_prune_mb_threshold:
del client_state[field][i]
sz = len(client_state[field])
else:
client_state[field][i] = client_state[field][i].cpu()
i += 1
else:
i += 1
# Remove reference for execution context
del H["execution_context"]
if 'exec_ctx' in H["args"]:
del H["args"]
# ==================================================================================================================
# Cleanup execution context resources
# ==================================================================================================================
execution_context.resetExecutionContext(exec_ctx)
# ==================================================================================================================
if execution_context.simulation_finish_fn is not None:
execution_context.simulation_finish_fn(H)
def runSimulation(cmdline, extra_=None):
# ==================================================================================================================
# ENTRY POINT FOR LAUNCH SIMULATION FROM GUI
# ==================================================================================================================
global CUDA_SUPPORT
args = parse_args(cmdline)
if args.hostname == "":
args.hostname = socket.gethostname()
Logger.setup_logging(args.loglevel, args.logfilter, logfile=args.logfile)
logger = Logger.get(args.run_id)
logger.info(f"CLI args original: {cmdline}")
if torch.cuda.device_count():
CUDA_SUPPORT = True
# torch.backends.cudnn.benchmark = True
else:
logger.warning('CUDA unsupported!!')
CUDA_SUPPORT = False
if not CUDA_SUPPORT:
args.gpu = "cpu"
if args.deterministic:
import torch.backends.cudnn as cudnn
import os
import random
# TODO: This settings are not thread safe in case of using cuda backend from several threads
# Project use execution context random generators for random and numpy in thread safe way
if CUDA_SUPPORT:
cudnn.deterministic = args.deterministic
cudnn.benchmark = not args.deterministic
torch.cuda.manual_seed(args.manual_init_seed)
torch.cuda.manual_seed_all(args.manual_init_seed)
# Turn off Tensor Cores if have been requested
torch.backends.cudnn.allow_tf32 = args.allow_use_nv_tensorcores
torch.backends.cuda.matmul.allow_tf32 = args.allow_use_nv_tensorcores
torch.manual_seed(args.manual_init_seed)
random.seed(args.manual_init_seed)
np.random.seed(args.manual_init_seed)
os.environ['PYTHONHASHSEED'] = str(args.manual_init_seed)
torch.backends.cudnn.deterministic = True
main(args, cmdline, extra_)
# ======================================================================================================================
g_Terminate = False
# ======================================================================================================================
def terminationWithSignal(*args):
global g_Terminate
g_Terminate = True
print("Request for terminate process has been obtained. Save results and terminate.")
def isSimulationNeedEarlyStopCmdLine(H):
global g_Terminate
return g_Terminate
# ======================================================================================================================
def saveResult(H):
"""Serialize server state into filesystem"""
args = H["args"]
job_id = args.run_id
fname = args.out
if fname == "":
fname = os.path.join(args.checkpoint_dir, f"{args.algorithm}_simulation_history.bin")
# Save formally list of tuples (job_id, server_state)
with open(fname, "wb") as f:
pickle.dump([(job_id, H)], f)
print(f"Final server state for {job_id} is serialized into: '{fname}'")
print(f"Program which uses '{args.algorithm}' algorithm for {args.rounds} rounds was finished")
class ClientThread(threading.Thread):
def __init__(self, clientSocket, listenPort):
threading.Thread.__init__(self)
self.clientSocket = clientSocket # OS-like socket to communicate with information
self.comSocket = comm_socket.CommSocket(clientSocket) # Light construction on top of socket to communicate
self.listenPort = listenPort # Listening port
def run(self):
cmd = self.comSocket.rawRecvString() # Obtain command
t = time.localtime()
cur_time = time.strftime("%H:%M:%S", t)
print(f"-- {cur_time}: The received command: {cmd}")
if cmd == "list_of_gpus": # List of resources in the system
gpus = len(gpu_utils.get_available_gpus())
self.comSocket.rawSendString(f"{gpus}")
elif cmd == "execute_work": # Execute work
devConfig = self.comSocket.rawRecvString() # 1-st param: device configuration
cmdLine = self.comSocket.rawRecv() # 2-nd param: command line
cmdLine = pickle.loads(cmdLine) # unpack cmdline
ip, port, dev_name, dev_number = devConfig.split(":")
# Remove unnecessary flags for worker
i = len(cmdLine) - 1
while i >= 0:
if cmdLine[i] == "--metric":
cmdLine[i + 1] = "none"
if cmdLine[i] == "--gpu":
cmdLine[i + 1] = dev_number # update device number
if cmdLine[i] in ['--wandb-key', '--eval-async-threads', '--save-async-threads',
'--threadpool-for-local-opt', '--external-devices', '--evaluate']:
del cmdLine[i] # remove flag
del cmdLine[i] # remove value for flag (originally at index i+1)
i -= 1
# Specify listen mode in which worker is
cmdLine.append("--worker-listen-mode")
cmdLine.append(str(self.listenPort))
runSimulation(cmdLine, extra_=self.comSocket)
else:
print(f"The received command {cmd} is not valid")
def executeRemoteWorkerSupportMode(listenPort):
"""
Listen for a specified port and wait for incoming connection. Maximum number of pending connection is 5.
"""
serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
serversocket.bind(('0.0.0.0', listenPort))
MAX_CONNECTIONS = 5
serversocket.listen(MAX_CONNECTIONS)
while True:
t = time.localtime()
cur_time = time.strftime("%H:%M:%S", t)
print(f"-- {cur_time}: Waiting for commands from the master by worker {socket.gethostname()}:{listenPort}")
# accept connections from outside and process it in separate threads
(clientsocket, address) = serversocket.accept()
ct = ClientThread(clientsocket, listenPort)
ct.run()
if __name__ == "__main__":
# ==================================================================================================================
# ENTRY POINT FOR CUI
# ==================================================================================================================
# Worker listener mode
for i in range(len(sys.argv)):
if sys.argv[i] == "--worker-listen-mode":
print(get_pretty_env_info())
print("")
executeRemoteWorkerSupportMode(int(sys.argv[i + 1]))
sys.exit(0)
# ==================================================================================================================
# Usual mode
# Setup signal for CTRL+C(SIGINT) or SIGTERM
signal.signal(signal.SIGINT, terminationWithSignal)
signal.signal(signal.SIGTERM, terminationWithSignal)
execution_context.is_simulation_need_earlystop_fn = isSimulationNeedEarlyStopCmdLine
execution_context.simulation_finish_fn = saveResult
runSimulation(sys.argv[1:])
sys.exit(0)
# ==================================================================================================================
Functions
def executeRemoteWorkerSupportMode(listenPort)
-
Listen for a specified port and wait for incoming connection. Maximum number of pending connection is 5.
Expand source code
def executeRemoteWorkerSupportMode(listenPort): """ Listen for a specified port and wait for incoming connection. Maximum number of pending connection is 5. """ serversocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) serversocket.bind(('0.0.0.0', listenPort)) MAX_CONNECTIONS = 5 serversocket.listen(MAX_CONNECTIONS) while True: t = time.localtime() cur_time = time.strftime("%H:%M:%S", t) print(f"-- {cur_time}: Waiting for commands from the master by worker {socket.gethostname()}:{listenPort}") # accept connections from outside and process it in separate threads (clientsocket, address) = serversocket.accept() ct = ClientThread(clientsocket, listenPort) ct.run()
def init_and_help_with_compute(args, raw_cmdline, trainloader, testloader, exec_ctx)
-
Expand source code
def init_and_help_with_compute(args, raw_cmdline, trainloader, testloader, exec_ctx): logger = Logger.get(args.run_id) model_dir = create_model_dir(args) # don't train if setup already exists if os.path.isdir(model_dir): logger.info(f"{model_dir} already exists.") logger.info("Skipping this setup.") return # create model directory os.makedirs(model_dir, exist_ok=True) # save used args as json to experiment directory with open(os.path.join(create_model_dir(args), 'args.json'), 'w') as f: json.dump(vars(args), f, indent=4) is_rnn = args.model in RNN_MODELS # TODO: figure out how to exploit DataParallel. Currently - parallelism across workers model, criterion, current_round = get_training_elements(args.model, args.dataset, trainloader.dataset, args, args.resume_from, args.load_best, args.device, args.loss, args.turn_off_batch_normalization_and_dropout) # ================================================================================================================== # Reset execution seeds for tunable runtime behaviour if args.deterministic: exec_ctx.random.seed(args.manual_runtime_seed) exec_ctx.np_random.seed(args.manual_runtime_seed) # ================================================================================================================== mutils.print_models_info(model, args) logger.info(f"Number of parameters in the model: {mutils.number_of_params(model):,d}\n") gpu_utils.print_info_about_used_gpu(args.device, args.run_id) # ================================================================================================================== # Initialize server state # ================================================================================================================== trainloader.dataset.set_client(None) algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False, evaluate_function=False, device=args.device, args=args) gradient_at_start = mutils.get_gradient(model) for p in model.parameters(): p.grad = None H = algorithms.initializeServerState(args, model, trainloader.dataset.num_clients, gradient_at_start, exec_ctx) logger.info(f'D: {H["D"]} / D_include_frozen : {H["D_include_frozen"]}') # Append all launch arguments (setup and default) H["args"] = args H["raw_cmdline"] = raw_cmdline H["execution_context"] = exec_ctx H["total_clients"] = trainloader.dataset.num_clients H["comment"] = args.comment H["group-name"] = args.group_name # ================================================================================================================== local_optimiser = get_optimiser(model.parameters(), args.local_optimiser, args.local_lr, args.local_momentum, args.local_weight_decay) socket = exec_ctx.extra_ state_dicts_thread_safe = Buffer() while True: # Main loop for obtaining commands from master cmd = socket.rawRecvString() if cmd != "finish_work" and cmd != "non_local_training": print("Unknown command: ", cmd) break if cmd == "finish_work": # Work termination break if cmd == "non_local_training": # Remote training. Param# 1- all params with local training msg = socket.rawRecv() serialized_args = pickle.loads(msg) (client_state, client_id, msg, model_dict_original, optimiser_dict_original, model_, train_loader_, criterion, local_optimiser_, device_, round, run_local_iters, number_of_local_steps, is_rnn, print_freq) = serialized_args # 1. Update to client specific parameters - compute device, execution context # 2. Move all tensors into device device = args.device client_state['device'] = args.device client_state['H']["execution_context"] = H["execution_context"] for k, v in client_state['H'].items(): if torch.is_tensor(v): client_state['H'][k] = v.to(device=device) model_dict_original = model_dict_original.to(device=device) res = local_training(None, client_state, client_id, msg, model_dict_original, optimiser_dict_original, model, trainloader, criterion, local_optimiser, device, round, run_local_iters, number_of_local_steps, is_rnn, print_freq, state_dicts_thread_safe, client_state['H']["run_id"]) # ========================================================================================================== # Training is finished - response socket.rawSendString("result_of_local_training") msg = state_dicts_thread_safe.popFront() # Temporary remove reference to server state from client h_ctx = msg['client_state']['H'] # Remove server state from sending del msg['client_state']['H'] for k, v in msg.items(): if torch.is_tensor(v): msg[k] = v.cpu() socket.rawSend(pickle.dumps(msg)) # Revert reference to server state from client msg['client_state']['H'] = h_ctx # ==========================================================================================================
def init_and_train_model(args, raw_cmdline, trainloader, testloader, exec_ctx)
-
Expand source code
def init_and_train_model(args, raw_cmdline, trainloader, testloader, exec_ctx): logger = Logger.get(args.run_id) model_dir = create_model_dir(args) # don't train if setup already exists if os.path.isdir(model_dir): logger.info(f"{model_dir} already exists.") logger.info("Skipping this setup.") return # create model directory os.makedirs(model_dir, exist_ok=True) # save used args as json to experiment directory with open(os.path.join(create_model_dir(args), 'args.json'), 'w') as f: json.dump(vars(args), f, indent=4) is_rnn = args.model in RNN_MODELS # TODO: figure out how to exploit DataParallel. Currently - parallelism across workers model, criterion, current_round = get_training_elements(args.model, args.dataset, trainloader.dataset, args, args.resume_from, args.load_best, args.device, args.loss, args.turn_off_batch_normalization_and_dropout) # ================================================================================================================== # Reset execution seeds for tunable runtime behaviour if args.deterministic: exec_ctx.random.seed(args.manual_runtime_seed) exec_ctx.np_random.seed(args.manual_runtime_seed) # ================================================================================================================== local_optimiser = get_optimiser(model.parameters(), args.local_optimiser, args.local_lr, args.local_momentum, args.local_weight_decay) local_scheduler = get_lr_scheduler(local_optimiser, args.rounds, args.local_lr_type) global_optimiser = get_optimiser(model.parameters(), args.global_optimiser, args.global_lr, args.global_momentum, args.global_weight_decay) global_scheduler = get_lr_scheduler(global_optimiser, args.rounds, args.global_lr_type) metric_to_optim = args.metric train_time_meter = 0 best_metric = -np.inf eval_metrics = {} mutils.print_models_info(model, args) logger.info(f"Number of parameters in the model: {mutils.number_of_params(model):,d}\n") gpu_utils.print_info_about_used_gpu(args.device, args.run_id) sampled_clients = get_sampled_clients(trainloader.dataset.num_clients, args, exec_ctx) # ============= Initialize worker threads=========================================================================== for th in exec_ctx.local_training_threads.threads: th.model_copy = copy.deepcopy(model) # th.model_copy, _ = initialise_model(args.model, args.dataset, trainloader.dataset, args, args.resume_from, # args.load_best) # th.model_copy.load_state_dict(copy.deepcopy(model.state_dict())) th.model_copy = th.model_copy.to(th.device) th.local_optimiser_copy = get_optimiser(th.model_copy.parameters(), args.local_optimiser, args.local_lr, args.local_momentum, args.local_weight_decay) th.local_scheduler = get_lr_scheduler(th.local_optimiser_copy, args.rounds, args.local_lr_type) for th in exec_ctx.eval_thread_pool.threads: th.model_copy = copy.deepcopy(model) # th.model_copy, _ = initialise_model(args.model, args.dataset, trainloader.dataset, args, args.resume_from, # args.load_best) # th.model_copy.load_state_dict(copy.deepcopy(model.state_dict())) th.model_copy = th.model_copy.to(th.device) for th in exec_ctx.saver_thread.threads: th.model_copy = copy.deepcopy(model) # th.model_copy, _ = initialise_model(args.model, args.dataset, trainloader.dataset, args, args.resume_from, # args.load_best) # th.model_copy.load_state_dict(copy.deepcopy(model.state_dict())) th.model_copy = th.model_copy.to(th.device) exec_ctx.saver_thread.best_metric = -np.inf exec_ctx.saver_thread.eval_metrics = {} exec_ctx.saver_thread.best_metric_lock = threading.Lock() # ================================================================================================================== optimiser_dict_original = copy.deepcopy(local_optimiser.state_dict()) # ================================================================================================================== logger.info(f"Use separate {exec_ctx.local_training_threads.workers()} worker threads for make local optimization") # Initialize server state if args.initialize_shifts_policy == "full_gradient_at_start": trainloader.dataset.set_client(None) algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False, evaluate_function=False, device=args.device, args=args) gradient_at_start = mutils.get_gradient(model) for p in model.parameters(): p.grad = None else: D = mutils.number_of_params(model) gradient_at_start = torch.zeros(D).to(args.device) H = algorithms.initializeServerState(args, model, trainloader.dataset.num_clients, gradient_at_start, exec_ctx) logger.info(f'D: {H["D"]} / D_include_frozen : {H["D_include_frozen"]}') # ================================================================================================================== # Append all launch arguments (setup and default) H["args"] = args H["raw_cmdline"] = raw_cmdline H["execution_context"] = exec_ctx H["total_clients"] = trainloader.dataset.num_clients H["comment"] = args.comment H["group-name"] = args.group_name # H["sampled_clients"] = sampled_clients # Probable pre-scale x0 if algorithms.has_experiment_option(H, "x0_norm"): H["x0"] = (H["x0"] / H["x0"].norm()) * algorithms.get_experiment_option_f(H, "x0_norm") # For task where we can obtain good estimation of L if "L" in dir(trainloader.dataset): H["L_compute"] = trainloader.dataset.L if "Li_all_clients" in dir(trainloader.dataset): H["Li_all_clients"] = trainloader.dataset.Li_all_clients H["Li_max"] = max(H["Li_all_clients"]) # Initialize metrics H['best_metric'] = best_metric H['eval_metrics'] = copy.deepcopy(eval_metrics) # Initialize starting point of optimization mutils.set_params(model, H['x0']) is_target_gpu = gpu_utils.is_target_dev_gpu(args.device) if execution_context.simulation_start_fn: execution_context.simulation_start_fn(H) extra_opts = [a.strip() for a in args.extra_track.split(",")] fed_dataset_all_clients = np.arange(H['total_clients']) # Main Algorithm Optimization loop for i in range(args.rounds): # Check requirement to force stop simulation if execution_context.is_simulation_need_earlystop_fn is not None: if execution_context.is_simulation_need_earlystop_fn(H): break # Check that all is ok with current run if i > 0: prev_history_stats = H['history'][i - 1] forceTermination = False for elem in ["x_before_round", "grad_sgd_server_l2"]: if elem in prev_history_stats: if math.isnan(prev_history_stats[elem]) or math.isinf(prev_history_stats[elem]): logger.error(f"Force early stop due to numerical problems with {elem} at round {i}.") forceTermination = True break if forceTermination: break # Update information about current round H['current_round'] = i start = time.time() # Generate client state # TODO: metrics meter to be used for Tensorboard/WandB metrics_meter, fed_dataset_clients = run_one_communication_round(H, model, trainloader, criterion, local_optimiser, optimiser_dict_original, global_optimiser, args.device, current_round, args.run_local_steps, args.number_of_local_iters, is_rnn, sampled_clients) # train_time = time.time() - start if H['algorithm'] == 'gradskip': train_time = H['time'] train_time_meter += train_time # Track timings for across epochs average logger.debug(f'Epoch train time: {train_time}') # H['history'][current_round]["xi_after"] = mutils.get_params(model) H['history'][current_round]["train_time"] = train_time H['history'][current_round]["time"] = time.time() - exec_ctx.context_init_time if H['algorithm'] == 'gradskip': if current_round == 0: H['history'][current_round]["time"] = train_time else: H['history'][current_round]["time"] = H['history'][current_round - 1]["time"] + train_time H['last_round_elapsed_sec'] = train_time if (i % args.eval_every == 0 or i == (args.rounds - 1)) and metric_to_optim != "none": # Save results obtained so far into backup file # ========================================================================================================== # Serialize server state into filesystem" makeBackupOfServerState(H, i) # ========================================================================================================== # Evaluate model if args.eval_async_threads > 0: defered_eval_and_save_checkpoint(model, criterion, args, current_round, is_rnn=is_rnn, metric_to_optim=metric_to_optim, exec_ctx=exec_ctx) # Update recent information about eval metrics exec_ctx.saver_thread.best_metric_lock.acquire() H['best_metric'] = max(H['best_metric'], exec_ctx.saver_thread.best_metric) H['eval_metrics'] = copy.deepcopy(exec_ctx.saver_thread.eval_metrics) exec_ctx.saver_thread.best_metric_lock.release() else: metrics = evaluate_model(model, testloader, criterion, args.device, current_round, print_freq=10, is_rnn=is_rnn, metric_to_optim=metric_to_optim) avg_metric = metrics[metric_to_optim].get_avg() cur_metrics = {"loss": metrics["loss"].get_avg(), "top_1_acc": metrics["top_1_acc"].get_avg(), "top_5_acc": metrics["top_5_acc"].get_avg(), "neq_perplexity": metrics["neq_perplexity"].get_avg() } eval_metrics.update({metrics['round']: cur_metrics}) # Save model checkpoint model_filename = '{model}_{run_id}_checkpoint_{round:0>2d}.pth.tar'.format(model=args.model, run_id=args.run_id, round=current_round) is_best = avg_metric > best_metric save_checkpoint(model, model_filename, is_best=is_best, args=args, metrics=metrics, metric_to_optim=metric_to_optim) if is_best: best_metric = avg_metric # Update recent information about eval metrics H['best_metric'] = max(H['best_metric'], best_metric) H['eval_metrics'] = copy.deepcopy(eval_metrics) if np.isnan(metrics['loss'].get_avg()): logger.critical('NaN loss detected, aborting training procedure.') # return logger.info(f'Current lrs global:{global_scheduler.get_last_lr()}') # ============================================================================================================== # Add number of clients in that round H['history'][current_round]["number_of_client_in_round"] = len(fed_dataset_clients) if args.log_gpu_usage: H['history'][current_round]["memory_gpu_used"] = 0 for dev in args.gpu: if gpu_utils.is_target_dev_gpu(dev): memory_gpu_used = torch.cuda.memory_stats(args.device)['reserved_bytes.all.current'] logger.info( f"GPU: Before round {i} we have used {memory_gpu_used / (1024.0 ** 2):.2f} MB from device {dev}") H['history'][current_round]["memory_gpu_used"] += memory_gpu_used / (1024.0 ** 2) if current_round % args.eval_every == 0 or i == (args.rounds - 1): if "full_gradient_norm_train" in extra_opts and "full_objective_value_train" in extra_opts: fed_dataset = trainloader.dataset gradient_avg = mutils.get_zero_gradient_compatible_with_model(model) fvalue_avg = 0.0 for c in fed_dataset_all_clients: fed_dataset.set_client(c) fvalue = algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False, evaluate_function=True, device=args.device, args=args) g = mutils.get_gradient(model) gradient_avg = (gradient_avg * c + g) / (c + 1) fvalue_avg = (fvalue_avg * c + fvalue) / (c + 1) H['history'][current_round]["full_gradient_norm_train"] = mutils.l2_norm_of_vec(gradient_avg) H['history'][current_round]["full_objective_value_train"] = fvalue_avg if "full_gradient_norm_train" in extra_opts and "full_objective_value_train" not in extra_opts: fed_dataset = trainloader.dataset gradient_avg = mutils.get_zero_gradient_compatible_with_model(model) for c in fed_dataset_all_clients: fed_dataset.set_client(c) algorithms.evaluateGradient(None, model, trainloader, criterion, is_rnn, update_statistics=False, evaluate_function=False, device=args.device, args=args) g = mutils.get_gradient(model) gradient_avg = (gradient_avg * c + g) / (c + 1) H['history'][current_round]["full_gradient_norm_train"] = mutils.l2_norm_of_vec(gradient_avg) if "full_objective_value_train" in extra_opts and "full_gradient_norm_train" not in extra_opts: fed_dataset = trainloader.dataset fvalue_avg = 0.0 for c in fed_dataset_all_clients: fed_dataset.set_client(c) fvalue = algorithms.evaluateFunction(None, model, trainloader, criterion, is_rnn, update_statistics=False, device=args.device, args=args) fvalue_avg = (fvalue_avg * c + fvalue) / (c + 1) H['history'][current_round]["full_objective_value_train"] = fvalue_avg # ========================================================================================================== if "full_gradient_norm_val" in extra_opts and "full_objective_value_val" in extra_opts: fvalue = algorithms.evaluateGradient(None, model, testloader, criterion, is_rnn, update_statistics=False, evaluate_function=True, device=args.device, args=args) H['history'][current_round]["full_gradient_norm_val"] = mutils.l2_norm_of_gradient_m(model) H['history'][current_round]["full_objective_value_val"] = fvalue if "full_gradient_norm_val" in extra_opts and "full_objective_value_val" not in extra_opts: algorithms.evaluateGradient(None, model, testloader, criterion, is_rnn, update_statistics=False, evaluate_function=False, device=args.device, args=args) H['history'][current_round]["full_gradient_norm_val"] = mutils.l2_norm_of_gradient_m(model) if "full_objective_value_val" in extra_opts and "full_gradient_norm_val" not in extra_opts: fvalue = algorithms.evaluateFunction(None, model, testloader, criterion, is_rnn, update_statistics=False, device=args.device, args=args) H['history'][current_round]["full_objective_value_val"] = fvalue # ========================================================================================================== # Interpolate # ========================================================================================================== if current_round > 0: look_behind = args.eval_every # ====================================================================================================== if i % args.eval_every == 0: look_behind = args.eval_every # Exactly look_behind previous evaluation had place to be elif i == args.rounds - 1: look_behind = i % args.eval_every # We are in case that we in last round "R-1". pass # Exactly i%args.eval_every round before previous eval had place to be else: assert (not "Impossible case") # ====================================================================================================== for s in range(1, look_behind): alpha = s / float(look_behind) prev_eval_round = current_round - look_behind inter_round = prev_eval_round + s # Linearly interpolate between a (alpha=0.0) and b(alpha=1.0) def lerp(a, b, alpha): return (1.0 - alpha) * a + alpha * b if "full_gradient_norm_train" in extra_opts: H['history'][inter_round]["full_gradient_norm_train"] = lerp( H['history'][prev_eval_round]["full_gradient_norm_train"], H['history'][current_round]["full_gradient_norm_train"], alpha) if "full_objective_value_train" in extra_opts: H['history'][inter_round]["full_objective_value_train"] = lerp( H['history'][prev_eval_round]["full_objective_value_train"], H['history'][current_round]["full_objective_value_train"], alpha) if "full_gradient_norm_val" in extra_opts: H['history'][inter_round]["full_gradient_norm_val"] = lerp( H['history'][prev_eval_round]["full_gradient_norm_val"], H['history'][current_round]["full_gradient_norm_val"], alpha) if "full_objective_value_val" in extra_opts: H['history'][inter_round]["full_objective_value_val"] = lerp( H['history'][prev_eval_round]["full_objective_value_val"], H['history'][current_round]["full_objective_value_val"], alpha) if current_round % args.eval_every == 0 or i == (args.rounds - 1): if execution_context.simulation_progress_steps_fn is not None: execution_context.simulation_progress_steps_fn(i / float(args.rounds) * 0.75, H) wandb_wrapper.logStatistics(H, current_round) # ============================================================================================================== # Save parameters of the last model # if i == args.rounds - 1: # xfinal = mutils.get_params(model) # H['xfinal'] = xfinal # ============================================================================================================== # Update schedulers if exec_ctx.local_training_threads.workers() > 0: # In this mode local optimizers for worker threads are used for th in exec_ctx.local_training_threads.threads: th.local_scheduler.step() else: # In this mode local optimizer is used local_scheduler.step() global_scheduler.step() # Increment rounder counter current_round += 1 # Empty caches. Force cleanup the cache after round (good to fix fragmentation issues) if args.per_round_clean_torch_cache: execution_context.torch_global_lock.acquire() torch.cuda.empty_cache() execution_context.torch_global_lock.release() xfinal = mutils.get_params(model) H['xfinal'] = xfinal logger.debug(f'Average epoch train time: {train_time_meter / args.rounds}') # ================================================================================================================== logger.info("Wait for threads from a training threadpool") exec_ctx.local_training_threads.stop() if execution_context.simulation_progress_steps_fn is not None: execution_context.simulation_progress_steps_fn(0.8, H) logger.info("Wait for threads from a evaluate threadpool") exec_ctx.eval_thread_pool.stop() if execution_context.simulation_progress_steps_fn is not None: execution_context.simulation_progress_steps_fn(0.95, H) logger.info("Wait for threads from a serialization threadpool") exec_ctx.saver_thread.stop() # Not need for use in fact, but we do it for purpose of excluding any bugs inside PyTorch logger.info("Synchronize all GPU streams") if is_target_gpu: torch.cuda.synchronize(args.device) # Update metrics based on eval and save results H['best_metric'] = max(H['best_metric'], exec_ctx.saver_thread.best_metric) H['best_metric'] = max(H['best_metric'], best_metric) if len(eval_metrics) > 0: H['eval_metrics'] = copy.deepcopy(eval_metrics) else: H['eval_metrics'] = copy.deepcopy(exec_ctx.saver_thread.eval_metrics) if execution_context.simulation_progress_steps_fn is not None: execution_context.simulation_progress_steps_fn(1.0, H) # ================================================================================================================== # Final prune HKeys = list(H.keys()) # Remove any tensor field huger then 1 MBytes from final solution tensor_prune_mb_threshold = 0.0 for field in HKeys: if torch.is_tensor(H[field]): size_in_mbytes = H[field].numel() * H[field].element_size() / (1024.0 ** 2) if size_in_mbytes > tensor_prune_mb_threshold: del H[field] continue else: # Convert any not removed tensor to CPU H[field] = H[field].cpu() if type(H[field]) == list: sz = len(H[field]) i = 0 while i < sz: if torch.is_tensor(H[field][i]): size_in_mbytes = H[field][i].numel() * H[field][i].element_size() / (1024.0 ** 2) if size_in_mbytes >= tensor_prune_mb_threshold: del H[field][i] sz = len(H[field]) else: # Convert any not removed tensor to CPU H[field][i] = H[field][i].cpu() i += 1 else: i += 1 # Remove local optimiser state, client_compressors, reference to server state not_to_remove_from_client_state = ["error"] for round, history_state in H['history'].items(): clients_history = history_state['client_states'] for client_id, client_state in clients_history.items(): client_state = client_state['client_state'] if 'H' in client_state: del client_state['H'] if 'optimiser' in client_state: del client_state['optimiser'] if 'buffers' in client_state: del client_state['buffers'] if 'client_compressor' in client_state: del client_state['client_compressor'] client_state_keys = list(client_state.keys()) for field in client_state_keys: if field in not_to_remove_from_client_state: continue if torch.is_tensor(client_state[field]): size_in_mbytes = client_state[field].numel() * client_state[field].element_size() / (1024.0 ** 2) if size_in_mbytes > tensor_prune_mb_threshold: del client_state[field] continue else: # Convert any not removed tensor to CPU client_state[field] = client_state[field].cpu() continue if type(client_state[field]) == list: sz = len(client_state[field]) i = 0 while i < sz: if torch.is_tensor(client_state[field][i]): size_in_mbytes = client_state[field][i].numel() * client_state[field][i].element_size() / ( 1024.0 ** 2) if size_in_mbytes >= tensor_prune_mb_threshold: del client_state[field][i] sz = len(client_state[field]) else: client_state[field][i] = client_state[field][i].cpu() i += 1 else: i += 1 # Remove reference for execution context del H["execution_context"] if 'exec_ctx' in H["args"]: del H["args"] # ================================================================================================================== # Cleanup execution context resources # ================================================================================================================== execution_context.resetExecutionContext(exec_ctx) # ================================================================================================================== if execution_context.simulation_finish_fn is not None: execution_context.simulation_finish_fn(H)
def isSimulationNeedEarlyStopCmdLine(H)
-
Expand source code
def isSimulationNeedEarlyStopCmdLine(H): global g_Terminate return g_Terminate
def main(args, raw_cmdline, extra_)
-
Expand source code
def main(args, raw_cmdline, extra_): # Init project and save in args url for the plots projectWB = None if len(args.wandb_key) > 0: projectWB = wandb_wrapper.initWandbProject( args.wandb_key, args.wandb_project_name, f"{args.algorithm}-rounds-{args.rounds}-{args.model}@{args.dataset}-runid-{args.run_id}", args ) if projectWB is not None: args.wandb_url = projectWB.url # In case of DataParallel for .to() to # This a device used by master for FL algorithms args.device = gpu_utils.get_target_device_str(args.gpu[0]) if type(args.gpu) == list else args.gpu # Instantiate execution context exec_ctx = execution_context.initExecutionContext() exec_ctx.extra_ = extra_ # Set all seeds if args.deterministic: exec_ctx.random.seed(args.manual_init_seed) exec_ctx.np_random.seed(args.manual_init_seed) # Setup extra options for experiments for option in args.algorithm_options.split(","): kv = option.split(":") if len(kv) == 1: exec_ctx.experimental_options[kv[0]] = True else: exec_ctx.experimental_options[kv[0]] = kv[1] # Load validation set trainset, testset = load_data(exec_ctx, args.data_path, args.dataset, args, load_trainset=True, download=True) test_batch_size = get_test_batch_size(args.dataset, args.batch_size) # Test set dataloader testloader = torch.utils.data.DataLoader(testset, batch_size=test_batch_size, num_workers=args.num_workers_test, shuffle=False, pin_memory=False, drop_last=False ) # Reset all seeds if args.deterministic: exec_ctx.random.seed(args.manual_init_seed) exec_ctx.np_random.seed(args.manual_init_seed) if not args.evaluate: # Training mode trainloader = torch.utils.data.DataLoader( trainset, batch_size=args.batch_size, num_workers=args.num_workers_train, shuffle=False, pin_memory=False, drop_last=False) # Is target device for master is GPU is_target_dev_gpu = gpu_utils.is_target_dev_gpu(args.device) # Initialize local training threads in case of using multithreading implementation exec_ctx.local_training_threads.adjust_num_workers(w=args.threadpool_for_local_opt, own_cuda_stream=True, device_list=args.gpu) for th in exec_ctx.local_training_threads.threads: th.trainset_copy = copy.deepcopy(trainset) th.train_loader_copy = torch.utils.data.DataLoader(th.trainset_copy, batch_size=args.batch_size, num_workers=args.num_workers_train, shuffle=False, pin_memory=False, drop_last=False) if hasattr(th.trainset_copy, "load_data"): th.trainset_copy.load_data() # Initialize remote clients - one remote connection per one remote computation device external_devices = [d.strip() for d in args.external_devices.split(",") if len(d.strip()) > 0] exec_ctx.remote_training_threads.adjust_num_workers(w=len(external_devices), own_cuda_stream=False, device_list=["cpu"]) for th_i in range(len(exec_ctx.remote_training_threads.threads)): th = exec_ctx.remote_training_threads.threads[th_i] dev = external_devices[th_i] th.external_dev_long = dev th.external_ip, th.external_port, th.external_dev, th.external_dev_dumber = dev.split(":") try: th.external_socket = comm_socket.CommSocket() th.external_socket.sock.connect((th.external_ip, int(th.external_port))) # connect with remote side th.external_socket.rawSendString("execute_work") # initiate work execution COMMAND th.external_socket.rawSendString(th.external_dev_long) # provide device specification th.external_socket.rawSend(pickle.dumps( raw_cmdline)) # provide original command line (it will be changed by the client) to be consistent with used devices th.external_socket_online = True except socket.error: th.external_socket_online = False # Initialize saver treads for information serialization (0 is acceptable) exec_ctx.saver_thread.adjust_num_workers(w=args.save_async_threads, own_cuda_stream=True, device_list=args.gpu) # Initialize local training threads in case of using multithreading implementation (0 is acceptable) exec_ctx.eval_thread_pool.adjust_num_workers(w=args.eval_async_threads, own_cuda_stream=True, device_list=args.gpu) for th in exec_ctx.eval_thread_pool.threads: th.testset_copy = copy.deepcopy(testset) th.testloader_copy = torch.utils.data.DataLoader(th.testset_copy, batch_size=test_batch_size, num_workers=args.num_workers_test, shuffle=False, pin_memory=False, drop_last=False ) if hasattr(th.testset_copy, "load_data"): th.testset_copy.load_data() if args.worker_listen_mode <= 0: # Path of execution for local simulation init_and_train_model(args, raw_cmdline, trainloader, testloader, exec_ctx) else: # Path of execution for assist simulation with a remote means init_and_help_with_compute(args, raw_cmdline, trainloader, testloader, exec_ctx) else: # Evaluation mode # TODO: figure out how to exploit DataParallel. Currently - parallelism across workers model, criterion, round = get_training_elements(args.model, args.dataset, args, args.resume_from, args.load_best, args.device, args.loss, args.turn_off_batch_normalization_and_dropout) metrics = evaluate_model(model, testloader, criterion, args.device, round, print_freq=10, metric_to_optim=args.metric, is_rnn=args.model in RNN_MODELS) metrics_dict = create_metrics_dict(metrics) logger.info(f'Validation metrics: {metrics_dict}') wandb_wrapper.finishProject(projectWB)
def makeBackupOfServerState(H, round)
-
Expand source code
def makeBackupOfServerState(H, round): args = H["args"] job_id = args.run_id fname = args.out if fname == "": fname = os.path.join(args.checkpoint_dir, f"{args.run_id}_{args.algorithm}_{args.global_lr}_backup.bin") else: fname = fname.replace(".bin", f"{args.run_id}_{args.algorithm}_{args.global_lr}_backup.bin") execution_context = H["execution_context"] H["execution_context"] = None exec_ctx_in_arg = ('exec_ctx' in H["args"]) exec_ctx = None if exec_ctx_in_arg: exec_ctx = H["args"]['exec_ctx'] H["execution_context"] = None with open(fname, "wb") as f: pickle.dump([(job_id, H)], f) H["execution_context"] = execution_context if exec_ctx_in_arg: H["args"]['exec_ctx'] = exec_ctx
def runSimulation(cmdline, extra_=None)
-
Expand source code
def runSimulation(cmdline, extra_=None): # ================================================================================================================== # ENTRY POINT FOR LAUNCH SIMULATION FROM GUI # ================================================================================================================== global CUDA_SUPPORT args = parse_args(cmdline) if args.hostname == "": args.hostname = socket.gethostname() Logger.setup_logging(args.loglevel, args.logfilter, logfile=args.logfile) logger = Logger.get(args.run_id) logger.info(f"CLI args original: {cmdline}") if torch.cuda.device_count(): CUDA_SUPPORT = True # torch.backends.cudnn.benchmark = True else: logger.warning('CUDA unsupported!!') CUDA_SUPPORT = False if not CUDA_SUPPORT: args.gpu = "cpu" if args.deterministic: import torch.backends.cudnn as cudnn import os import random # TODO: This settings are not thread safe in case of using cuda backend from several threads # Project use execution context random generators for random and numpy in thread safe way if CUDA_SUPPORT: cudnn.deterministic = args.deterministic cudnn.benchmark = not args.deterministic torch.cuda.manual_seed(args.manual_init_seed) torch.cuda.manual_seed_all(args.manual_init_seed) # Turn off Tensor Cores if have been requested torch.backends.cudnn.allow_tf32 = args.allow_use_nv_tensorcores torch.backends.cuda.matmul.allow_tf32 = args.allow_use_nv_tensorcores torch.manual_seed(args.manual_init_seed) random.seed(args.manual_init_seed) np.random.seed(args.manual_init_seed) os.environ['PYTHONHASHSEED'] = str(args.manual_init_seed) torch.backends.cudnn.deterministic = True main(args, cmdline, extra_)
def saveResult(H)
-
Serialize server state into filesystem
Expand source code
def saveResult(H): """Serialize server state into filesystem""" args = H["args"] job_id = args.run_id fname = args.out if fname == "": fname = os.path.join(args.checkpoint_dir, f"{args.algorithm}_simulation_history.bin") # Save formally list of tuples (job_id, server_state) with open(fname, "wb") as f: pickle.dump([(job_id, H)], f) print(f"Final server state for {job_id} is serialized into: '{fname}'") print(f"Program which uses '{args.algorithm}' algorithm for {args.rounds} rounds was finished")
def terminationWithSignal(*args)
-
Expand source code
def terminationWithSignal(*args): global g_Terminate g_Terminate = True print("Request for terminate process has been obtained. Save results and terminate.")
Classes
class ClientThread (clientSocket, listenPort)
-
A class that represents a thread of control.
This class can be safely subclassed in a limited fashion. There are two ways to specify the activity: by passing a callable object to the constructor, or by overriding the run() method in a subclass.
This constructor should always be called with keyword arguments. Arguments are:
group should be None; reserved for future extension when a ThreadGroup class is implemented.
target is the callable object to be invoked by the run() method. Defaults to None, meaning nothing is called.
name is the thread name. By default, a unique name is constructed of the form "Thread-N" where N is a small decimal number.
args is the argument tuple for the target invocation. Defaults to ().
kwargs is a dictionary of keyword arguments for the target invocation. Defaults to {}.
If a subclass overrides the constructor, it must make sure to invoke the base class constructor (Thread.init()) before doing anything else to the thread.
Expand source code
class ClientThread(threading.Thread): def __init__(self, clientSocket, listenPort): threading.Thread.__init__(self) self.clientSocket = clientSocket # OS-like socket to communicate with information self.comSocket = comm_socket.CommSocket(clientSocket) # Light construction on top of socket to communicate self.listenPort = listenPort # Listening port def run(self): cmd = self.comSocket.rawRecvString() # Obtain command t = time.localtime() cur_time = time.strftime("%H:%M:%S", t) print(f"-- {cur_time}: The received command: {cmd}") if cmd == "list_of_gpus": # List of resources in the system gpus = len(gpu_utils.get_available_gpus()) self.comSocket.rawSendString(f"{gpus}") elif cmd == "execute_work": # Execute work devConfig = self.comSocket.rawRecvString() # 1-st param: device configuration cmdLine = self.comSocket.rawRecv() # 2-nd param: command line cmdLine = pickle.loads(cmdLine) # unpack cmdline ip, port, dev_name, dev_number = devConfig.split(":") # Remove unnecessary flags for worker i = len(cmdLine) - 1 while i >= 0: if cmdLine[i] == "--metric": cmdLine[i + 1] = "none" if cmdLine[i] == "--gpu": cmdLine[i + 1] = dev_number # update device number if cmdLine[i] in ['--wandb-key', '--eval-async-threads', '--save-async-threads', '--threadpool-for-local-opt', '--external-devices', '--evaluate']: del cmdLine[i] # remove flag del cmdLine[i] # remove value for flag (originally at index i+1) i -= 1 # Specify listen mode in which worker is cmdLine.append("--worker-listen-mode") cmdLine.append(str(self.listenPort)) runSimulation(cmdLine, extra_=self.comSocket) else: print(f"The received command {cmd} is not valid")
Ancestors
- threading.Thread
Methods
def run(self)
-
Method representing the thread's activity.
You may override this method in a subclass. The standard run() method invokes the callable object passed to the object's constructor as the target argument, if any, with sequential and keyword arguments taken from the args and kwargs arguments, respectively.
Expand source code
def run(self): cmd = self.comSocket.rawRecvString() # Obtain command t = time.localtime() cur_time = time.strftime("%H:%M:%S", t) print(f"-- {cur_time}: The received command: {cmd}") if cmd == "list_of_gpus": # List of resources in the system gpus = len(gpu_utils.get_available_gpus()) self.comSocket.rawSendString(f"{gpus}") elif cmd == "execute_work": # Execute work devConfig = self.comSocket.rawRecvString() # 1-st param: device configuration cmdLine = self.comSocket.rawRecv() # 2-nd param: command line cmdLine = pickle.loads(cmdLine) # unpack cmdline ip, port, dev_name, dev_number = devConfig.split(":") # Remove unnecessary flags for worker i = len(cmdLine) - 1 while i >= 0: if cmdLine[i] == "--metric": cmdLine[i + 1] = "none" if cmdLine[i] == "--gpu": cmdLine[i + 1] = dev_number # update device number if cmdLine[i] in ['--wandb-key', '--eval-async-threads', '--save-async-threads', '--threadpool-for-local-opt', '--external-devices', '--evaluate']: del cmdLine[i] # remove flag del cmdLine[i] # remove value for flag (originally at index i+1) i -= 1 # Specify listen mode in which worker is cmdLine.append("--worker-listen-mode") cmdLine.append(str(self.listenPort)) runSimulation(cmdLine, extra_=self.comSocket) else: print(f"The received command {cmd} is not valid")