def train(self,
config_file: str,
total_timesteps: int,
tot_trace_cnt: int,
tb_log_name: str = "",
validation_flag: bool = False,
training_traces: List[Trace] = [],
validation_traces: List[Trace] = []):
assert isinstance(self.model, PPO1)
if not config_file and not training_traces:
raise ValueError("Neither configuration file nor training_traces "
"are provided. Please provide one.")
elif config_file and training_traces:
raise ValueError("Both configuration file and training_traces are "
"provided. Please choose one.")
elif config_file and not training_traces:
training_traces = generate_traces(config_file,
tot_trace_cnt,
duration=30)
# generate validation traces
if not config_file and not validation_traces:
raise ValueError("Neither configuration file nor validation_traces"
" are provided. Please provide one.")
elif config_file and validation_traces:
raise ValueError("Both configuration file and validation_traces "
"are provided. Please choose one.")
elif config_file and not validation_traces:
validation_traces = generate_traces(config_file, 20, duration=30)
# config = read_json_file(config_file)[-1]
# validation_traces = [generate_trace(config['duration'],
# config['bandwidth_lower_bound'],
# config['bandwidth_upper_bound'], config['delay'],
# config['loss'], config['queue'], config['T_s'],
# config['delay_noise']) for _ in range(20)]
# self.original_training_traces = training_traces
# training_traces = self.original_training_traces[0:36]
env = gym.make('PccNs-v0',
traces=training_traces,
train_flag=True,
delta_scale=self.delta_scale,
config_file=config_file)
env.seed(self.seed)
# env = Monitor(env, self.log_dir)
self.model.set_env(env)
# Create the callback: check every n steps and save best model
callback = SaveOnBestTrainingRewardCallback(
self,
check_freq=self.timesteps_per_actorbatch,
log_dir=self.log_dir,
steps_trained=self.steps_trained,
val_traces=validation_traces,
config_file=config_file,
validation_flag=validation_flag)
self.model.learn(total_timesteps=total_timesteps,
tb_log_name=tb_log_name,
callback=callback)
def get_trace(self):
if self.aurora.callback.n_calls % self.aurora.callback.check_freq == 0:
i = int(self.aurora.callback.n_calls //
(self.aurora.callback.check_freq * 20))
self.config_file = self.config_files[i]
# print(self.n_calls, i, self.model.env.config_file)
return generate_traces(self.config_file, 1, duration=30)[0]
def get_trace(self):
if self.traces and np.random.uniform(0, 1) < self.percent:
return np.random.choice(self.traces)
elif self.config_file:
return generate_traces(self.config_file, 1, duration=30)[0]
else:
raise ValueError
def __init__(self,
count: int,
config_file: Optional[str],
config=None,
seed: int = 42):
set_seed(seed)
self.count = count
self.traces = []
self.config_file = config_file
self.config = config
if self.config_file:
self.traces = generate_traces(self.config_file, self.count, 30)
elif self.config:
self.traces = generate_traces_from_config(self.config, self.count,
30)
def sample_trace(self, difficulty):
if difficulty == 0:
target_difficulty = difficulty
elif difficulty == 1:
prob = np.random.uniform(0, 1, 1).item()
if prob < 0.7:
target_difficulty = 0
else:
target_difficulty = 1
elif difficulty == 2:
prob = np.random.uniform(0, 1, 1).item()
if prob < 0.49:
target_difficulty = 0
elif 0.49 < prob < 0.7:
target_difficulty = 1
else:
target_difficulty = 2
elif difficulty == 3:
prob = np.random.uniform(0, 1, 1).item()
if prob < 0.343:
target_difficulty = 0
elif 0.343 < prob < 0.49:
target_difficulty = 1
elif 0.49 < prob < 0.7:
target_difficulty = 2
else:
target_difficulty = 3
else:
prob = np.random.uniform(0, 1, 1).item()
if prob < 0.2401:
target_difficulty = 0
elif 0.2401 < prob < 0.343:
target_difficulty = 1
elif 0.343 < prob < 0.49:
target_difficulty = 2
elif 0.49 < prob < 0.7:
target_difficulty = 3
else:
target_difficulty = 4
while not self.difficulty_trace_cache[target_difficulty]:
traces = generate_traces(self.config_file, 1, duration=30)
self.insert_to_difficulty_cache(traces)
trace = np.random.choice(
self.difficulty_trace_cache[target_difficulty])
self.difficulty_trace_cache[target_difficulty].remove(trace)
print('select from level', target_difficulty)
return trace
def reset(self):
self.steps_taken = 0
self.net.reset()
# old snippet start
# self.current_trace = np.random.choice(self.traces)
# old snippet end
# choose real trace with a probability. otherwise, use synthetic trace
if self.train_flag and self.config_file:
self.current_trace = generate_traces(self.config_file, 1, duration=30)[0]
if random.uniform(0, 1) < self.real_trace_prob and self.traces:
config_syn = np.array(self.current_trace.real_trace_configs(normalized=True)).reshape(1, -1)
assert self.real_trace_configs is not None
dists = np.linalg.norm(self.real_trace_configs - config_syn, axis=1)
target_idx = np.argmin(dists)
real_trace = self.traces[target_idx]
# real_trace = np.random.choice(self.traces) # randomly select a real trace
real_trace.queue_size = self.current_trace.queue_size
real_trace.loss_rate = self.current_trace.loss_rate
self.current_trace = real_trace
else:
self.current_trace = np.random.choice(self.traces)
# if self.train_flag and not self.config_file:
# bdp = np.max(self.current_trace.bandwidths) / BYTES_PER_PACKET / \
# BITS_PER_BYTE * 1e6 * np.max(self.current_trace.delays) * 2 / 1000
# self.current_trace.queue_size = max(2, int(bdp * np.random.uniform(0.2, 3.0))) # hard code this for now
# loss_rate_exponent = float(np.random.uniform(np.log10(0+1e-5), np.log10(0.5+1e-5), 1))
# if loss_rate_exponent < -4:
# loss_rate = 0
# else:
# loss_rate = 10**loss_rate_exponent
# self.current_trace.loss_rate = loss_rate
self.current_trace.reset()
self.create_new_links_and_senders()
self.net = Network(self.senders, self.links, self)
self.episodes_run += 1
# old code snippet start
# if self.train_flag and self.config_file is not None and self.episodes_run % 100 == 0:
# self.traces = generate_traces(self.config_file, 10, duration=30)
# old code snippet end
self.net.run_for_dur(self.run_dur)
self.reward_ewma *= 0.99
self.reward_ewma += 0.01 * self.reward_sum
self.reward_sum = 0.0
return self._get_all_sender_obs()
def reset(self):
self.steps_taken = 0
self.net.reset()
if self.config_file:
self.current_trace = self.sample_trace(self.curr_diff)
# print(self.current_trace.optimal_reward - self.baseline.test(self.current_trace, "")[1])
self.current_trace.reset()
self.create_new_links_and_senders()
self.net = Network(self.senders, self.links, self)
self.episodes_run += 1
if self.train_flag and self.config_file is not None:
self.traces = generate_traces(self.config_file, 10, duration=30)
self.net.run_for_dur(self.run_dur)
self.reward_ewma *= 0.99
self.reward_ewma += 0.01 * self.reward_sum
self.reward_sum = 0.0
return self._get_all_sender_obs()
def reset(self):
self.steps_taken = 0
self.net.reset()
# old snippet start
# self.current_trace = np.random.choice(self.traces)
# old snippet end
# choose real trace with a probability. otherwise, use synthetic trace
if self.train_flag and self.config_file:
self.current_trace = generate_traces(self.config_file,
1,
duration=30)[0]
if random.uniform(0, 1) < self.real_trace_prob and self.traces:
cur_config = read_json_file(self.config_file)[0]
cur_config['weight'] = 1
self.current_trace = generate_trace_from_config([cur_config])
import pdb
pdb.set_trace()
else:
self.current_trace = np.random.choice(self.traces)
# if self.train_flag and not self.config_file:
# bdp = np.max(self.current_trace.bandwidths) / BYTES_PER_PACKET / \
# BITS_PER_BYTE * 1e6 * np.max(self.current_trace.delays) * 2 / 1000
# self.current_trace.queue_size = max(2, int(bdp * np.random.uniform(0.2, 3.0))) # hard code this for now
# loss_rate_exponent = float(np.random.uniform(np.log10(0+1e-5), np.log10(0.5+1e-5), 1))
# if loss_rate_exponent < -4:
# loss_rate = 0
# else:
# loss_rate = 10**loss_rate_exponent
# self.current_trace.loss_rate = loss_rate
self.current_trace.reset()
self.create_new_links_and_senders()
self.net = Network(self.senders, self.links, self)
self.episodes_run += 1
# old code snippet start
# if self.train_flag and self.config_file is not None and self.episodes_run % 100 == 0:
# self.traces = generate_traces(self.config_file, 10, duration=30)
# old code snippet end
self.net.run_for_dur(self.run_dur)
self.reward_ewma *= 0.99
self.reward_ewma += 0.01 * self.reward_sum
self.reward_sum = 0.0
return self._get_all_sender_obs()
def main():
args = parse_args()
set_seed(args.seed)
if args.heuristic == 'cubic':
cc = Cubic()
elif args.heuristic == 'bbr_old':
cc = BBR_old()
elif args.heuristic == 'bbr':
cc = BBR()
else:
raise NotImplementedError
# if 'large' in args.config_file:
# config = read_json_file(args.config_file)
# config[0]['bandwidth_lower_bound'] = (1, 1)
# config[0]['bandwidth_upper_bound'] = (1, 100)
# traces = generate_traces_from_config(config, 50, 30)
# else:
if not args.config_file:
dataset = PantheonDataset('../../data', 'all')
traces = dataset.get_traces(0, 50)
save_dirs = [os.path.join(
args.save_dir, link_conn_type, link_name, trace_name)
for link_conn_type, (link_name, trace_name) in
zip(dataset.link_conn_types, dataset.trace_names)]
else:
traces = generate_traces(args.config_file, 50, 30)
save_dirs = [os.path.join(args.save_dir, "trace_{:02d}".format(i))
for i in range(len(traces))]
cc_save_dirs = [os.path.join(save_dir, cc.cc_name)
for save_dir in save_dirs]
cc_res = cc.test_on_traces(traces, cc_save_dirs, plot_flag=False, n_proc=16)
aurora_save_dirs = [os.path.join(save_dir, 'aurora')
for save_dir in save_dirs]
aurora_res = test_on_traces(args.model_path, traces, aurora_save_dirs,
nproc=16, seed=42, record_pkt_log=False,
plot_flag=False)
print(cc.cc_name, np.mean([res[1] for res in cc_res]))
print('aurora', np.mean([res[1] for res in aurora_res]))
for i, (trace, save_dir) in enumerate(zip(traces, save_dirs)):
trace.dump(os.path.join(save_dir, 'trace_{:02d}.json'.format(i)))
def main():
args = parse_args()
set_seed(args.seed)
if args.save_dir:
os.makedirs(args.save_dir, exist_ok=True)
if args.trace_file is not None and args.trace_file.endswith('.json'):
test_traces = [Trace.load_from_file(args.trace_file)]
elif args.trace_file is not None and args.trace_file.endswith('.log'):
test_traces = [
Trace.load_from_pantheon_file(args.trace_file, args.delay,
args.loss, args.queue)
]
elif args.config_file is not None:
test_traces = generate_traces(args.config_file,
1,
args.duration,
constant_bw=not args.time_variant_bw)
else:
test_traces = [
generate_trace((args.duration, args.duration),
(args.bandwidth, args.bandwidth),
(args.delay, args.delay), (args.loss, args.loss),
(args.queue, args.queue), (60, 60), (60, 60),
constant_bw=not args.time_variant_bw)
]
# print(test_traces[0].bandwidths)
aurora = Aurora(seed=args.seed,
timesteps_per_actorbatch=10,
log_dir=args.save_dir,
pretrained_model_path=args.model_path,
delta_scale=args.delta_scale)
results, pkt_logs = aurora.test_on_traces(test_traces, [args.save_dir])
for pkt_log in pkt_logs:
with open(os.path.join(args.save_dir, "aurora_packet_log.csv"), 'w',
1) as f:
pkt_logger = csv.writer(f, lineterminator='\n')
pkt_logger.writerows(pkt_log)
def __init__(self, traces, history_len=10,
# features="sent latency inflation,latency ratio,send ratio",
features="sent latency inflation,latency ratio,recv ratio",
train_flag=False, delta_scale=1.0, config_file=None,
record_pkt_log: bool = False, real_trace_prob: float = 0):
"""Network environment used in simulation.
congestion_control_type: aurora is pcc-rl. cubic is TCPCubic.
"""
self.real_trace_prob = real_trace_prob
self.record_pkt_log = record_pkt_log
self.config_file = config_file
self.delta_scale = delta_scale
self.traces = traces
self.train_flag = train_flag
if self.config_file:
self.current_trace = generate_traces(self.config_file, 1, 30)[0]
elif self.traces:
self.current_trace = np.random.choice(self.traces)
else:
raise ValueError
if self.train_flag and self.traces:
self.real_trace_configs = []
for trace in self.traces:
self.real_trace_configs.append(trace.real_trace_configs(True))
self.real_trace_configs = np.array(self.real_trace_configs).reshape(-1, 4)
else:
self.real_trace_configs = None
self.use_cwnd = False
self.history_len = history_len
# print("History length: %d" % history_len)
self.features = features.split(",")
# print("Features: %s" % str(self.features))
self.links = None
self.senders = None
self.create_new_links_and_senders()
self.net = Network(self.senders, self.links, self)
self.run_dur = None
self.run_period = 0.1
self.steps_taken = 0
self.debug_thpt_changes = False
self.last_thpt = None
self.last_rate = None
if self.use_cwnd:
self.action_space = spaces.Box(
np.array([-1e12, -1e12]), np.array([1e12, 1e12]), dtype=np.float32)
else:
self.action_space = spaces.Box(
np.array([-1e12]), np.array([1e12]), dtype=np.float32)
self.observation_space = None
# use_only_scale_free = True
single_obs_min_vec = sender_obs.get_min_obs_vector(self.features)
single_obs_max_vec = sender_obs.get_max_obs_vector(self.features)
self.observation_space = spaces.Box(np.tile(single_obs_min_vec, self.history_len),
np.tile(single_obs_max_vec,
self.history_len),
dtype=np.float32)
self.reward_sum = 0.0
self.reward_ewma = 0.0
self.episodes_run = -1
model_path_after = os.path.join(
ROOT, 'seed_{}'.format(seed), 'config_{:02d}'.format(config_id),
'model_step_{}.ckpt'.format(int(val_log['num_timesteps'][idx])))
if not os.path.exists(model_path_after + '.index') or not os.path.exists(
os.path.join(ROOT, 'seed_{}'.format(seed),
'config_{:02d}'.format(config_id),
'model_step_64800.ckpt.meta')):
print('skip', model_path_after)
continue
# config = read_json_file(os.path.join(ROOT, 'seed_{}'.format(seed), "bo_{}.json".format(bo)))[-1]
# config['weight'] = 1
# config = [config]
# config_file = '../../config/gap_vs_improvement_pretrained/config_{:02d}.json'.format(config_id)
config_file = '../../config/gap_vs_improvement/config_{:02d}.json'.format(
config_id)
traces = generate_traces(config_file, TRACE_CNT, 30)
save_dirs = [
os.path.join(SAVE_ROOT, 'seed_{}'.format(seed),
'config_{:02d}'.format(config_id),
'trace_{:05d}'.format(i)) for i in range(len(traces))
]
before_save_dirs = [
os.path.join(save_dir, 'before') for save_dir in save_dirs
]
after_save_dirs = [
os.path.join(save_dir, 'after_best_pkt_level_reward')
for save_dir in save_dirs
]
bbr_old_save_dirs = [
os.path.join(save_dir, 'bbr_old') for save_dir in save_dirs
]
