diff --git a/workers/executor/executors/constants.py b/workers/executor/executors/constants.py index 9eddab8423..868a47c773 100644 --- a/workers/executor/executors/constants.py +++ b/workers/executor/executors/constants.py @@ -51,6 +51,11 @@ class PromptServiceConstants: CHALLENGE = "challenge" ENABLE_CHALLENGE = "enable_challenge" EXTRACTION = "extraction" + # Reserved namespace for pipeline-level metrics so they cannot collide with + # user-defined output/prompt names at the top level of the metrics dict. + PIPELINE = "_pipeline" + # Wire-format key every metrics producer reports its duration under. + TIME_TAKEN = "time_taken(s)" SUMMARIZE = "summarize" SINGLE_PASS_EXTRACTION = "single-pass-extraction" SIMPLE_PROMPT_STUDIO = "simple-prompt-studio" diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index ce7fbea0d1..07b7d8c4f1 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -619,10 +619,12 @@ def _failure(child_result: ExecutionResult) -> ExecutionResult: ) step = 1 + extraction_metrics: dict = {} try: # ---- Step 1: Extract ---- if not skip_extraction: step += 1 + extraction_start = time.monotonic() extract_ctx = ExecutionContext( executor_name=context.executor_name, operation=Operation.EXTRACT.value, @@ -640,6 +642,15 @@ def _failure(child_result: ExecutionResult) -> ExecutionResult: return _failure(extract_result) _absorb(extract_result) extracted_text = extract_result.data.get(IKeys.EXTRACTED_TEXT, "") + extraction_time = time.monotonic() - extraction_start + # Nest under a reserved "_pipeline" namespace so the metric never + # collides with a user-defined output named "extraction" during + # the top-level metrics merge (see _merge_pipeline_metrics). + extraction_metrics = { + PSKeys.PIPELINE: { + PSKeys.EXTRACTION: {PSKeys.TIME_TAKEN: extraction_time} + } + } # ---- Step 2: Summarize (if enabled) ---- if is_summarization: @@ -700,6 +711,7 @@ def _failure(child_result: ExecutionResult) -> ExecutionResult: source_file_name=source_file_name, extracted_text=extracted_text, index_metrics=index_metrics, + extraction_metrics=extraction_metrics, ) output_map = structured_output.get(PSKeys.OUTPUT, {}) or {} @@ -800,17 +812,29 @@ def _finalize_pipeline_result( source_file_name: str, extracted_text: str, index_metrics: dict, + extraction_metrics: dict[str, dict] | None = None, ) -> None: - """Populate metadata/metrics in structured_output after pipeline completion.""" + """Populate metadata/metrics in structured_output after pipeline completion. + + Args: + extraction_metrics: Pipeline-level extraction timing, shaped as + ``{"_pipeline": {"extraction": {"time_taken(s)": float}}}``. + Nested under the reserved ``_pipeline`` namespace to avoid + colliding with user-defined output names during the merge. + ``None``/empty when extraction is skipped. + """ if "metadata" not in structured_output: structured_output["metadata"] = {} structured_output["metadata"]["file_name"] = source_file_name if extracted_text: structured_output["metadata"]["extracted_text"] = extracted_text - if index_metrics: + new_metrics = self._merge_pipeline_metrics( + index_metrics or {}, extraction_metrics or {} + ) + if new_metrics: existing_metrics = structured_output.get("metrics", {}) structured_output["metrics"] = self._merge_pipeline_metrics( - existing_metrics, index_metrics + existing_metrics, new_metrics ) def _run_pipeline_summarize( @@ -955,8 +979,6 @@ def _index_pipeline_output( index_records: list[dict], ) -> None: """Index a single structure-pipeline output entry in-place.""" - import datetime - chunk_size = output.get("chunk-size", 0) if chunk_size == 0: return @@ -977,7 +999,7 @@ def _index_pipeline_output( return seen_params.add(param_key) - indexing_start = datetime.datetime.now() + indexing_start = time.monotonic() logger.info( "Pipeline indexing: chunk_size=%s chunk_overlap=%s vector_db=%s", chunk_size, @@ -1029,9 +1051,9 @@ def _index_pipeline_output( if child_records: index_records.extend(child_records) - elapsed = (datetime.datetime.now() - indexing_start).total_seconds() + elapsed = time.monotonic() - indexing_start output_name = output.get("name", "") - index_metrics[output_name] = {"indexing": {"time_taken(s)": elapsed}} + index_metrics[output_name] = {"indexing": {PSKeys.TIME_TAKEN: elapsed}} @staticmethod def _merge_pipeline_metrics(metrics1: dict, metrics2: dict) -> dict: diff --git a/workers/tests/test_phase5d.py b/workers/tests/test_phase5d.py index 8b3252c9e6..72f1715242 100644 --- a/workers/tests/test_phase5d.py +++ b/workers/tests/test_phase5d.py @@ -304,6 +304,126 @@ def test_index_metrics_merged(self, executor): assert "indexing" in metrics["field_a"] +# --------------------------------------------------------------------------- +# Tests — Extraction timing metric (UN-2771) +# --------------------------------------------------------------------------- + + +class TestExtractionMetrics: + """Extraction duration is reported under the reserved _pipeline namespace. + + The metric is nested as ``metrics["_pipeline"]["extraction"]`` rather than a + bare top-level ``"extraction"`` key, because sibling keys in this dict are + user-defined output/prompt names — a prompt literally named "extraction" + would otherwise collide with it. + """ + + def _executor_with_extraction(self, executor, answer_metrics=None): + executor._handle_extract = MagicMock( + return_value=ExecutionResult( + success=True, data={"extracted_text": "text"} + ) + ) + executor._handle_index = MagicMock( + return_value=ExecutionResult(success=True, data={"doc_id": "d1"}) + ) + executor._handle_answer_prompt = MagicMock( + return_value=ExecutionResult( + success=True, + data={"output": {}, "metrics": answer_metrics or {}}, + ) + ) + return executor + + def _run(self, executor): + ctx = _make_pipeline_context({ + "extract_params": _base_extract_params(), + "index_template": _base_index_template(), + "answer_params": _base_answer_params(), + "pipeline_options": _base_pipeline_options(), + }) + return executor._handle_structure_pipeline(ctx) + + def test_extraction_metrics_recorded(self, executor): + """A normal run reports a non-negative float extraction duration.""" + self._executor_with_extraction(executor) + result = self._run(executor) + + assert result.success + time_taken = result.data["metrics"]["_pipeline"]["extraction"][ + "time_taken(s)" + ] + assert isinstance(time_taken, float) + assert time_taken >= 0 + + def test_extraction_metric_does_not_collide_with_output_name(self, executor): + """A prompt named "extraction" coexists with the pipeline metric.""" + self._executor_with_extraction( + executor, + answer_metrics={"extraction": {"llm": {"time_taken(s)": 2.0}}}, + ) + result = self._run(executor) + + metrics = result.data["metrics"] + # The user's prompt metrics are untouched... + assert metrics["extraction"] == {"llm": {"time_taken(s)": 2.0}} + # ...and the pipeline timing lives in its own namespace. + assert "time_taken(s)" in metrics["_pipeline"]["extraction"] + + def test_index_and_extraction_metrics_merged(self, executor): + """Per-output indexing and pipeline extraction metrics coexist.""" + self._executor_with_extraction( + executor, + answer_metrics={"field_a": {"llm": {"time_taken(s)": 2.0}}}, + ) + executor._run_pipeline_index = MagicMock( + return_value=({"field_a": {"indexing": {"time_taken(s)": 0.5}}}, []) + ) + result = self._run(executor) + + metrics = result.data["metrics"] + assert "llm" in metrics["field_a"] + assert "indexing" in metrics["field_a"] + assert "time_taken(s)" in metrics["_pipeline"]["extraction"] + + def test_skip_extraction_records_no_extraction_metric(self, executor): + """Smart-table path skips extract, so no extraction metric is added.""" + executor._handle_extract = MagicMock() + executor._handle_index = MagicMock() + executor._handle_answer_prompt = MagicMock( + return_value=ExecutionResult(success=True, data={"output": {}}) + ) + opts = _base_pipeline_options() + opts["skip_extraction_and_indexing"] = True + + ctx = _make_pipeline_context({ + "extract_params": _base_extract_params(), + "index_template": _base_index_template(), + "answer_params": _base_answer_params(), + "pipeline_options": opts, + }) + result = executor._handle_structure_pipeline(ctx) + + assert result.success + executor._handle_extract.assert_not_called() + assert "_pipeline" not in result.data.get("metrics", {}) + + def test_extract_failure_records_no_extraction_metric(self, executor): + """Timing is taken after the failure early-return, so a failed extract + must not report a duration.""" + executor._handle_extract = MagicMock( + return_value=ExecutionResult.failure(error="x2text error") + ) + executor._handle_index = MagicMock() + executor._handle_answer_prompt = MagicMock() + executor._finalize_pipeline_result = MagicMock() + + result = self._run(executor) + + assert not result.success + executor._finalize_pipeline_result.assert_not_called() + + # --------------------------------------------------------------------------- # Tests — Embedding usage record propagation # ---------------------------------------------------------------------------