Function evaluate

experiments/benchmark_TTFT.py:86–116

(
        graph: Graph,
        *,
        samples: int,
        output_dir: str,
        **kwargs
        )

async def evaluate(
        graph: Graph,
        *,
        samples: int,
        output_dir: str,
        **kwargs
        ) -> List[Dict[str, Any]]:

    graph.spatial_logits.requires_grad_ = False
    graph.temporal_logits.requires_grad_ = False

    data = [{"task": _make_random_token_sequence(1000)} for _ in range(samples)]

    all_results: List[Dict[str, Any]] = []
    for i, input_dict in enumerate(data):
        print(80*'-')

        realized_graph = copy.deepcopy(graph)
        realized_graph.spatial_logits = graph.spatial_logits
        realized_graph.temporal_logits = graph.temporal_logits
        tasks = [asyncio.create_task(realized_graph.arun(input_dict, **kwargs))]
        raw_results = await asyncio.gather(*tasks)
        all_results.extend(raw_results)
    print("Done!")

    try:
        _write_per_agent_latency(output_dir)
    except Exception as e:
        logger.warning("Failed to write per-agent latency JSONs: {}", e)

    return all_results


def _write_per_agent_latency(output_dir: str) -> None: