diff --git a/README.md b/README.md index b9d2ded3..843efca5 100644 --- a/README.md +++ b/README.md @@ -62,8 +62,8 @@ source $HOME/.local/bin/env 3. Set up the environment: ```bash -sudo apt install iperf -y -uv venv +sudo apt install iperf -y # or conda install -c conda-forge iperf +uv venv --python 3.12.1 source .venv/bin/activate uv sync --extra all git submodule update --init --recursive diff --git a/src/zeroband/comms.py b/src/zeroband/comms.py index ca3d7ce6..218694e3 100644 --- a/src/zeroband/comms.py +++ b/src/zeroband/comms.py @@ -29,6 +29,24 @@ BENCH_TENSOR_SIZE = 1_000_000 +def _detect_iperf_version() -> str: + """Detect which version of iperf is available on the system. + + Returns: + str: Either 'iperf3' or 'iperf' depending on which is available. + Prefers iperf3 if both are available. + """ + try: + subprocess.run(["iperf3", "--version"], capture_output=True, check=True) + return "iperf3" + except (subprocess.SubprocessError, FileNotFoundError): + try: + subprocess.run(["iperf", "--version"], capture_output=True, check=True) + return "iperf" + except (subprocess.SubprocessError, FileNotFoundError): + raise RuntimeError("Neither iperf nor iperf3 is available on the system") + + class ElasticDeviceMesh: """A class to manage the process groups for elastic training without restarts. @@ -500,10 +518,11 @@ def _start_iperf_server(self) -> None: iperf_addr = get_ip_address(IPERF_IFNAME) iperf_port = IPERF_PORT + self.world_info.global_rank - cmd: List[str] = ["iperf", "-s", "-p", str(iperf_port)] + iperf_cmd = _detect_iperf_version() + cmd: List[str] = [iperf_cmd, "-s", "-p", str(iperf_port)] self.server_process = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) self.god_store.set(f"iperf_{self.world_info.global_unique_id}", f"{iperf_addr}:{iperf_port}") - self._logger.info(f"Started iperf server on {iperf_addr} with port {iperf_port}") + self._logger.info(f"Started {iperf_cmd} server on {iperf_addr} with port {iperf_port}") except Exception as e: self._logger.error(f"Failed to start iperf server: {str(e)}") raise @@ -529,8 +548,9 @@ def measure_bandwidth(self, target_host: str, target_port: int) -> int: int: The time taken to transfer 10Tb of data in seconds """ try: + iperf_cmd = _detect_iperf_version() cmd: List[str] = [ - "iperf", + iperf_cmd, "-c", target_host, "-p",