Keysight AI DCB collective communications

[lthompson-keysight]

Snapshot of Keysight's model work, from internal commit 5a19c55eb586, based on upstream c48c7ea. Generated artifacts are built by CI from the sources in this commit.

Feature Overview

• Related Issue: N/A
• Brief Description:
  Add support for benchmarking collective communications systems

---

Feature Details

• Redocly docs for this feature/branch
• New Locations/Nomenclature:

config body: /collective_communications
config body: /artifacts
set_control_state body: /collective_communications
get_state body: /collective_communications
get_metrics body: /collective_communications
new api: generate_config

• Dev-Snappi branch reference

---

Code snippets

"""
Basic sample for running a collective benchmark trial via the OTG API.
"""

import time
from typing import Literal

import pandas as pd
from pandas.core.frame import DataFrame

import snappi
from snappi import (
    Api,
    CollectiveCommunicationsStateStopped,
    Config,
    CollectiveCommunicationsInfrastructureRef,
    MetricsResponse,
)

from keysight_dse_common.schemas.summary_metrics_schema import (
    SummaryMetricsColumnIds,
)

# Metric types selectable under the single collective application.
SingleCollectiveMetric = Literal[
    "summary",
    "data_chunk",
    "flows",
    "qpairs",
    "ports",
    "iterations_summary",
    "iteration",
    "packet_drop",
]

DSE_SERVER_ADDRESS = "localhost:50001"
NUM_HOSTS = 4
XPUS_PER_HOST = 8
CHASSIS_ADDRESS = "1.2.3.4"
PLATFORM_NAME = "hw"
IP_ADDRESS_TEMPLATE = "32.0.{rank}.1"
IP_GATEWAY_TEMPLATE = "32.0.{rank}.254"
IP_PREFIX = 24
HOST_NAME = "generic_host"

# Workload-level RoCEv2 transport settings, applied uniformly to every rank
# pair using this transport.
ROCEV2_RDMA_MESSAGE_SIZE = 1 << 17  # 128 KiB
ROCEV2_VERB = "write"
ROCEV2_REUSE_QPS = True
TRANSPORT_QPS_PER_RANK_PAIR = 1


def fill_in_nic_infra_ref(
    out_ref: CollectiveCommunicationsInfrastructureRef, host_index: int, xpu_index: int
) -> None:
    """Fill in an InfraRef structure with the data for refering to a NIC within the topology."""
    out_ref.device_instance_name = HOST_NAME
    out_ref.device_index = host_index
    out_ref.component_name = "nic"
    out_ref.component_index = xpu_index


def fill_in_xpu_infra_ref(
    out_ref: CollectiveCommunicationsInfrastructureRef, host_index: int, xpu_index: int
) -> None:
    """Fill in an InfraRef structure with the data for refering to an XPU within the topology."""
    out_ref.device_instance_name = HOST_NAME
    out_ref.device_index = host_index
    out_ref.component_name = "xpu"
    out_ref.component_index = xpu_index


def set_up_infrastructure(mut_config: Config) -> None:
    cc = mut_config.collective_communications

    bindings = cc.bindings.basic_binding

    rank_id = 0
    for host_index in range(NUM_HOSTS):
        device = mut_config.devices.add()
        device.name = f"host{host_index}-device"

        for xpu_index in range(XPUS_PER_HOST):
            # Register test chassis port that will be used for this rank's traffic
            port = mut_config.ports.add()
            port.name = f"host{host_index}-nic{xpu_index}-port"
            port.location = f"{CHASSIS_ADDRESS};{rank_id + 1}"

            # Configure the IPv4 settings for the port
            ethernet = device.ethernets.add()
            ethernet.name = f"host{host_index}-nic{xpu_index}-eth"
            ethernet.connection.choice = ethernet.connection.PORT_NAME
            ethernet.connection.port_name = port.name
            ipv4 = ethernet.ipv4_addresses.add()
            ipv4.name = f"{ethernet.name}-ipv4"
            ipv4.address = IP_ADDRESS_TEMPLATE.format(rank=rank_id)
            ipv4.prefix = IP_PREFIX
            ipv4.gateway = IP_GATEWAY_TEMPLATE.format(rank=rank_id)

            # Associate the rank with the XPU and connect the emulated XPU to a emulated NIC
            rank_binding = bindings.rank_bindings.add()
            rank_binding.rank_id = rank_id
            fill_in_xpu_infra_ref(
                rank_binding.infrastructure_ref, host_index, xpu_index
            )
            fill_in_nic_infra_ref(rank_binding.nic_refs.add(), host_index, xpu_index)

            # Associate the emulated NIC to a defined ethernet with its settings
            nic_binding = bindings.nic_bindings.add()
            fill_in_nic_infra_ref(nic_binding.infrastructure_ref, host_index, xpu_index)
            nic_binding.ethernet_name = ethernet.name

            # Associate the emulated NIC to the previously defined physical port
            physical_binding = bindings.physical_bindings.add()
            fill_in_nic_infra_ref(
                physical_binding.infrastructure_ref, host_index, xpu_index
            )
            physical_binding.platform_name = PLATFORM_NAME
            physical_binding.port_name = port.name

            rank_id += 1

    # Configure Layer 1 settings for all ports
    layer1 = mut_config.layer1.add()
    layer1.name = "all-ports-l1"
    layer1.port_names = [port.name for port in mut_config.ports]
    layer1.speed = layer1.SPEED_400_GBPS
    layer1.signaling = layer1.PAM4_106_GBPS
    layer1.fec_mode = layer1.RS_FEC_KP4


def set_up_benchmark(mut_config: Config) -> None:
    cc = mut_config.collective_communications

    platform = cc.platforms.add()
    platform.name = PLATFORM_NAME
    platform.platform_type = "keysight_hw"

    single_collective = cc.workload.single_collective

    algorithm_type = single_collective.collective_algorithm.algorithm_type
    algorithm_type.system = algorithm_type.ALL_REDUCE_VECTOR_HALVING_DOUBLING

    single_collective.datasizes.list.size_bytes = [
        1 << 27,
        1 << 30,
        1 << 32,
    ]  # 128 MiB, 1 GiB, 4 GiB

    single_collective.iterations = 10

    rocev2 = single_collective.transport.rocev2
    rocev2.rdma_message_size = ROCEV2_RDMA_MESSAGE_SIZE
    rocev2.verb = rocev2.WRITE
    rocev2.reuse_qps = ROCEV2_REUSE_QPS
    rocev2.qps_per_rankpair = TRANSPORT_QPS_PER_RANK_PAIR


def run_and_wait_for_result(
    api: Api, poll_interval_s: float = 2.0
) -> CollectiveCommunicationsStateStopped:
    control_state = api.control_state()
    execution = control_state.collective_communications.execution
    execution.state = execution.START
    api.set_control_state(control_state)

    states_request = api.states_request()
    states_request.collective_communications.choice = (
        states_request.collective_communications.EXECUTION
    )

    while True:
        ai_state = api.get_states(states_request).collective_communications.execution
        if (
            ai_state.choice == ai_state.STOPPED
            and ai_state.stopped.choice != ai_state.stopped.UNRUN
        ):
            return ai_state.stopped
        time.sleep(poll_interval_s)


def print_result(stopped: CollectiveCommunicationsStateStopped) -> None:
    match stopped.choice:
        case stopped.SUCCESS:
            warnings = stopped.success.warnings or []
            print(f"Run succeeded with {len(warnings)} warning(s):")
            for warning in warnings:
                print(f"  - {warning}")
        case stopped.FAILURE:
            print("Run failed:")
            for warning in stopped.failure.warnings or []:
                print(f"  - {warning}")
        case stopped.MANUAL:
            print("Run was stopped manually.")
        case _:
            print(f"Run stopped with unexpected reason: {stopped.choice}")


def fetch_collective_communications_metrics(
    api: Api,
    metric_selection: SingleCollectiveMetric,
    column_names: list[str],
    pagination_start: int,
    pagination_count: int,
) -> MetricsResponse:
    """Fetch collective benchmark metrics of the selected type for the given page of rows."""
    request = api.metrics_request()
    request.choice = request.COLLECTIVE_COMMUNICATIONS

    cc = request.collective_communications
    cc.choice = cc.SINGLE_COLLECTIVE

    single_collective = cc.single_collective
    single_collective.choice = metric_selection
    getattr(single_collective, metric_selection).column_names = column_names

    request.pagination.start = pagination_start
    request.pagination.count = pagination_count

    return api.get_metrics(request)


def fetch_port_metrics(
    api: Api,
    port_names: list[str],
    column_names: list[str],
    pagination_start: int,
    pagination_count: int,
) -> MetricsResponse:
    """Fetch generic OTG port metrics for the given ports and page of rows."""
    request = api.metrics_request()
    request.choice = request.PORT

    port = request.port
    port.port_names = port_names
    port.column_names = column_names

    request.pagination.start = pagination_start
    request.pagination.count = pagination_count

    return api.get_metrics(request)

# Maps each CollectiveCommunications.Summary.Metric field to its summary schema column id.
# All names match the column id member name except pfc_rx, whose member is pfc_fx.
SUMMARY_FIELD_TO_COLUMN_ID = {
    "data_size": SummaryMetricsColumnIds.data_size,
    "collective": SummaryMetricsColumnIds.collective,
    "iterations": SummaryMetricsColumnIds.iterations,
    "completion_time": SummaryMetricsColumnIds.completion_time,
    "algo_bw": SummaryMetricsColumnIds.algo_bw,
    "bus_bw": SummaryMetricsColumnIds.bus_bw,
    "ideal": SummaryMetricsColumnIds.ideal,
    "pfc_rx": SummaryMetricsColumnIds.pfc_fx,
    "ecn_ce_rx": SummaryMetricsColumnIds.ecn_ce_rx,
    "frames_retx": SummaryMetricsColumnIds.frames_retx,
    "total_time": SummaryMetricsColumnIds.total_time,
    "duty_cycle": SummaryMetricsColumnIds.duty_cycle,
    "total_bytes_tx": SummaryMetricsColumnIds.total_bytes_tx,
    "total_bytes_rx": SummaryMetricsColumnIds.total_bytes_rx,
    "cnp_tx": SummaryMetricsColumnIds.cnp_tx,
    "cnp_rx": SummaryMetricsColumnIds.cnp_rx,
}


def summary_metrics_to_dataframe(response: MetricsResponse) -> pd.DataFrame:
    """Convert the collective benchmark summary metrics in a response into a DataFrame
    whose columns are the summary schema display names (SummaryMetricsColumnIds)."""
    summary_metrics = (
        response.collective_communications_metrics.single_collective.summary
    )

    schema_columns = [
        column_id.value for column_id in SUMMARY_FIELD_TO_COLUMN_ID.values()
    ]
    rows = [
        {
            column_id.value: getattr(summary_metric, field_name)
            for field_name, column_id in SUMMARY_FIELD_TO_COLUMN_ID.items()
        }
        for summary_metric in summary_metrics
    ]
    return pd.DataFrame(rows, columns=schema_columns)


def main() -> None:
    api = snappi.api(location=DSE_SERVER_ADDRESS, transport="grpc")

    config = api.config()
    set_up_infrastructure(config)
    set_up_benchmark(config)
    api.set_config(config)

    print_result(run_and_wait_for_result(api))

    summary_metrics_response = fetch_collective_communications_metrics(
        api,
        metric_selection="summary",
        column_names=[],
        pagination_start=0,
        pagination_count=100,
    )
    summary_dataframe: DataFrame = summary_metrics_to_dataframe(
        summary_metrics_response
    )

    port_metrics_response: MetricsResponse = fetch_port_metrics(
        api,
        port_names=["1.1", "1.2"],
        column_names=["frames_tx_rate", "frames_rx_rate"],
        pagination_start=0,
        pagination_count=100,
    )


if __name__ == "__main__":
    main()

---

[JoyChWang]

Not sure if flow_control_config really depend on which algorithm is chosen. Would it make more sense to make it a sibling of collective_algorithm and transport instead of nested inside Algorithm?

[JoyChWang]

For a dual-stack device, how to specify the IP version for the Rocev2?

[lthompson-keysight]

So, the otg rocev2 device settings define a list of IPv4 peers and a list of IPv6 peers. Each of those peers has the name of an otg ipv4 or ipv6 device (that is defined on the ipv4_interfaces and/or ipv6_interfaces of an ethernets device). The nic bindings have the name of the otg ethernets device, which can have both ipv4 and ipv6 devices. If the transport is rocev2, then it uses the rocev2 peer wired up to the named tcp device on that ethernets device. Otherwise, if it's a tcp transport, then it uses the ethernet.ipv4/ipv6_interfaces directly.

Here's an ascii diagram:
dual_stack_otg.txt

It's a bit complicated, but this is the "otg way"

[JoyChWang]

There are two separate places that specify the connection type — Rocev2.ConnectionType in device/rocev2/qps.yaml (per-QP) and Rocev2.QPConnectionType in port/rocev2.yaml (per-port) — what happens if they're configured differently? Also, is there a plan to add other connection types (e.g. unreliable_datagram) to either of them?

[lthompson-keysight]

The port values are overrides. In other words the connection type in port/rocev2 overrides the value in device/rocev2 for the specified ports. Any unlisted ports get the value from device/rocev2.

[JoyChWang]

What is the difference between completion_time and total_time in CollectiveCommunications.SingleCollective.Summary.Metric.
Since total_time covers all iterations, is completion_time meant to be the average per-iteration time?

[lthompson-keysight]

This slipped through. We must not change the existing OTG definitions (adding to them with defaultable fields is ok).