From 2b8c9367c8528d48032ed89a31c04498ca803f66 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 28 Aug 2026 23:58:58 +0800 Subject: [PATCH] perf: instrument mapnav episode timings --- .../services/retrieval/nav/nav_agent.py | 27 ++++++- .../shared/services/retrieval/nav/nav_llm.py | 74 ++++++++++++------- .../services/retrieval/nav/nav_orchestrate.py | 37 ++++++++-- 3 files changed, 104 insertions(+), 34 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py index f4d0d12d..197af5af 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_agent.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_agent.py @@ -1,7 +1,8 @@ from __future__ import annotations -import time +import logging import os +import time from typing import Any, List, Optional, Sequence, Tuple from ._compat import AgentStep, EpisodeResult @@ -37,6 +38,7 @@ # Back-compat aliases for tests / callers. _evidence_owner_section_id = evidence_owner_section_id _unit_score_for_evidence_chunk = unit_score_for_evidence_chunk +_logger = logging.getLogger(__name__) def _chunks_to_retrieved_nodes(chunks: List[Chunk]) -> List[str]: @@ -398,6 +400,7 @@ def _run_nav_episode_body( state = NavState(doc_id=episode_doc, query=query, task_type=task_type) steps: List[AgentStep] = [] + map_started = time.perf_counter() if namespace_mode: section_ids = list(ts.sections_for_doc("")) state.map_scores, state.unit_scores = compute_corpus_map_and_unit_scores( @@ -408,6 +411,12 @@ def _run_nav_episode_body( state.map_scores, state.unit_scores = compute_map_and_unit_scores( ts, doc_id=episode_doc, query=query, root_ids=section_ids ) + _logger.info( + "retrieval mapnav phase=map_scoring seconds=%.3f documents=%d sections=%d", + time.perf_counter() - map_started, + len(corpus_ids), + len(section_ids), + ) state.highlight_ids = select_map_highlights( state.unit_scores, k=int(cfg.collect_top_k) ) @@ -434,6 +443,11 @@ def _run_nav_episode_body( }, t0=plan_t0), ) ) + _logger.info( + "retrieval mapnav phase=planner seconds=%.3f subgoals=%d", + time.perf_counter() - plan_t0, + len(retrieval_plan.subgoals), + ) # Checklist: wave orchestration; navigate mode: classic single navigate. if cfg.is_checklist and state.retrieval_plan is not None: @@ -452,6 +466,11 @@ def _run_nav_episode_body( detail=stamp_step_detail(orch_detail, t0=orch_t0), ) ) + _logger.info( + "retrieval mapnav phase=orchestration seconds=%.3f waves=%d", + time.perf_counter() - orch_t0, + len(orch_detail.get("waves", [])), + ) else: navigate( ts, @@ -464,6 +483,7 @@ def _run_nav_episode_body( steps_out=steps, ) + evidence_started = time.perf_counter() fill = pack_nav_evidence( _dedupe_scored(list(state.collected)), ts, @@ -471,6 +491,11 @@ def _run_nav_episode_body( cfg, budget_chars=budget_chars, ) + _logger.info( + "retrieval mapnav phase=evidence_pack seconds=%.3f chunks=%d", + time.perf_counter() - evidence_started, + len(fill.kept_chunks), + ) scored_chunks = list(fill.scored_chunks) retrieval_seconds = time.perf_counter() - retrieval_t0 composed = "" diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py index f8ba3187..767c1675 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_llm.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_llm.py @@ -20,7 +20,9 @@ from __future__ import annotations +import logging import os +import time from contextlib import contextmanager from contextvars import ContextVar from typing import Any, Callable, Dict, Iterator, Optional, Sequence @@ -30,6 +32,7 @@ NavChatBackend = Callable[..., Dict[str, Any]] _backend: Optional[NavChatBackend] = None +_logger = logging.getLogger(__name__) _DS_DEFAULT_MODEL = "deepseek-v4-flash" _DEFAULT_PLANNER_THINK_MAX = 16384 @@ -190,25 +193,34 @@ def nav_chat( raise NavTokenLimit(used=nav_tokens_used(), limit=nav_token_limit()) merged_extra = _merge_thinking_extra(extra, role=thinking_role, model=model) + call_started = time.perf_counter() if _backend is not None: - result = _backend( - purpose=purpose, - messages=list(messages), - model=model, - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - thinking_role=thinking_role, - context=context, - api_key_env=api_key_env, - base_url_env=base_url_env, - timeout=timeout, - usage_tag=usage_tag, - ) - record_episode_tokens((result or {}).get("usage")) - return result + try: + result = _backend( + purpose=purpose, + messages=list(messages), + model=model, + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + extra=merged_extra, + thinking_role=thinking_role, + context=context, + api_key_env=api_key_env, + base_url_env=base_url_env, + timeout=timeout, + usage_tag=usage_tag, + ) + record_episode_tokens((result or {}).get("usage")) + return result + finally: + _logger.info( + "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", + purpose, + thinking_role, + time.perf_counter() - call_started, + ) from ._compat import cached_chat_completion # type: ignore from ._compat import ( # type: ignore @@ -230,16 +242,24 @@ def nav_chat( f"(model={model!r}; need DS_KEY for deepseek-* or OPENAI_API_KEY)." ) client = make_openai_client(api_key=key, base_url=base_url, timeout=timeout) - cached = cached_chat_completion( - client, - purpose=purpose, - model=model, - messages=list(messages), - temperature=temperature, - max_tokens=max_tokens, - response_format=response_format, - extra=merged_extra, - ) + try: + cached = cached_chat_completion( + client, + purpose=purpose, + model=model, + messages=list(messages), + temperature=temperature, + max_tokens=max_tokens, + response_format=response_format, + extra=merged_extra, + ) + finally: + _logger.info( + "retrieval mapnav llm_call purpose=%s role=%s seconds=%.3f", + purpose, + thinking_role, + time.perf_counter() - call_started, + ) if usage_tag: record_usage(usage_tag, cached.get("usage")) record_episode_tokens(cached.get("usage")) diff --git a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py index 91f45e1c..2e82be08 100644 --- a/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py +++ b/packages/shared-python/shared/services/retrieval/nav/nav_orchestrate.py @@ -8,14 +8,14 @@ from __future__ import annotations -from .nav_token_budget import stamp_step_detail - +import logging import re import time from contextlib import contextmanager from dataclasses import asdict from typing import Any, Dict, Iterator, List, Optional, Sequence, Set, Tuple +from .nav_token_budget import stamp_step_detail from .nav_navigate import navigate from .nav_plan import ( RetrievalPlan, @@ -29,6 +29,7 @@ from .nav_verify import apply_bindings_from_result, build_subgoal_result _SLOT_STRIP_RE = re.compile(r"\{\{\s*[^}]+\s*\}\}") +_logger = logging.getLogger(__name__) def ready_subgoal_ids( @@ -483,8 +484,18 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) ) # Serial wave execution (parallel fan-out retired with ThreadPoolExecutor). + wave_started = time.perf_counter() for sid in ready: - outputs.append(_run_one(sid, state, steps_out)) + subgoal_started = time.perf_counter() + try: + outputs.append(_run_one(sid, state, steps_out)) + finally: + _logger.info( + "retrieval mapnav harvest subgoal=%s wave=%d seconds=%.3f", + sid, + wave_idx, + time.perf_counter() - subgoal_started, + ) # Bookkeeping shared by both decision paths. for item in outputs: @@ -508,9 +519,17 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) state.subgoal_attempt_counts.get(sid, 0) ) + 1 - control_detail = _apply_plan_control( - ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out - ) + control_started = time.perf_counter() + try: + control_detail = _apply_plan_control( + ts, state, config, plan=plan, outputs=outputs, by_id=by_id, steps_out=steps_out + ) + finally: + _logger.info( + "retrieval mapnav plan_control wave=%d seconds=%.3f", + wave_idx, + time.perf_counter() - control_started, + ) wave_detail["plan_control"] = control_detail replan_requested = bool(control_detail.get("replan")) if control_detail.get("done"): @@ -525,6 +544,12 @@ def _run_one(sid: str, working_state: NavState, out_steps: Optional[List[Any]]) ) ) summary["waves"].append(wave_detail) + _logger.info( + "retrieval mapnav wave=%d ready=%d seconds=%.3f", + wave_idx, + len(ready), + time.perf_counter() - wave_started, + ) if replan_requested: cap = int(getattr(config, "max_replans", 0) or 0)