The `query_graph` API

Overview

Once the infrastructure has been defined with set_graph and, optionally, extended with annotate_graph, the query_graph API is used to retrieve nodes, edges, or graph-level attributes that match a set of filters, or to resolve the shortest path between two nodes.

A single QueryRequest is a choice of exactly one of:

attribute_query — one or more named node/edge/graph filters, evaluated independently.
shortest_path_query — a source/destination pair to resolve via the graph's shortest path.

The QueryResponse mirrors this with a matching choice of attribute_query or shortest_path_query.

Attribute Query Requests

Query.Request.Filter accepts arrays of node_filters and edge_filters, plus a single graph_filter:

query = QueryRequest()
query.attribute_query.node_filters.add(name="...")   # 0 or more
query.attribute_query.edge_filters.add(name="...")   # 0 or more
query.attribute_query.graph_filter.attributes.add(attribute="...", value="...")

Every node/edge filter requires a unique name. Names are what let you tell the results of one filter apart from another when a request carries several — the response groups matches by that name rather than merging everything into one list.

Node and Edge Matching Rules

Each node_filter combines node_identifiers (or endpoints for edges) with attribute_filters. Both are optional, but at least one must be set:

`node_identifiers` / `endpoints`	`attribute_filters`	Result
unset	unset	No results.
set	unset	All matching nodes/edges, with every attribute.
unset	set	Every node/edge in the graph that matches the attribute filter.
set	set	Matching nodes/edges that also satisfy the attribute filter.

node_identifiers and edge endpoints support the same slicing operator used by annotate_graph, so a single entry like server[0:2]xpu[0:3] expands to every matching node.

attribute_filters takes one or more attribute/value pairs and an optional logic (and/or, defaults to and) to combine them.

attribute_filters is a single filter shared across every entry in node_identifiers/endpoints within that filter — it cannot be set differently per identifier. If you need different attribute criteria for different node identifiers or edge endpoints, use separate node_filters/edge_filters entries (each with its own name), one per distinct attribute criteria.

Default Schema Properties

Every node and edge carries a set of attributes derived directly from the infrastructure schema, in addition to anything added later via annotate_graph. attribute_filters can match against these the same way it matches user-added annotations.

Node (component) attributes:

Attribute	Description
`type`	The component type of this node (e.g. `xpu`, `nic`, `cpu`).
`device`	Name of the device model this component belongs to (e.g. `dgx_h100`).
`instance`	The parent instance path this component belongs to.
`instance_idx`	Numeric index of the parent instance.
`composed_device`	Full instance path this component was composed under.

Edge attributes:

Attribute	Description
`link`	Name of the link/interconnect connecting the two endpoints (e.g. `pcie`, `nvlink`).
`bandwidth`	Present when the link's physical bandwidth is defined in the schema, e.g. `"1400 Gbps"`.
`latency`	Present when the link's physical latency is defined in the schema, e.g. `"5 ns"`.

type, device, instance, instance_idx, composed_device, and link are immutable — annotate_graph rejects attempts to overwrite them. bandwidth and latency are not immutable and may be overwritten.

query = QueryRequest()
node_filter = query.attribute_query.node_filters.add(name="smart_nic_filter")
node_filter.node_identifiers = ["dgx_h100"]
node_filter.attribute_filters.attributes.add(attribute="cx7_type", value="smart")

query_response = service.query_graph(query)
result = query_response.attribute_query.nodes[0]
for node in result.nodes:
    print(node.name, {a.attribute: a.value for a in node.attributes})

query = QueryRequest()
edge_filter = query.attribute_query.edge_filters.add(name="nvlink_filter")
edge_filter.attribute_filters.attributes.add(attribute="link_type", value="nvlink")

query_response = service.query_graph(query)
result = query_response.attribute_query.edges[0]
for edge in result.edges:
    print(edge.ep1, edge.ep2, {a.attribute: a.value for a in edge.attributes})

query = QueryRequest()
query.attribute_query.graph_filter.attributes.add(attribute="region", value="us-east")

query_response = service.query_graph(query)
graph_attrs = {a.attribute: a.value for a in query_response.attribute_query.graph}

Multiple Filters in One Request

Because node_filters/edge_filters are arrays of named filters, a single request can ask for several independent things at once, and the response keeps each one's matches separate:

query = QueryRequest()

xpus = query.attribute_query.node_filters.add(name="xpus")
xpus.attribute_filters.attributes.add(attribute="type", value="xpu")

nics = query.attribute_query.node_filters.add(name="smart_nics")
nics.attribute_filters.attributes.add(attribute="cx7_type", value="smart")

query_response = service.query_graph(query)
for result in query_response.attribute_query.nodes:
    print(result.name, len(result.nodes))

Shortest Path Requests

query = QueryRequest()
query.shortest_path_query.name = "rank0-rank1"
query.shortest_path_query.source = service.get_endpoints("rank", "0")[0]
query.shortest_path_query.destination = service.get_endpoints("rank", "1")[0]

query_response = service.query_graph(query)
path = [node.name for node in query_response.shortest_path_query.nodes]

source and destination must be exact node IDs (not slice expressions) already present in the graph.

Full Examples

Node, edge, and graph attribute filter queries

import json
import pytest
import yaml
import uuid
import networkx
from infragraph import *
from datetime import datetime
from infragraph.blueprints.fabrics.clos_fat_tree_fabric import ClosFatTreeFabric
from infragraph.blueprints.fabrics.single_tier_fabric import SingleTierFabric
from infragraph.blueprints.devices.nvidia.dgx import NvidiaDGX
from infragraph.blueprints.devices.generic.generic_switch import Switch
from infragraph.blueprints.fabrics.closfabric import ClosFabric
from infragraph.infragraph_service import InfraGraphService
from infragraph.visualizer.visualize import run_visualizer

def print_graph(service):
    g = service.get_networkx_graph()
    # validations

    print("\nAnnotated node attributes:")
    for node, data in g.nodes(data=True):
        print(f"  {node}: {data}")

    print("\nAnnotated edge attributes:")
    for u, v, data in g.edges(data=True):
        print(f"  {u} -- {v}: {data}")

    print("\nAnnotated graph attributes:")
    print(f"  {g.graph}")

def visualize_dgx(infrastructure, annotation, output=None):
    """Build the annotated DGX infrastructure and launch the visualizer on it.

    A fresh output folder is generated on every call so repeated runs never
    clobber a prior visualization.
    """
    if output is None:
        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output = f"./viz_dgx_{stamp}_{uuid.uuid4().hex[:6]}"

    run_visualizer(infrastructure=infrastructure, annotations=annotation, output=output)

def _annotate_graph(service, **kwargs):
    annotation = Annotation()
    for attr, value in kwargs.items():
        annotation.graph.add(attribute=attr, value=str(value))
    service.annotate_graph(annotation)
    return annotation

def _annotate_node(service, node_name, **kwargs):
    annotation = Annotation()
    node_annotation = annotation.nodes.add(name=node_name)
    for attr, value in kwargs.items():
        node_annotation.attributes.add(attribute=attr, value=str(value))
    service.annotate_graph(annotation)
    return annotation

def _annotate_edge(service, ep1, ep2, **kwargs):
    annotation = Annotation()
    edge_annotation = annotation.edges.add(ep1=ep1, ep2=ep2)
    for attr, value in kwargs.items():
        edge_annotation.attributes.add(attribute=attr, value=str(value))
    service.annotate_graph(annotation)
    return annotation

def _annotate_topology(service):
    _annotate_graph(service, fabric="Single Tier Topology", hosts=["dgx_h100[0]", "dgx_h100[1]"], region="us-east")
    # annotate nodes here - the devices
    _annotate_node(service, node_name="dgx_h100[0]", location="rack 0")
    _annotate_node(service, node_name="dgx_h100[1]", location="rack 8")
    # switch asic name
    _annotate_node(service, node_name="switch[0]asic[0]", vendor="intel tofino")

    # add device type?
    _annotate_node(service, node_name="switch", device_type="switch")
    _annotate_node(service, node_name="dgx_h100", device_type="host")

    # set ranks?
    for i in range(0, 16):
        dev_index = 0 if i < 8 else 1
        comp_index = i % 8
        _annotate_node(service, node_name=f"dgx_h100[{str(dev_index)}]xpu[{str(comp_index)}]", rank=str(i))

    # set dgx 0 cx7 annotation to smart cx7
    _annotate_node(service, node_name="dgx_h100[0]cx7", cx7_type="smart cx7")

    # set dgx 1 cpu annotation to hyperthreaded
    _annotate_node(service, node_name="dgx_h100[1]cpu", cpu_type="hyper threaded RISC")

    # add annotation of nvlink to both edges
    _annotate_edge(service, ep1="dgx_h100[0]xpu", ep2="dgx_h100[0]nvsw", latency="0.01", link_type="nvlink", error_rate="6")

    _annotate_edge(service, ep1="dgx_h100[1]xpu", ep2="dgx_h100[1]nvsw", latency="0.08", link_type="nvlink", error_rate="20")

@pytest.fixture
def service():
    dgx = NvidiaDGX()
    clos_fat_tree = SingleTierFabric(dgx, 2)
    svc = InfraGraphService()
    svc.set_graph(clos_fat_tree)
    _annotate_topology(svc)
    return svc

def test_node_filter_attribute_query(service):
    query = QueryRequest()
    node_filter = query.attribute_query.node_filters.add(name="cpu_type_filter")
    node_filter.attribute_filters.attributes.add(attribute="cpu_type", value="hyper threaded")
    query_response = service.query_graph(query).attribute_query
    nodes = query_response.nodes[0].nodes
    assert query_response.nodes[0].name == "cpu_type_filter"
    assert len(nodes) == 2
    assert "dgx_h100.1.cpu." in nodes[0].name
    assert len(query_response.edges) == 0
    assert len(query_response.graph) == 0

def test_query_node_attribute(service):
    # get all smart nics
    query = QueryRequest()
    node_filter = query.attribute_query.node_filters.add(name="cx7_type_filter")
    node_filter.node_identifiers = ["dgx_h100"]
    node_filter.attribute_filters.attributes.add(attribute="cx7_type", value="smart")
    query_response = service.query_graph(query).attribute_query
    nodes = query_response.nodes[0].nodes
    assert len(nodes) == 8
    assert "dgx_h100" in nodes[0].name
    assert len(query_response.edges) == 0
    assert len(query_response.graph) == 0


def test_query_rank_node_attribute(service):
    # get all smart nics
    query = QueryRequest()
    node_filter = query.attribute_query.node_filters.add(name="rank_filter")
    node_filter.attribute_filters.attributes.add(attribute="rank", value="")
    query_response = service.query_graph(query).attribute_query
    nodes = query_response.nodes[0].nodes
    assert len(nodes) == 16
    assert "xpu" in nodes[0].name
    assert len(query_response.edges) == 0
    assert len(query_response.graph) == 0
    # print_graph(service)
    # visualize_dgx(service.infrastructure, None, "clos_visual")

def test_query_nic_node_attribute(service):
    # get all smart nics
    query = QueryRequest()
    node_filter = query.attribute_query.node_filters.add(name="nic_type_filter")
    node_filter.attribute_filters.attributes.add(attribute="type", value="nic")
    query_response = service.query_graph(query).attribute_query
    nodes = query_response.nodes[0].nodes
    assert len(nodes) == 16
    assert "cx7" in nodes[0].name
    assert len(query_response.edges) == 0
    assert len(query_response.graph) == 0

    # get specific smart nics
    query = QueryRequest()
    node_filter = query.attribute_query.node_filters.add(name="nic_type_filter")
    node_filter.node_identifiers = ["dgx_h100[1]"]
    node_filter.attribute_filters.attributes.add(attribute="type", value="nic")
    query_response = service.query_graph(query).attribute_query
    nodes = query_response.nodes[0].nodes
    assert len(nodes) == 8
    assert "cx7" in nodes[0].name
    assert len(query_response.edges) == 0
    assert len(query_response.graph) == 0
    # print_graph(service)
    # visualize_dgx(service.infrastructure, None, "clos_visual")

def test_edge_filter_attribute_query(service):
    query = QueryRequest()
    edge_filter = query.attribute_query.edge_filters.add(name="nvlink_filter")
    edge_filter.attribute_filters.attributes.add(attribute="link_type", value="nvlink")
    query_response = service.query_graph(query).attribute_query
    edges = query_response.edges[0].edges
    assert len(edges) == 64
    attrs = {a.attribute: a.value for a in edges[0].attributes}
    assert attrs["link_type"] == "nvlink"
    assert len(query_response.nodes) == 0
    assert len(query_response.graph) == 0

def test_query_graph_attribute(service):
    # get all smart nics
    query = QueryRequest()
    query.attribute_query.graph_filter.attributes.add(attribute="region", value="us-east")
    query_response = service.query_graph(query).attribute_query
    assert len(query_response.graph) > 0
    assert len(query_response.edges) == 0
    assert len(query_response.nodes) == 0


if __name__ == "__main__":
    pytest.main(["-s", __file__])

Shortest path query

from typing import Tuple
import pytest
from infragraph import *
from infragraph.blueprints.fabrics.closfabric import ClosFabric
from infragraph.infragraph_service import InfraGraphService


@pytest.mark.asyncio
@pytest.mark.parametrize("ranks", [(i, i + 1) for i in range(0, 7)])
async def test_shortest_path(ranks: Tuple[int, int]):
    """Test resolving the shortest path from one rank to another"""
    service = InfraGraphService()
    service.set_graph(ClosFabric().serialize())

    # add ranks
    npu_endpoints = service.get_endpoints("type", Component.XPU)
    annotation = Annotation()
    for idx, npu_endpoint in enumerate(npu_endpoints):
        annotation_node = annotation.nodes.add(
            name=npu_endpoint
        )
        annotation_node.attributes.add(attribute="rank", value=str(idx))
    service.annotate_graph(annotation.serialize())

    # find shortest path from one rank to another
    query = QueryRequest()
    query.shortest_path_query.name = "rank0-rank1"


    query.shortest_path_query.source = service.get_endpoints("rank", str(ranks[0]))[0]
    query.shortest_path_query.destination = service.get_endpoints("rank", str(ranks[1]))[0]

    query_response = service.query_graph(query)
    shortest_route = ""
    for node in query_response.shortest_path_query.nodes:
        shortest_route = shortest_route + " -> " + node.name

    print(f"\nShortest Path between rank {ranks[0]} and rank {ranks[1]}")
    print(shortest_route)


if __name__ == "__main__":
    pytest.main(["-s", __file__])

The `query_graph` API

Overview

Attribute Query Requests

Node and Edge Matching Rules

Default Schema Properties

Node Attribute Filter

Edge Attribute Filter

Graph Attribute Filter

Multiple Filters in One Request

Shortest Path Requests

Full Examples

The query_graph API

Overview

Attribute Query Requests

Node and Edge Matching Rules

Default Schema Properties

Node Attribute Filter

Edge Attribute Filter

Graph Attribute Filter

Multiple Filters in One Request

Shortest Path Requests

Full Examples

The `query_graph` API