Truncate compression dict training if it gets too big.

Author: TallJimbo

doc/changes/DM-55232.bugfix.md

@@ -0,0 +1 @@
+Fix zstd failure in `butler aggregate-graph` when logs are really huge.

python/lsst/pipe/base/quantum_graph/aggregator/_config.py

@@ -136,6 +136,11 @@ class AggregatorConfig(pydantic.BaseModel):
     logs, and metadata blocks encountered.
     """
 
+    zstd_dict_input_max_bytes: int = 1_048_576
+    """Maximum total size of the ZStandard compression dictionary training
+    sample.
+    """
+
     mock_storage_classes: bool = False
     """Enable support for storage classes by created by the
     lsst.pipe.base.tests.mocks package.

python/lsst/pipe/base/quantum_graph/aggregator/_writer.py

@@ -172,15 +172,34 @@ def make_compression_dictionary(self) -> zstandard.ZstdCompressionDict:
         # chop out the datastore records since those don't appear in the
         # provenance graph.
         n_pred_quanta: int = 0
+        training_sample_size: int = 0
         for predicted_quantum in self.predicted.quantum_datasets.values():
             if len(training_inputs) == self.comms.config.zstd_dict_n_inputs:
                 break
             predicted_quantum.datastore_records.clear()
-            training_inputs.append(predicted_quantum.model_dump_json().encode())
+            block = predicted_quantum.model_dump_json().encode()
+            training_sample_size += len(block)
+            if training_sample_size >= self.comms.config.zstd_dict_input_max_bytes:
+                self.comms.log.warning(
+                    "Reached compression dict training sample size limit at %d predicted quanta.",
+                    len(training_inputs),
+                )
+                break
+            training_inputs.append(block)
+            training_sample_size += len(block)
             n_pred_quanta += 1
         # Add the provenance quanta, metadata, and logs we've accumulated.
         for write_request in self.pending_compression_training:
             assert not write_request.is_compressed, "We can't compress without the compression dictionary."
+            training_sample_size += len(write_request.quantum)
+            training_sample_size += len(write_request.metadata)
+            training_sample_size += len(write_request.logs)
+            if training_sample_size >= self.comms.config.zstd_dict_input_max_bytes:
+                self.comms.log.warning(
+                    "Reached compression dict training sample size limit at %d blocks.",
+                    len(training_inputs),
+                )
+                break
             training_inputs.append(write_request.quantum)
             training_inputs.append(write_request.metadata)
             training_inputs.append(write_request.logs)