Merge pull request #2 from kaajalbgupta/notif_kthread_iterations

Notif kthread iterations

Author: kaajalbgupta

frontend/enso/ensogen.py

@@ -364,3 +364,305 @@ def get_pcap_mean_pkt_size(self, pcap_path: str) -> float:
 
         mean_pcap_pkt_size = float(output)
         return mean_pcap_pkt_size
+
+
+class EnsoGenSwitch(Pktgen):
+    """Python wrapper for the EnsōGen packet generator."""
+
+    def __init__(
+        self,
+        nic: EnsoNic,
+        core_id: int = 0,
+        queues: int = 4,
+        single_core: bool = False,
+        rtt: bool = False,
+        rtt_hist: bool = False,
+        rtt_hist_offset: Optional[int] = None,
+        rtt_hist_len: Optional[int] = None,
+        stats_file: Optional[str] = None,
+        hist_file: Optional[str] = None,
+        stats_delay: Optional[int] = None,
+        pcie_addr: Optional[str] = None,
+        log_file: Union[bool, TextIO] = False,
+        check_tx_rate=False,
+    ) -> None:
+        super().__init__()
+
+        self.nic = nic
+        self.queues = queues
+
+        self.core_id = core_id
+        self.single_core = single_core
+        self.rtt = rtt
+        self.rtt_hist = rtt_hist
+        self.rtt_hist_offset = rtt_hist_offset
+        self.rtt_hist_len = rtt_hist_len
+        self.stats_delay = stats_delay
+
+        if pcie_addr is not None and pcie_addr.count(":") == 1:
+            pcie_addr = f"0000:{pcie_addr}"  # Add domain.
+        self.pcie_addr = pcie_addr
+
+        self.stats_file = stats_file or "stats.csv"
+        self.hist_file = hist_file or "hist.csv"
+
+        self.log_file = log_file
+        self.check_tx_rate = check_tx_rate
+
+        self.pktgen_cmd = None
+
+        self.clean_stats()
+
+    def set_params(
+        self, pkt_size: int, nb_src: int, nb_dst: int, nb_pcaps: int
+    ) -> None:
+        self.pcap_paths = [None * 1024]
+        self.nb_pcaps = 0
+
+        nb_pkts = nb_src * nb_dst
+
+        remote_dir_path = Path(self.nic.enso_path)
+
+        pcap_gen_cmd = remote_dir_path / Path(PCAP_GEN_CMD)
+        for i in range(nb_pcaps):
+            dst_start = i * nb_dst
+            pcap_name = (
+                f"{nb_pkts}_{pkt_size}_{nb_src}_{nb_dst}_{dst_start}.pcap"
+            )
+            pcap_dst = remote_dir_path / Path(PCAPS_DIR) / Path(pcap_name)
+
+            pcap_gen_cmd = (
+                f"{pcap_gen_cmd} {nb_pkts} {pkt_size} {nb_src} "
+                f"{nb_dst} {dst_start}"
+            )
+
+            pcap_gen_cmd = self.nic.host.run_command(
+                pcap_gen_cmd, print_command=self.log_file
+            )
+            pcap_gen_cmd.watch(stdout=self.log_file, stderr=self.log_file)
+
+            status = pcap_gen_cmd.recv_exit_status()
+            if status != 0:
+                raise RuntimeError("Error generating pcap")
+
+            self.pcap_paths[i] = pcap_dst
+            self.nb_pcaps += 1
+
+    def start(self, throughput: float, nb_pkts: int, window_size: int) -> None:
+        """Start packet generation.
+
+        Args:
+            throughput: Throughput in bits per second.
+            nb_pkts: Number of packets to transmit.
+        """
+        if self.pcap_paths is None:
+            raise RuntimeError("No pcaps were configured")
+
+        bits_per_pkt = (self.mean_pcap_pkt_size + ETHERNET_OVERHEAD) * 8
+        pkts_per_sec = throughput / bits_per_pkt
+        flits_per_pkt = math.ceil(self.mean_pcap_pkt_size / 64)
+
+        hardware_rate = pkts_per_sec * flits_per_pkt / FPGA_RATELIMIT_CLOCK
+
+        rate_frac = Fraction(hardware_rate).limit_denominator(1000)
+
+        self.current_throughput = throughput
+        self.current_goodput = pkts_per_sec * self.mean_pcap_pkt_size * 8
+        self.current_pps = pkts_per_sec
+        self.expected_tx_duration = nb_pkts / pkts_per_sec
+
+        # Make sure remote stats file does not exist.
+        remove_stats_file = self.nic.host.run_command(
+            f"rm -f {self.stats_file}",
+            print_command=False,
+        )
+        remove_stats_file.watch(stdout=self.log_file, stderr=self.log_file)
+        status = remove_stats_file.recv_exit_status()
+        if status != 0:
+            raise RuntimeError(
+                f"Error removing remote stats file ({self.stats_file})"
+            )
+
+        command = (
+            f"sudo {self.nic.enso_path}/{ENSOGEN_CMD}"
+            f" {rate_frac.numerator} {rate_frac.denominator}"
+            + f" --pkts-per-pcap {window_size}"
+            f" --count {nb_pkts}"
+            f" --core {self.core_id}"
+            f" --queues {self.queues}"
+            f" --save {self.stats_file}"
+        )
+
+        for i in range(self.nb_pcaps):
+            command += f" --pcap-file {self.pcap_paths[i]}"
+
+        if self.single_core:
+            command += " --single-core"
+
+        if self.rtt:
+            command += " --rtt"
+
+        if self.rtt_hist:
+            command += f" --rtt-hist {self.hist_file}"
+
+        if self.rtt_hist_offset is not None:
+            command += f" --rtt-hist-offset {self.rtt_hist_offset}"
+
+        if self.rtt_hist_len is not None:
+            command += f" --rtt-hist-len {self.rtt_hist_len}"
+
+        if self.stats_delay is not None:
+            command += f" --stats-delay {self.stats_delay}"
+
+        if self.pcie_addr is not None:
+            command += f" --pcie-addr {self.pcie_addr}"
+
+        self.pktgen_cmd = self.nic.host.run_command(
+            command,
+            print_command=self.log_file,
+            pty=True,
+        )
+
+    def wait_transmission_done(self) -> None:
+        pktgen_cmd = self.pktgen_cmd
+
+        if pktgen_cmd is None:
+            # Pktgen is not running.
+            return
+
+        pktgen_cmd.watch(
+            stdout=self.log_file,
+            stderr=self.log_file,
+            keyboard_int=lambda: pktgen_cmd.send(b"\x03"),
+        )
+        status = pktgen_cmd.recv_exit_status()
+        if status != 0:
+            raise RuntimeError("Error running EnsōGen")
+
+        self.update_stats()
+
+        self.pktgen_cmd = None
+
+    def update_stats(self) -> None:
+        # Make sure transmission rate matches specification for sufficiently
+        # high rates (i.e., >50Gbps).
+        calculate_tx_mean = (
+            self.check_tx_rate
+            and (self.current_throughput > 50e9)
+            and (self.expected_tx_duration > 4)
+        )
+
+        # Retrieve the latest stats.
+        with tempfile.TemporaryDirectory() as tmp:
+            local_stats = f"{tmp}/stats.csv"
+            download_file(
+                self.nic.host_name, self.stats_file, local_stats, self.log_file
+            )
+            parsed_stats = EnsoGenStats(local_stats)
+
+            stats_summary = parsed_stats.get_summary(calculate_tx_mean)
+
+        # Check TX rate.
+        if calculate_tx_mean:
+            tx_goodput_mbps = stats_summary["tx_mean_goodput_mbps"]
+            tx_goodput_gbps = tx_goodput_mbps / 1e3
+            expected_goodput_gbps = self.current_goodput / 1e9
+
+            if abs(tx_goodput_gbps - expected_goodput_gbps) > 1.0:
+                raise RuntimeError(
+                    f"TX goodput ({tx_goodput_gbps} Gbps) does not match "
+                    f"specification ({expected_goodput_gbps} Gbps)"
+                )
+
+        self.mean_rx_goodput = stats_summary.get("rx_mean_goodput_mbps", 0)
+        self.mean_rx_goodput *= 1_000_000
+        self.mean_tx_goodput = stats_summary.get("tx_mean_goodput_mbps", 0)
+        self.mean_tx_goodput *= 1_000_000
+
+        self.mean_rx_rate = stats_summary.get("rx_mean_rate_kpps", 0) * 1000
+        self.mean_tx_rate = stats_summary.get("tx_mean_rate_kpps", 0) * 1000
+
+        self.nb_rx_pkts += stats_summary.get("rx_packets", 0)
+        self.nb_tx_pkts += stats_summary.get("tx_packets", 0)
+
+        self.nb_rx_bytes += stats_summary.get("rx_bytes", 0)
+        self.nb_tx_bytes += stats_summary.get("tx_bytes", 0)
+
+        self.mean_rtt = stats_summary.get("mean_rtt_ns", 0)
+
+    @property
+    def pcap_path(self) -> None:
+        return self._pcap_path
+
+    @pcap_path.setter
+    def pcap_path(self, pcap_path) -> None:
+        self.mean_pcap_pkt_size = self.get_pcap_mean_pkt_size(pcap_path)
+        self._pcap_path = pcap_path
+
+    @property
+    def queues(self) -> int:
+        return self._queues
+
+    @queues.setter
+    def queues(self, queues) -> None:
+        self._queues = queues
+
+    def get_nb_rx_pkts(self) -> int:
+        return self.nb_rx_pkts
+
+    def get_nb_tx_pkts(self) -> int:
+        return self.nb_tx_pkts
+
+    def get_rx_throughput(self) -> int:
+        return self.mean_rx_goodput + self.mean_rx_rate * 24 * 8
+
+    def get_tx_throughput(self) -> int:
+        return self.mean_tx_goodput + self.mean_tx_rate * 24 * 8
+
+    def clean_stats(self) -> None:
+        self.nb_rx_pkts = 0
+        self.nb_rx_bytes = 0
+        self.mean_rx_goodput = 0
+        self.mean_rx_rate = 0
+        self.nb_tx_pkts = 0
+        self.mean_tx_goodput = 0
+        self.mean_tx_rate = 0
+        self.nb_tx_bytes = 0
+        self.mean_rtt = 0
+
+    def stop(self) -> None:
+        if self.pktgen_cmd is None:
+            # Pktgen is not running.
+            return
+
+        if self.pktgen_cmd.exit_status_ready():
+            # Pktgen has already exited.
+            self.pktgen_cmd = None
+            return
+
+        self.pktgen_cmd.send(b"\x03")
+
+        self.pktgen_cmd.watch(stdout=self.log_file, stderr=self.log_file)
+        status = self.pktgen_cmd.recv_exit_status()
+        if status != 0:
+            raise RuntimeError("Error stopping EnsōGen")
+
+        self.update_stats()
+
+        self.pktgen_cmd = None
+
+    def close(self) -> None:
+        # No need to close here.
+        pass
+
+    def get_pcap_mean_pkt_size(self, pcap_path: str) -> float:
+        cmd_str = f"{self.nic.enso_path}/{GET_PCAP_SIZE_CMD} {pcap_path}"
+        cmd = self.nic.host.run_command(cmd_str, pty=True)
+        output = cmd.watch(keyboard_int=lambda: cmd.send("\x03"))
+        status = cmd.recv_exit_status()
+
+        if status != 0:
+            raise RuntimeError(f'Error processing pcap: "{output}"')
+
+        mean_pcap_pkt_size = float(output)
+        return mean_pcap_pkt_size