[FEATURE] Declare the restart parameters on extract() and document what a restart costs

docs-site/src/content/docs/lexical-graph/indexing.mdx

@@ -244,6 +244,23 @@ The run records three things under `s3://<bucket>/<key_prefix>/<collection_id>/_
 
 The plan does not yet cover the auto-tuned batch extraction path; a run with a `run_id` and `auto_tune=True` raises a `ValueError`.
 
+**What a restart costs.** A restart of a run whose Bedrock batch jobs have all finished costs seconds. The figures below cover the extract stage only: NTSB accident narratives extracted with Haiku 4.5 on Bedrock batch inference in `us-west-2`, staged to S3. The 100 and 500-document runs used 2 extraction workers, the 1,000-document runs 8.
+
+| Documents | Restarts | First run | All restarts | Per restart |
+| ------------- | ------------- | ------------- | ------------- | ------------- |
+| 100 | 10 | 3h 8m 35s | 2m 21s | 14.1 s |
+| 100 | 99 (one per document) | 2h 29m 28s | 21m 30s | 13.0 s |
+| 500 | 10 | 10h 11m 24s | 5m 14s | 31.4 s |
+| 1,000 | 10 | 8h 43m 22s | 8m 5s | 48.5 s |
+
+Three things to know before quoting these figures:
+
+- **Each per-restart figure includes that run's share of staging**, which an uninterrupted run also pays, so it overstates restart overhead. At 100 documents about 12.9 s of the 14.1 s is overhead; at 500 documents a restart stages about 45 documents for the first time.
+- **The two worker counts are not comparable with each other.** The rise from 14.1 s to 48.5 s spans a change of worker count and compute as well as a change of corpus size.
+- **Use the absolute cost, not a percentage of the run.** Wall clock is set by Bedrock batch queue and job time: the same 1,000-document extraction, same documents and the same 11,131 chunks, took 8h 43m once and 3h 11m another time. The 8m 5s of restarts is 1.5% of the first and 4.2% of the second.
+
+Every interruption in these runs came after the last Bedrock batch job had finished, so the figures cover downloading completed job output and staging, not waiting on a running job or submitting one again. Each run was interrupted by raising inside staging, which waits for chunk uploads already in flight; a process killed outright does not, so the figures are slightly optimistic against a real crash. All runs used batch inference; synchronous extraction was not measured.
+
 The identity under which the lexical-graph runs needs `s3:PutObject`, `s3:GetObject` and `s3:DeleteObject` on `arn:aws:s3:::<bucket>/<key_prefix>/*`, and `s3:ListBucket` on `arn:aws:s3:::<bucket>`. `s3:DeleteObject` is used when a run with a `run_id` finishes and its partition records are folded into the rollup. Both stages list the objects under a collection: the extract stage to decide what is already stored, the build stage to read a document's chunks and its markers. Without `s3:ListBucket` these listings fail with an access denied error.
 
 A restart trusts what it reads under this prefix: the plan, the partition records and the completion markers. Anyone who can write there can make a restart reuse other batch output in the same bucket, or skip a document it should store. Limit write access to the identity that runs the lexical-graph. The bucket always comes from your configuration, never from a record.

lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py

@@ -25,6 +25,7 @@
 from graphrag_toolkit.lexical_graph.indexing.extract import TopicExtractor, BatchTopicExtractorSync
 from graphrag_toolkit.lexical_graph.indexing.extract import ExtractionPipeline
 from graphrag_toolkit.lexical_graph.indexing.extract import InferClassifications, InferClassificationsConfig
+from graphrag_toolkit.lexical_graph.indexing.extract import RunPlanStore
 from graphrag_toolkit.lexical_graph.indexing.build import BuildPipeline
 from graphrag_toolkit.lexical_graph.indexing.build import VectorIndexing
 from graphrag_toolkit.lexical_graph.indexing.build import GraphConstruction
@@ -448,6 +449,8 @@ def extract(
             checkpoint: Optional[Checkpoint] = None,
             show_progress: Optional[bool] = False,
             progress_monitor: Optional[ProgressMonitor] = None,
+            run_id: Optional[str] = None,
+            run_plan_store: Optional[RunPlanStore] = None,
             **kwargs: Any) -> None:
         """
         Executes the extraction process for a given set of nodes using the specified handler,
@@ -467,6 +470,13 @@ def extract(
                 callbacks as documents complete LLM extraction. Receives
                 increment_llm_processed_documents and increment_llm_processed_chunks calls at
                 document boundaries.
+            run_id (Optional[str]): Id that makes the run restartable. Extracting again with the
+                same run_id takes up where the interrupted run stopped. Needs run_plan_store, and
+                handler=run_plan_store.staging_handler(run_id) so a restart also skips the
+                documents already staged. The plan does not yet cover the auto-tuned batch
+                extraction path; a run with a run_id and auto_tune=True raises a ValueError.
+            run_plan_store (Optional[RunPlanStore]): Where the run records its plan and the Bedrock
+                batch job it submitted for each partition, so a restart can reuse them.
             **kwargs (Any): Additional keyword arguments for pipeline configurations or overrides.
 
         Example::
@@ -496,6 +506,8 @@ def increment_llm_processed_documents(self, count=1):
             tenant_id=DEFAULT_TENANT_ID,
             extraction_filters=self.indexing_config.extraction.extraction_filters,
             source_id_width=self._source_id_width(),
+            run_id=run_id,
+            run_plan_store=run_plan_store,
             **kwargs
         )
 
@@ -595,6 +607,8 @@ def extract_and_build(
             checkpoint: Optional[Checkpoint] = None,
             show_progress: Optional[bool] = False,
             progress_monitor: Optional[ProgressMonitor] = None,
+            run_id: Optional[str] = None,
+            run_plan_store: Optional[RunPlanStore] = None,
             **kwargs: Any
     ) -> None:
         """
@@ -612,6 +626,13 @@ def extract_and_build(
                 callbacks across both extraction and build stages. The same instance receives
                 LLM increment calls during extraction, then graph/vector increment calls
                 during build.
+            run_id (Optional[str]): Id that makes the extraction restartable. Running again with
+                the same run_id reuses the Bedrock batch jobs the interrupted run submitted; the
+                build stage runs over every document again. Needs run_plan_store. The plan does not yet cover the
+                auto-tuned batch extraction path; a run with a run_id and auto_tune=True raises a
+                ValueError.
+            run_plan_store (Optional[RunPlanStore]): Where the run records its plan and the Bedrock
+                batch job it submitted for each partition, so a restart can reuse them.
             **kwargs (Any): Additional parameters to pass to pipelines.
         """
 
@@ -629,6 +650,8 @@ def extract_and_build(
             tenant_id=DEFAULT_TENANT_ID,
             extraction_filters=self.indexing_config.extraction.extraction_filters,
             source_id_width=self._source_id_width(stored_width),
+            run_id=run_id,
+            run_plan_store=run_plan_store,
             **kwargs
         )
 

lexical-graph/tests/unit/test_lexical_graph_index.py

@@ -103,3 +103,81 @@ def test_init_propagates_graph_store_init_failure(self):
                 LexicalGraphIndex(graph_store='dummy://', vector_store='dummy://')
 
         graph_store.init.assert_called_once_with()
+
+
+class TestRestartParameters:
+    """run_id and run_plan_store are named parameters of extract() and
+    extract_and_build(), and reach the extraction pipeline only. The build
+    pipeline has no use for them."""
+
+    MODULE = 'graphrag_toolkit.lexical_graph.lexical_graph_index'
+
+    @staticmethod
+    def _index():
+        with (
+            patch(f'{TestRestartParameters.MODULE}.GraphStoreFactory.for_graph_store', return_value=Mock()),
+            patch(f'{TestRestartParameters.MODULE}.MultiTenantGraphStore.wrap', return_value=Mock()),
+            patch(f'{TestRestartParameters.MODULE}.VectorStoreFactory.for_vector_store', return_value=Mock()),
+            patch(f'{TestRestartParameters.MODULE}.MultiTenantVectorStore.wrap', return_value=Mock()),
+            patch.object(LexicalGraphIndex, '_configure_extraction_pipeline', return_value=([], [])),
+        ):
+            return LexicalGraphIndex(graph_store='dummy://', vector_store='dummy://')
+
+    def _extract(self, **kwargs):
+        index = self._index()
+
+        with (
+            patch.object(LexicalGraphIndex, '_source_id_width', return_value=None),
+            patch(f'{self.MODULE}.ExtractionPipeline.create') as extraction_create,
+            patch(f'{self.MODULE}.BuildPipeline.create') as build_create,
+        ):
+            index.extract([], **kwargs)
+
+        return extraction_create.call_args.kwargs, build_create.call_args.kwargs
+
+    def test_extract_passes_the_restart_parameters_to_the_extraction_pipeline(self):
+        run_plan_store = Mock()
+
+        extraction_kwargs, build_kwargs = self._extract(run_id='a-run', run_plan_store=run_plan_store)
+
+        assert extraction_kwargs['run_id'] == 'a-run'
+        assert extraction_kwargs['run_plan_store'] is run_plan_store
+        assert 'run_id' not in build_kwargs
+        assert 'run_plan_store' not in build_kwargs
+
+    def test_extract_without_restart_parameters_passes_none(self):
+        extraction_kwargs, _ = self._extract()
+
+        assert extraction_kwargs['run_id'] is None
+        assert extraction_kwargs['run_plan_store'] is None
+
+    def _extract_and_build(self, **kwargs):
+        index = self._index()
+
+        with (
+            patch(f'{self.MODULE}.read_stored_source_id_width', return_value=None),
+            patch.object(LexicalGraphIndex, '_source_id_width', return_value=None),
+            patch(f'{self.MODULE}.GraphConstruction.for_graph_store', return_value=Mock()),
+            patch(f'{self.MODULE}.VectorIndexing.for_vector_store', return_value=Mock()),
+            patch(f'{self.MODULE}.ExtractionPipeline.create') as extraction_create,
+            patch(f'{self.MODULE}.BuildPipeline.create') as build_create,
+        ):
+            index.extract_and_build([], enable_versioning=False, **kwargs)
+
+        return extraction_create.call_args.kwargs, build_create.call_args.kwargs
+
+    def test_extract_and_build_passes_the_restart_parameters_to_the_extraction_pipeline(self):
+        run_plan_store = Mock()
+
+        extraction_kwargs, build_kwargs = self._extract_and_build(run_id='a-run', run_plan_store=run_plan_store)
+
+        assert extraction_kwargs['run_id'] == 'a-run'
+        assert extraction_kwargs['run_plan_store'] is run_plan_store
+        assert 'run_id' not in build_kwargs
+        assert 'run_plan_store' not in build_kwargs
+
+    def test_extract_and_build_without_restart_parameters_passes_none(self):
+        extraction_kwargs, _ = self._extract_and_build()
+
+        assert extraction_kwargs['run_id'] is None
+        assert extraction_kwargs['run_plan_store'] is None