clean codes

gen_ckpt.py

@@ -80,7 +80,6 @@ def __init__(self):
         model.eval()
         self.model=model
         self.lr=10
-        # self.bucket_thold = dataset_config["bucket_thold"]
         self.enable_context = dataset_config.get("enable_context", False)
         self.enable_log=training_config["enable_log"]
         bdp_dict_db = {}
@@ -102,7 +101,7 @@ def run_inference(self,idx):
         size_bucket_list = self.bdp_dict_db[n_hosts]
         size_bucket_list_output = self.bdp_dict_db_output[n_hosts]
         print(f"spec-1: {spec}, {topo_type}")   
-        spec=f'data_lr10Gbps_{n_hosts}'
+        spec=f'data_lr10Gbps'
         print(f"spec-2: {spec}")
 
         dir_input_tmp = f"{self.dir_input}/{spec}"

util/consts.py

@@ -18,11 +18,7 @@
     5: 10 * MTU,
     7: 15 * MTU,
 }
-# LINK_TO_DELAY_DICT={
-#     3:np.array([0,2*DELAY_PROPAGATION_BASE,0]),
-#     5:np.array([0,2*DELAY_PROPAGATION_BASE,2*DELAY_PROPAGATION_BASE,2*DELAY_PROPAGATION_BASE,0]),
-#     7:np.array([0,2*DELAY_PROPAGATION_BASE,2*DELAY_PROPAGATION_BASE,2*DELAY_PROPAGATION_BASE,2*DELAY_PROPAGATION_BASE,2*DELAY_PROPAGATION_BASE,0]),
-# }
+
 LINK_TO_DELAY_DICT={
     3:np.array([0,0,0]),
     5:np.array([0,0,1*DELAY_PROPAGATION_BASE,0,0]),
@@ -35,17 +31,6 @@
 MTU_IN_BYTE = MTU + np.ceil(MTU / MTU) * HEADER_SIZE
 MTU_IN_BIT = (MTU + np.ceil(MTU / MTU) * HEADER_SIZE) * BYTE_TO_BIT
 EPS = 1e-12
-# SIZE_BUCKET_LIST = (80 * (1.5 ** np.arange(19))).astype(np.int32)
-# SIZE_BUCKET_LIST = np.array([MTU, BDP, 2 * BDP, 5 * BDP, 10 * BDP, 20 * BDP, 50 * BDP])
-# P99_DIVIDENT=[np.arange(0,1),np.arange(1,2),np.arange(2,4),np.arange(4,8)]
-
-# P99_DIVIDENT=[np.arange(0,1),np.arange(1,2),np.arange(2,3),np.arange(3,4)]
-# SIZE_BUCKET_LIST = np.array([MTU,BDP,5 * BDP])
-
-# SIZE_BUCKET_LIST = np.array(
-#     [MTU // 4, MTU // 2, MTU, 2 * MTU, 5 * MTU, BDP, 2 * BDP, 3 * BDP, 5 * BDP]
-# )
-# SIZE_BUCKET_LIST_LABEL = ["(0, 0.25MTU)","(0.25MTU, 0.5MTU)","(0.5MTU, MTU)","(MTU, 2MTU)","(2MTU, 5MTU)","(5MTU, BDP)","(BDP, 2BDP)","(2BDP, 3BDP)","(3BDP, 5BDP)","(5BDP, INF)"]
 SIZE_BUCKET_LIST = np.array(
     [
         MTU // 4,
@@ -60,7 +45,6 @@
     ]
 )
 P99_DIVIDENT = [np.arange(0, 4), np.arange(4, 8), np.arange(8, 9), np.arange(9, 10)]
-# P99_DIVIDENT = [np.arange(0, 2), np.arange(2, 4), np.arange(4, 6), np.arange(6, 8)]
 
 def get_size_bucket_list(mtu, bdp):
     return np.array(
@@ -110,29 +94,8 @@ def get_base_delay_pmn(sizes, n_links_passed, lr_bottleneck,flow_idx_target,flow
 
 SIZE_BUCKET_LIST_LABEL_OUTPUT = ["(0, MTU)", "(MTU, BDP)", "(BDP, 5BDP)", "(5BDP, INF)"]
 
-# P99_PERCENTILE_LIST = np.array(
-#     [30, 50, 60, 70, 75, 80, 82, 84, 86,88, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99]
-# )
-# P99_PERCENTILE_LIST = np.array(
-#     [10, 20, 30, 40, 50, 60, 70, 75, 80, 85, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99]
-# )
-# P99_PERCENTILE_LIST = np.array(
-#     [1, 25, 40, 55, 70, 75, 80, 85, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 100]
-# )
-
 P99_PERCENTILE_LIST = np.arange(1, 101, 1)
-# P99_PERCENTILE_LIST = np.arange(1, 102, 1)
-# P99_PERCENTILE_LIST[-1]=100
-
-# P99_PERCENTILE_LIST = np.array(
-#     [0, 10, 20, 30, 40, 50, 60, 70, 80, 85, 90, 91, 92, 93, 94, 95, 96, 97, 98, 98.2, 98.4, 98.6, 98.8, 99, 99.2,99.4, 99.6, 99.8, 100, 100]
-# )
-
-# P99_PERCENTILE_LIST = np.array(
-#     [1, 10, 20, 30, 40, 50, 55, 60, 65, 70, 
-#      75, 80, 85, 90, 92, 94, 96, 98, 98.2, 98.4, 
-#      98.6, 98.8, 99, 99.2, 99.4, 99.6, 99.8, 100, 100]
-# )
+
 PERCENTILE_METHOD='nearest'
 
 color_list = [
@@ -143,33 +106,11 @@ def get_base_delay_pmn(sizes, n_links_passed, lr_bottleneck,flow_idx_target,flow
     "seagreen",
     "black",
 ]
-# color_list = [
-#     "C1",
-#     "C2",
-#     "C3",
-#     "C4",
-#     "C5",
-#     "C6",
-#     "C7",
-# ]
+
 hatch_list = ["o", "x", "/", ".", "*", "-", "\\"]
 linestyle_list = ["solid", "dashed", "dashdot", "dotted"]
 markertype_list = ["o", "^", "x", "x", "|"]
 
-# FLOWSIZE_BUCKET_DICT = {
-#     0: "0<size<=MTU",
-#     1: "MTU<size<=BDP",
-#     2: "BDP<size<=10BDP",
-#     3: "10BDP<size",
-# }
-# FLOWSIZE_BUCKET_DICT = {
-#     0: "(0,MTU]",
-#     1: "(MTU,BDP]",
-#     2: "(BDP,5BDP]",
-#     3: "(5BDP,inf)",
-# }
-# SIZE_BIN_LIST = [MTU, BDP, 5 * BDP]
-
 SIZEDIST_LIST_EMPIRICAL = [
     "GoogleRPC2008",
     "AliStorage2019",

util/dataset.py

@@ -30,7 +30,6 @@
 import os
 
 def my_collate(batch):
-    # batch contains a list of tuples of structure (sequence, target)
     sizebucket_to_sldn_flowsim = [item[0] for item in batch]
     sizebucket_to_sldn_flowsim = np.concatenate(sizebucket_to_sldn_flowsim, 0)
 
@@ -42,7 +41,6 @@ def my_collate(batch):
     spec = np.array([item[4] for item in batch])
     sizebucket_to_sldn_flowsim_idx = np.array([item[5] for item in batch])
     src_dst_pair_target_str = np.array([item[6] for item in batch])
-    # src_dst_pair_target = np.array([item[7] for item in batch])
     return (
         torch.tensor(sizebucket_to_sldn_flowsim),
         torch.tensor(num_flows_per_cell_flowsim),
@@ -51,7 +49,6 @@ def my_collate(batch):
         spec,
         sizebucket_to_sldn_flowsim_idx,
         src_dst_pair_target_str,
-        # torch.tensor(src_dst_pair_target),
     )
 
 class PathDataModule(LightningDataModule):
@@ -143,80 +140,45 @@ def setup(self, stage: str):
             self.__dump_data_list(self.dir_output)
 
         if stage == "test":
-            if self.test_on_manual:
+            if self.test_on_empirical:
                 data_list_test = []
-                for shard in np.arange(0, 3000):
+                for shard in np.arange(10000, 10200):
                     for n_flows in [30000]:
                         for n_hosts in [2, 3, 4, 5, 6, 7, 8]:
                             topo_type_cur = self.topo_type.replace(
                                 "x-x", f"{n_hosts}-{n_hosts}"
                             )
                             spec = f"shard{shard}_nflows{n_flows}_nhosts{n_hosts}_lr{self.lr}Gbps"
                             dir_input_tmp = f"{self.dir_input}/{spec}"
-                            if not os.path.exists(f"{dir_input_tmp}/flow_src_dst.npy"):
+                            if not os.path.exists(
+                                f"{dir_input_tmp}/flow_src_dst.npy"
+                            ):
                                 continue
-                            flow_src_dst = np.load(f"{dir_input_tmp}/flow_src_dst.npy")
+                            flow_src_dst = np.load(
+                                f"{dir_input_tmp}/flow_src_dst.npy"
+                            )
                             stats = decode_dict(
                                 np.load(
                                     f"{dir_input_tmp}/stats.npy",
                                     allow_pickle=True,
                                     encoding="bytes",
                                 ).item()
                             )
-
                             n_flows_total = stats["n_flows"]
-                            if len(flow_src_dst) == n_flows_total:
-                                target_idx = stats["host_pair_list"].index(
-                                    (0, n_hosts - 1)
-                                )
-                                size_dist = stats["size_dist_candidates"][
-                                    target_idx
-                                ].decode("utf-8")
-                                if size_dist != "gaussian":
-                                    continue
+                            if (
+                                n_flows_total < 5000000
+                                and len(flow_src_dst) == n_flows_total
+                            ):
                                 data_list_test.append(
                                     (spec, (0, n_hosts - 1), topo_type_cur)
                                 )
+
             else:
-                if self.test_on_empirical:
-                    data_list_test = []
-                    for shard in np.arange(10000, 10200):
-                        for n_flows in [30000]:
-                            for n_hosts in [2, 3, 4, 5, 6, 7, 8]:
-                                topo_type_cur = self.topo_type.replace(
-                                    "x-x", f"{n_hosts}-{n_hosts}"
-                                )
-                                spec = f"shard{shard}_nflows{n_flows}_nhosts{n_hosts}_lr{self.lr}Gbps"
-                                dir_input_tmp = f"{self.dir_input}/{spec}"
-                                if not os.path.exists(
-                                    f"{dir_input_tmp}/flow_src_dst.npy"
-                                ):
-                                    continue
-                                flow_src_dst = np.load(
-                                    f"{dir_input_tmp}/flow_src_dst.npy"
-                                )
-                                stats = decode_dict(
-                                    np.load(
-                                        f"{dir_input_tmp}/stats.npy",
-                                        allow_pickle=True,
-                                        encoding="bytes",
-                                    ).item()
-                                )
-                                n_flows_total = stats["n_flows"]
-                                if (
-                                    n_flows_total < 5000000
-                                    and len(flow_src_dst) == n_flows_total
-                                ):
-                                    data_list_test.append(
-                                        (spec, (0, n_hosts - 1), topo_type_cur)
-                                    )
-
+                data_list = self.__read_data_list(self.dir_output)
+                if self.test_on_train:
+                    data_list_test = data_list["train"]
                 else:
-                    data_list = self.__read_data_list(self.dir_output)
-                    if self.test_on_train:
-                        data_list_test = data_list["train"]
-                    else:
-                        data_list_test = data_list["test"]
+                    data_list_test = data_list["test"]
             self.test = self.__create_dataset(
                 data_list_test,
                 self.dir_input,
@@ -520,5 +482,4 @@ def __getitem__(self, idx):
             spec,
             n_input,
             src_dst_pair_target_str,
-            # np.array(src_dst_pair_target),
         )

util/model.py

@@ -1,5 +1,4 @@
 import torch
-import torch.distributed as dist
 
 from pytorch_lightning import LightningModule
 import torch.nn as nn
@@ -273,7 +272,6 @@ def step(self, batch, batch_idx, tag=None):
             spec,
             sizebucket_to_sldn_flowsim_idx,
             src_dst_pair_target_str,
-            # src_dst_pair_target,
         ) = batch
 
         if self.enable_const_opt:
@@ -300,7 +298,6 @@ def step(self, batch, batch_idx, tag=None):
                     idx_start + 1 : idx_start + idx_interval
                 ]
 
-                # tmp=torch.flatten(tmp).long()
                 sizebucket_to_sldn_background, _ = self.model_transformer(
                     tmp[None, :]
                 )
@@ -315,7 +312,6 @@ def step(self, batch, batch_idx, tag=None):
         else:
             sizebucket_to_sldn_foreground = sizebucket_to_sldn_flowsim[:, 0, :]
             sizebucket_to_sldn_input = sizebucket_to_sldn_foreground
-        # Instead of sizebucket_to_sldn_est = self.model_mlp(sizebucket_to_sldn_input) + 1.0
         sizebucket_to_sldn_est = self.model_mlp(sizebucket_to_sldn_input)
         sizebucket_to_sldn_est.add_(1.0)  # In-place addition
 
@@ -348,53 +344,32 @@ def step(self, batch, batch_idx, tag=None):
                 prog_bar=True,
                 batch_size=self.batch_size,
             )
-        # logging.info(f"step-{batch_idx}-{tag}_loss: {loss}")
 
         if tag == "test":
             test_dir = f"{self.save_dir}/{spec[0]}_{src_dst_pair_target_str[0]}"
-            # logging.info(f"save to {test_dir}")
             os.makedirs(test_dir, exist_ok=True)
             sizebucket_to_sldn_flowsim = sizebucket_to_sldn_flowsim.cpu().numpy()[0]
             sizebucket_to_sldn_est = sizebucket_to_sldn_est.cpu().numpy()[0]
             sizebucket_to_sldn = sizebucket_to_sldn.cpu().numpy()[0]
             num_flows_per_cell = num_flows_per_cell.cpu().numpy()[0]
-            # num_flows_per_cell_flowsim_ori=num_flows_per_cell_flowsim_ori.cpu().numpy()[0]
-            # num_flows_per_cell_flowsim=num_flows_per_cell_flowsim.cpu().numpy()[0]
-            # error = np.divide(
-            #     abs(sizebucket_to_sldn_est - sizebucket_to_sldn),
-            #     sizebucket_to_sldn,
-            #     out=np.zeros_like(sizebucket_to_sldn),
-            #     where=sizebucket_to_sldn != 0,
-            # )
-            # logging.info(
-            #     np.round(np.nanmin(error), 3),
-            #     np.round(np.nanpercentile(error, 50), 3),
-            #     np.round(np.nanmax(error), 3),
-            # )
             np.savez(
                 f"{test_dir}/res.npz",
                 sizebucket_to_sldn_est=sizebucket_to_sldn_est,
                 sizebucket_to_sldn_flowsim=sizebucket_to_sldn_flowsim,
                 sizebucket_to_sldn=sizebucket_to_sldn,
                 num_flows_per_cell=num_flows_per_cell,
-                # num_flows_per_cell_flowsim_ori=num_flows_per_cell_flowsim_ori,
-                # num_flows_per_cell_flowsim=num_flows_per_cell_flowsim,
             )
         return loss
 
     def configure_optimizers(self):
         optimizer = self.model_transformer.configure_optimizers(
             self.weight_decay, self.learning_rate, self.betas
         )
-        # optimizer_mlp = torch.optim.Adam(
-        #     self.model.parameters(), lr=self.learning_rate
-        # )
         optimizer.add_param_group(
             {"params": self.model_mlp.parameters(), "weight_decay": 0.0}
         )
         if self.enable_const_opt:
             optimizer.add_param_group(
                 {"params": self.const_tensor, "weight_decay": 0.0}
             )
-        # return optimizer_transformer,optimizer_mlp
         return optimizer

util/model_llama.py

@@ -1,10 +1,8 @@
 import math
-import struct
 import inspect
 from dataclasses import dataclass
 from typing import Any, Optional, Tuple
 
-import numpy as np
 import torch
 import torch.nn.functional as F
 from torch import nn