diff --git a/miles/rollout/generate_hub/agentic_tool_call.py b/miles/rollout/generate_hub/agentic_tool_call.py index feba568164..908361472a 100644 --- a/miles/rollout/generate_hub/agentic_tool_call.py +++ b/miles/rollout/generate_hub/agentic_tool_call.py @@ -112,6 +112,15 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput: for s in samples: s.metadata.update(agent_metadata or {}) + # If the agent function reports wall-clock time spent outside policy generation + # (env/tool steps), surface it on Sample.non_generation_time so throughput + # accounting subtracts it. Must be equal across all turn-samples: merge_samples + # collapses them with _merge_equal_value, which asserts the values match. + ngt = ((agent_metadata or {}).get("agent_metrics") or {}).get("total_tool_time") + if ngt is not None: + for s in samples: + s.non_generation_time = ngt + if max_seq_len is not None: samples = truncate_samples_by_total_tokens(samples, max_seq_len, input.state.tokenizer)