apollo_network: broadcast network stress test run draft

Author: sirandreww-starkware

crates/apollo_network_benchmark/README.md

@@ -0,0 +1,11 @@
+# does not build, only copies the binary and entrypoint.sh from the local machine
+FROM ubuntu:24.04
+
+RUN apt update && apt -y install iproute2 kmod chrony && apt clean && rm -rf /var/lib/apt/lists/*
+
+COPY --from=tmp broadcast_network_stress_test_node /usr/local/bin/broadcast_network_stress_test_node
+COPY ./crates/apollo_network_benchmark/run/entrypoint.sh /entrypoint.sh
+
+RUN chmod +x /entrypoint.sh
+
+ENTRYPOINT ["/entrypoint.sh"]

crates/apollo_network_benchmark/run/Dockerfile.fast

@@ -0,0 +1,37 @@
+# syntax = devthefuture/dockerfile-x
+#
+# Run with the project root context:
+# docker build -f crates/apollo_network/src/bin/broadcast_network_stress_test_node/cluster/Dockerfile .
+#
+# For time synchronization, you can either:
+# 1. Use the built-in chrony service (configured below), or
+# 2. Mount host time at runtime with: -v /etc/localtime:/etc/localtime:ro -v /etc/timezone:/etc/timezone:ro
+
+INCLUDE deployments/images/base/Dockerfile
+
+# --- Stage 1: Install cargo-chef and prepare recipe ---
+FROM base AS planner
+WORKDIR /app
+COPY . .
+RUN cargo chef prepare --recipe-path recipe.json
+
+# --- Stage 2: Build dependencies ---
+FROM base AS builder
+WORKDIR /app
+COPY --from=planner /app/recipe.json recipe.json
+RUN cargo chef cook --recipe-path recipe.json
+COPY . .
+# using tokio_unstable to get additional tokio metrics
+RUN RUSTFLAGS="--cfg tokio_unstable" cargo build --release --bin broadcast_network_stress_test_node
+
+# --- Final Stage: Runtime image ---
+FROM ubuntu:24.04 AS final_stage
+
+RUN apt update && apt -y install iproute2 kmod chrony && apt clean && rm -rf /var/lib/apt/lists/*
+
+COPY --from=builder /app/target/release/broadcast_network_stress_test_node /usr/local/bin/broadcast_network_stress_test_node
+COPY --from=builder /app/crates/apollo_network_benchmark/run/entrypoint.sh /entrypoint.sh
+
+RUN chmod +x /entrypoint.sh
+
+ENTRYPOINT ["/entrypoint.sh"]

crates/apollo_network_benchmark/run/Dockerfile.slow

@@ -0,0 +1,97 @@
+#!/bin/bash
+
+set -e
+
+echo "Starting container with hostname: $(hostname)"
+
+# For Indexed Jobs: Set stable hostname based on JOB_COMPLETION_INDEX
+# This enables stable DNS names like: broadcast-network-stress-test-0.broadcast-network-stress-test-headless
+if [ ! -z "$JOB_COMPLETION_INDEX" ]; then
+    NEW_HOSTNAME="broadcast-network-stress-test-${JOB_COMPLETION_INDEX}"
+    hostname "$NEW_HOSTNAME" || echo "Warning: Could not set hostname to $NEW_HOSTNAME"
+    echo "Set hostname to: $(hostname) (based on JOB_COMPLETION_INDEX=$JOB_COMPLETION_INDEX)"
+fi
+
+# ********************************* machine information *********************************
+
+echo "Machine identification:"
+echo "  Container ID: $(cat /proc/self/cgroup 2>/dev/null | head -1 | cut -d/ -f3 | cut -c1-12 || echo 'N/A')"
+echo "  Host IP addresses:"
+ip addr show | grep -E 'inet [0-9]' | awk '{print "    " $2}' || echo "    IP info unavailable"
+echo "  Machine ID: $(cat /etc/machine-id 2>/dev/null || echo 'N/A')"
+echo "  Kernel: $(uname -r)"
+echo "  Architecture: $(uname -m)"
+if [ -n "$NODE_NAME" ]; then
+    echo "  Kubernetes Node: $NODE_NAME"
+fi
+if [ -n "$KUBERNETES_NODE_NAME" ]; then
+    echo "  K8s Node Name: $KUBERNETES_NODE_NAME"
+fi
+
+# ***************************** throttling connection start *****************************
+
+set -e
+
+INTERFACE="eth0"  # Default Docker interface
+
+# Load ifb module for ingress shaping
+modprobe ifb || echo "ifb module already loaded or not needed"
+
+# Set up ifb0 for ingress
+ip link add ifb0 type ifb || true
+ip link set ifb0 up || true
+
+# Redirect all ingress traffic to ifb0
+tc qdisc add dev $INTERFACE ingress handle ffff: || true
+tc filter add dev $INTERFACE parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev ifb0 || true
+
+# Function to apply shaping (htb for bandwidth + netem for latency)
+apply_shaping() {
+  local dev=$1
+  local parent=$2
+  local handle=$3  # e.g., 1 (no trailing :)
+
+  # If throughput is set, calculate rate in kbit/s (assuming THROUGHPUT in KB/s)
+  if [ ! -z "${THROUGHPUT}" ]; then
+    RATE=$((THROUGHPUT * 8))
+    tc qdisc add dev $dev $parent handle ${handle}: htb default 1 || true
+    tc class add dev $dev parent ${handle}: classid ${handle}:1 htb rate ${RATE}kbit ceil ${RATE}kbit || true
+    netem_parent="${handle}:1"
+  else
+    netem_parent="root"
+  fi
+
+  # If latency is set, add netem (delay in ms)
+  if [ ! -z "${LATENCY}" ]; then
+    # Use maximum queue limit to prevent packet drops
+    LIMIT=1000000 # 1.4 GiB
+    tc qdisc add dev $dev parent $netem_parent netem delay ${LATENCY}ms limit ${LIMIT} || true
+    echo "Applied netem delay ${LATENCY}ms with queue limit ${LIMIT} packets"
+  fi
+}
+
+# Apply to egress (eth0)
+# apply_shaping $INTERFACE "root" "1"
+
+# Apply to ingress (ifb0)
+apply_shaping ifb0 "root" "1"
+
+# ***************************** throttling connection end *****************************
+
+# Call broadcast_network_stress_test_node
+# Use ID from environment variable if set, otherwise use JOB_COMPLETION_INDEX (for Indexed Jobs)
+if [ -z "$ID" ]; then
+    if [ ! -z "$JOB_COMPLETION_INDEX" ]; then
+        # Used when running in Indexed Job
+        export ID=$JOB_COMPLETION_INDEX
+        echo "ID not set in environment, using JOB_COMPLETION_INDEX: $ID"
+    else
+        # Fallback: extract from hostname (legacy StatefulSet support)
+        export ID=$(hostname | grep -o '[0-9]*$')
+        echo "ID not set in environment, extracted from hostname: $ID"
+    fi
+else
+    # Used when running locally
+    echo "Using ID from environment variable: $ID"
+fi
+exec broadcast_network_stress_test_node

crates/apollo_network_benchmark/run/entrypoint.sh

@@ -0,0 +1,111 @@
+use std::collections::HashMap;
+
+use apollo_network_benchmark::node_args::UserArgs;
+use clap::Parser;
+use serde::Serialize;
+
+/// Configuration for the orchestrator that spawns multiple nodes
+#[derive(Parser, Debug, Clone, Serialize)]
+pub struct SharedArgs {
+    /// Number of nodes to run
+    #[arg(long, default_value = "3")]
+    pub num_nodes: u32,
+
+    /// Sets the multi-addresses to use UDP/QUIC instead of TCP
+    #[arg(long, default_value = "false")]
+    pub quic: bool,
+
+    #[command(flatten)]
+    pub user: UserArgs,
+}
+
+pub fn get_arguments(
+    id: Option<u32>,
+    metric_port: u16,
+    p2p_port: u16,
+    bootstrap_nodes: &[String],
+    args: &SharedArgs,
+) -> Vec<(String, String)> {
+    let broadcaster =
+        args.user.broadcaster.and_then(|b| u32::try_from(b).ok()).unwrap_or(args.num_nodes - 1);
+
+    let mut result = vec![];
+
+    if let Some(id) = id {
+        result.push(("--id".to_string(), id.to_string()));
+    }
+
+    result.extend_from_slice(&[
+        ("--metric-port".to_string(), metric_port.to_string()),
+        ("--p2p-port".to_string(), p2p_port.to_string()),
+        ("--bootstrap".to_string(), bootstrap_nodes.join(",")),
+        ("--timeout".to_string(), args.user.timeout.to_string()),
+        ("--verbosity".to_string(), args.user.verbosity.to_string()),
+        ("--buffer-size".to_string(), args.user.buffer_size.to_string()),
+        ("--message-size-bytes".to_string(), args.user.message_size_bytes.to_string()),
+        ("--heartbeat-millis".to_string(), args.user.heartbeat_millis.to_string()),
+        ("--mode".to_string(), args.user.mode.to_string()),
+        ("--network-protocol".to_string(), args.user.network_protocol.to_string()),
+        ("--broadcaster".to_string(), broadcaster.to_string()),
+        ("--round-duration-seconds".to_string(), args.user.round_duration_seconds.to_string()),
+        (
+            "--explore-cool-down-duration-seconds".to_string(),
+            args.user.explore_cool_down_duration_seconds.to_string(),
+        ),
+        (
+            "--explore-run-duration-seconds".to_string(),
+            args.user.explore_run_duration_seconds.to_string(),
+        ),
+        (
+            "--explore-min-throughput-byte-per-seconds".to_string(),
+            args.user.explore_min_throughput_byte_per_seconds.to_string(),
+        ),
+        (
+            "--explore-min-message-size-bytes".to_string(),
+            args.user.explore_min_message_size_bytes.to_string(),
+        ),
+        ("--num-nodes".to_string(), args.num_nodes.to_string()),
+    ]);
+
+    result
+}
+
+pub fn get_env_vars(
+    id: Option<u32>,
+    metric_port: u16,
+    p2p_port: u16,
+    bootstrap_nodes: &[String],
+    args: &SharedArgs,
+    latency: Option<u32>,
+    throughput: Option<u32>,
+) -> Vec<HashMap<String, String>> {
+    let arguments = get_arguments(id, metric_port, p2p_port, bootstrap_nodes, args);
+
+    let mut env_vars = vec![];
+
+    // Convert arguments to environment variables
+    for (name, value) in arguments {
+        let env_name = name[2..].replace("-", "_").to_uppercase();
+        let mut env_map = HashMap::new();
+        env_map.insert("name".to_string(), env_name);
+        env_map.insert("value".to_string(), value);
+        env_vars.push(env_map);
+    }
+
+    // Add latency and throughput if provided
+    if let Some(latency) = latency {
+        let mut env_map = HashMap::new();
+        env_map.insert("name".to_string(), "LATENCY".to_string());
+        env_map.insert("value".to_string(), latency.to_string());
+        env_vars.push(env_map);
+    }
+
+    if let Some(throughput) = throughput {
+        let mut env_map = HashMap::new();
+        env_map.insert("name".to_string(), "THROUGHPUT".to_string());
+        env_map.insert("value".to_string(), throughput.to_string());
+        env_vars.push(env_map);
+    }
+
+    env_vars
+}

crates/apollo_network_benchmark/src/bin/apollo_network_benchmark_run/args.rs

@@ -0,0 +1,45 @@
+use crate::mod_utils::{cluster_deployment_file_path, read_deployment_file, run_cmd};
+
+pub fn run() -> Result<(), String> {
+    let deployment_data = read_deployment_file(&cluster_deployment_file_path())?;
+
+    let num_nodes = deployment_data
+        .get("args")
+        .and_then(|a| a.get("shared"))
+        .and_then(|s| s.get("num_nodes"))
+        .and_then(|n| n.as_u64())
+        .ok_or("Failed to get num_nodes from deployment file")?;
+
+    let namespace_name = deployment_data
+        .get("namespace")
+        .and_then(|n| n.as_str())
+        .ok_or("Failed to get namespace from deployment file")?;
+
+    run_cmd(
+        &format!("kubectl get pods -n {}", namespace_name),
+        "Check if pods are running",
+        false,
+    )?;
+
+    for i in 0..num_nodes {
+        run_cmd(
+            &format!(
+                "timeout 5 kubectl logs -n {} -l \
+                 app=broadcast-network-stress-test,batch.kubernetes.io/job-completion-index={} > \
+                 /tmp/broadcast-network-stress-test-{}.logs.txt",
+                namespace_name, i, i
+            ),
+            &format!("Check logs for node {}", i),
+            true,
+        )?;
+    }
+
+    run_cmd(
+        &format!("kubectl get pods -n {}", namespace_name),
+        "Check if pods are running",
+        false,
+    )?;
+
+    println!("Cluster logs have been saved to /tmp/broadcast-network-stress-test-*.logs.txt");
+    Ok(())
+}

crates/apollo_network_benchmark/src/bin/apollo_network_benchmark_run/cluster_logs.rs

@@ -0,0 +1,54 @@
+use std::process::{Command, Stdio};
+use std::thread;
+
+use crate::mod_utils::{cluster_deployment_file_path, connect_to_cluster, read_deployment_file};
+use crate::pr;
+use crate::yaml_maker::PROMETHEUS_SERVICE_NAME;
+
+fn port_forward(service_name: &str, local_port: u16, remote_port: u16, namespace: &str) {
+    let _ = Command::new("kubectl")
+        .args([
+            "port-forward",
+            &format!("service/{}", service_name),
+            &format!("{}:{}", local_port, remote_port),
+            "-n",
+            namespace,
+        ])
+        .stdout(Stdio::inherit())
+        .stderr(Stdio::inherit())
+        .status();
+}
+
+pub fn run() -> Result<(), String> {
+    let deployment_data = read_deployment_file(&cluster_deployment_file_path())?;
+
+    let namespace = deployment_data
+        .get("namespace")
+        .and_then(|n| n.as_str())
+        .ok_or("No namespace found in deployment file")?;
+
+    connect_to_cluster()?;
+
+    pr!("Port forwarding in namespace: {}", namespace);
+    pr!("  → Grafana:    http://localhost:3000");
+    pr!("  → Prometheus: http://localhost:9090");
+    pr!("\nPress Ctrl+C to stop\n");
+
+    let namespace_grafana = namespace.to_string();
+    let namespace_prometheus = namespace.to_string();
+
+    // Start both port-forwards in separate threads
+    let grafana_thread = thread::spawn(move || {
+        port_forward("grafana-service", 3000, 3000, &namespace_grafana);
+    });
+
+    let prometheus_thread = thread::spawn(move || {
+        port_forward(PROMETHEUS_SERVICE_NAME, 9090, 9090, &namespace_prometheus);
+    });
+
+    // Wait for threads (they'll run until interrupted)
+    let _ = grafana_thread.join();
+    let _ = prometheus_thread.join();
+
+    Ok(())
+}