Support Zstandard-compressed Spark event logs

.github/workflows/yarn-resource-cost-python.yml

@@ -27,11 +27,12 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Compile
         run: python -m py_compile yarn-resource-cost/*.py
+      - name: Install
+        run: python -m pip install ./yarn-resource-cost
       - name: Unit tests
         run: python -m unittest discover -s yarn-resource-cost -p 'test*.py'
       - name: Installed console error handling
         run: |
-          python -m pip install ./yarn-resource-cost
           python - <<'PY'
           import shutil
           import subprocess

yarn-resource-cost/CONTRIBUTING.md

@@ -10,6 +10,7 @@ explicitly sanitized fixtures.
 Before submitting a change, run:
 
 ```bash
+python3 -m pip install ./yarn-resource-cost
 python3 -m py_compile yarn-resource-cost/*.py
 python3 -m unittest discover -s yarn-resource-cost -p 'test*.py'
 python3 yarn-resource-cost/package_yarn_job_cost.py

yarn-resource-cost/README.md

@@ -16,7 +16,8 @@ they are not treated as the allocation ledger.
 
 ## Requirements
 
-- Python 3.10 or newer; the accounting code has no third-party Python packages.
+- Python 3.10 or newer. Installing the package also installs `zstandard` for
+  compressed Spark event logs.
 - Archived Spark rolling event logs.
 - ResourceManager logs containing node registrations, allocations, terminal
   transitions, application summaries, and ResourceCalculator evidence.
@@ -74,13 +75,17 @@ print(usage.instance_seconds_by_type)
 `event_log_uri` accepts an S3 URI, a plain local path, or a local `file://` URI.
 Passing a pre-materialized local file or rolling-event-log directory avoids an
 S3 download; the selected event log is still streamed to extract accounting
-metadata.
+metadata. Plain, `.lz4`, and `.zstd` event-log segments are supported. Individual
+event records are limited to 64 MiB of decompressed data.
 
 Incomplete archived logs return `complete=False` and indicate whether a later
 retry can help. Missing summaries, allocations, terminal transitions, or node
 registration metadata are retryable; contradictory or unsupported accounting
-policies are not. Authentication and transport errors propagate from boto3.
-The caller decides whether and how to translate instance-seconds into currency.
+policies are not. Unreadable event-log segments are isolated to their owning
+application: corrupt compressed data is retryable, while a missing decoder or
+an oversized event record is not. Authentication and transport errors propagate
+from boto3. The caller decides whether and how to translate instance-seconds
+into currency.
 
 Memory, vcores, `yarn.io/gpu`, and arbitrary numeric custom resources are
 parsed generically. Heterogeneous node classes remain separate in structured

yarn-resource-cost/package_yarn_job_cost.py

@@ -112,6 +112,7 @@ def validate_bundle(root: Path) -> None:
         raise ValueError("Bundle CONTRIBUTING.md does not contain DCO 1.1")
     run([sys.executable, "-m", "unittest", "discover", "-s", ".", "-p", "test*.py"], cwd=root)
     run([sys.executable, "yarn_resource_cost.py", "--help"], cwd=root)
+    run([sys.executable, "-S", "yarn_resource_cost.py", "--help"], cwd=root)
 
 
 def remove_bytecode(root: Path) -> None:

yarn-resource-cost/pyproject.toml

@@ -9,6 +9,7 @@ description = "Portable YARN resource accounting for Spark applications"
 readme = "README.md"
 requires-python = ">=3.10,<3.13"
 license = "Apache-2.0"
+dependencies = ["zstandard>=0.23.0"]
 classifiers = [
   "Programming Language :: Python :: 3",
   "Programming Language :: Python :: 3.10",

yarn-resource-cost/test_calculate_yarn_job_cost.py

@@ -19,6 +19,8 @@
 from unittest import mock
 from pathlib import Path
 
+import zstandard
+
 
 SCRIPT = Path(__file__).with_name("calculate_yarn_job_cost.py")
 DISCOVERY = Path(__file__).with_name("yarn_job_cost_discovery.py")
@@ -726,6 +728,72 @@ def test_event_logs_derive_identity_cluster_versions_and_duration(self):
                 metadata["application_123_0002"].as_metadata()["spark_duration_seconds"],
             )
 
+    def test_event_logs_read_zstd_and_ignore_non_object_json(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            app_dir = root / f"eventlog_v2_{APP_ID}"
+            app_dir.mkdir()
+            records = [
+                1,
+                "not an event",
+                [],
+                None,
+                {"Event": "SparkListenerLogStart", "Spark Version": "4.0.2"},
+                {
+                    "Event": "SparkListenerApplicationStart",
+                    "App ID": APP_ID,
+                    "App Name": "/run/j0144__job.py",
+                    "Timestamp": 1000,
+                },
+                {"Event": "SparkListenerApplicationEnd", "Timestamp": 6000},
+            ]
+            payload = (
+                "\n".join(json.dumps(record) for record in records) + "\n"
+            ).encode()
+            compressed = zstandard.ZstdCompressor().compress(payload)
+            (app_dir / f"events_1_{APP_ID}.zstd").write_bytes(compressed)
+
+            metadata = MODULE.read_event_log_metadata(root)[APP_ID]
+
+            self.assertEqual("4.0.2", metadata.spark_version)
+            self.assertEqual("144", metadata.job_id)
+            self.assertEqual(5.0, metadata.as_metadata()["spark_duration_seconds"])
+
+    def test_event_logs_isolate_corrupt_application(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            healthy_app_id = "application_123_0002"
+            healthy_dir = root / f"eventlog_v2_{healthy_app_id}"
+            healthy_dir.mkdir()
+            healthy_events = [
+                {
+                    "Event": "SparkListenerApplicationStart",
+                    "App ID": healthy_app_id,
+                    "App Name": "/run/j0144__job.py",
+                    "Timestamp": 1000,
+                },
+                {"Event": "SparkListenerApplicationEnd", "Timestamp": 6000},
+            ]
+            (healthy_dir / f"events_1_{healthy_app_id}").write_text(
+                "\n".join(json.dumps(event) for event in healthy_events) + "\n"
+            )
+
+            corrupt_app_id = "application_123_0003"
+            corrupt_dir = root / f"eventlog_v2_{corrupt_app_id}"
+            corrupt_dir.mkdir()
+            (corrupt_dir / f"events_1_{corrupt_app_id}.zstd").write_bytes(
+                b"not a zstandard frame"
+            )
+
+            metadata = MODULE.read_event_log_metadata(root)
+
+            self.assertEqual("144", metadata[healthy_app_id].job_id)
+            corrupt = metadata[corrupt_app_id]
+            self.assertTrue(corrupt.event_log_read_retryable)
+            self.assertIn("Unable to decompress", corrupt.event_log_read_errors[0])
+            self.assertIn(
+                corrupt.event_log_read_errors[0], corrupt.task_metric_warnings
+            )
 
     def test_event_logs_sum_successful_task_duration_and_require_all_segments(self):
         with tempfile.TemporaryDirectory() as directory:
@@ -1170,7 +1238,11 @@ def test_portable_eventlog_reader_handles_raw_lz4block(self):
         )
         self.assertEqual(
             ["first", "second"],
-            list(EVENTLOG.iter_text_lines(io.BytesIO(header + payload), True)),
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(header + payload), EVENTLOG.LZ4_CODEC
+                )
+            ),
         )
 
     def test_portable_eventlog_reader_handles_literal_lz4_block(self):
@@ -1180,6 +1252,64 @@ def test_portable_eventlog_reader_handles_literal_lz4_block(self):
             payload, EVENTLOG.lz4_decompress_block(compressed, len(payload))
         )
 
+    def test_portable_eventlog_reader_classifies_truncated_lz4_as_retryable(self):
+        with self.assertRaisesRegex(ValueError, "Truncated LZ4Block") as raised:
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(EVENTLOG.LZ4_BLOCK_MAGIC), EVENTLOG.LZ4_CODEC
+                )
+            )
+        self.assertTrue(raised.exception.retryable)
+
+    def test_portable_eventlog_reader_handles_zstd(self):
+        payload = b"first\nsecond\n"
+        compressed = zstandard.ZstdCompressor().compress(payload)
+        self.assertEqual(
+            ["first", "second"],
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(compressed), EVENTLOG.ZSTD_CODEC
+                )
+            ),
+        )
+
+    def test_portable_eventlog_reader_preserves_split_utf8(self):
+        payload = '{"name":"caf\u00e9"}\n'.encode()
+        split = payload.index("\u00e9".encode()) + 1
+        with mock.patch.object(
+            EVENTLOG,
+            "iter_zstd_chunks",
+            return_value=iter((payload[:split], payload[split:])),
+        ):
+            lines = list(
+                EVENTLOG.iter_text_lines(io.BytesIO(), EVENTLOG.ZSTD_CODEC)
+            )
+        self.assertEqual(['{"name":"caf\u00e9"}'], lines)
+
+    def test_portable_eventlog_reader_rejects_oversized_line(self):
+        with self.assertRaisesRegex(ValueError, "8-byte limit") as raised:
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(b"123456789\n"), None, max_line_bytes=8
+                )
+            )
+        self.assertFalse(raised.exception.retryable)
+
+    def test_portable_eventlog_reader_reports_missing_zstandard(self):
+        with mock.patch.dict(sys.modules, {"zstandard": None}):
+            with self.assertRaisesRegex(ValueError, "pip install") as raised:
+                list(EVENTLOG.iter_zstd_chunks(io.BytesIO(b"not-zstd")))
+        self.assertFalse(raised.exception.retryable)
+
+    def test_portable_eventlog_reader_classifies_corrupt_zstandard_as_retryable(self):
+        with self.assertRaisesRegex(ValueError, "Unable to decompress") as raised:
+            list(
+                EVENTLOG.iter_text_lines(
+                    io.BytesIO(b"not a zstandard frame"), EVENTLOG.ZSTD_CODEC
+                )
+            )
+        self.assertTrue(raised.exception.retryable)
+
     def test_portable_eventlog_reader_handles_tar_bundle(self):
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
@@ -1195,11 +1325,38 @@ def test_portable_eventlog_reader_handles_tar_bundle(self):
                     ),
                 )
             streams = []
-            for app_dir, _, stream, compressed in EVENTLOG.iter_eventlog_streams(
+            for app_dir, _, stream, codec in EVENTLOG.iter_eventlog_streams(
+                archive
+            ):
+                streams.append(
+                    (app_dir, list(EVENTLOG.iter_text_lines(stream, codec)))
+                )
+            self.assertEqual(
+                [("eventlog_v2_application_123_0002", ["one", "two"])], streams
+            )
+
+    def test_portable_eventlog_reader_handles_zstd_tar_member(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            source = root / "events_1_application_123_0002.zstd"
+            source.write_bytes(
+                zstandard.ZstdCompressor().compress(b"one\ntwo\n")
+            )
+            archive = root / "events.tar.gz"
+            with tarfile.open(archive, "w:gz") as bundle:
+                bundle.add(
+                    source,
+                    arcname=(
+                        "eventlog_v2_application_123_0002/"
+                        "events_1_application_123_0002.zstd"
+                    ),
+                )
+            streams = []
+            for app_dir, _, stream, codec in EVENTLOG.iter_eventlog_streams(
                 archive
             ):
                 streams.append(
-                    (app_dir, list(EVENTLOG.iter_text_lines(stream, compressed)))
+                    (app_dir, list(EVENTLOG.iter_text_lines(stream, codec)))
                 )
             self.assertEqual(
                 [("eventlog_v2_application_123_0002", ["one", "two"])], streams

yarn-resource-cost/test_portable_yarn_resource_cost.py

@@ -132,6 +132,7 @@ def test_on_prem_fixture_end_to_end_with_catalog(self):
             completed = subprocess.run(
                 [
                     sys.executable,
+                    "-S",
                     str(SCRIPT),
                     "--adapter",
                     "on-prem",

yarn-resource-cost/test_yarn_job_cost_api.py

@@ -9,6 +9,10 @@
 import io
 import unittest
 from pathlib import Path
+from types import SimpleNamespace
+from unittest import mock
+
+import zstandard
 
 from yarn_job_cost_api import (
     EmrApplicationUsageRequest,
@@ -170,7 +174,7 @@ def test_missing_node_instance_type_is_retryable(self):
         self.assertTrue(result.retryable)
         self.assertIn("unknown instance type", " | ".join(result.warnings))
 
-    def test_s3_rolling_event_log_segments_stay_grouped(self):
+    def test_s3_zstd_rolling_event_log_segments_stay_grouped(self):
         event_prefix = "spark-events/eventlog_v2_application_1_0001"
         segment_1 = "\n".join(
             (
@@ -190,10 +194,15 @@ def test_s3_rolling_event_log_segments_stay_grouped(self):
                 '{"Event":"SparkListenerApplicationEnd","Timestamp":11000}',
             )
         )
+        compressor = zstandard.ZstdCompressor()
         s3 = FakeS3Client(
             {
-                f"{event_prefix}/events_1_application_1_0001": segment_1.encode(),
-                f"{event_prefix}/events_2_application_1_0001": segment_2.encode(),
+                f"{event_prefix}/events_1_application_1_0001.zstd": (
+                    compressor.compress(segment_1.encode())
+                ),
+                f"{event_prefix}/events_2_application_1_0001.zstd": (
+                    compressor.compress(segment_2.encode())
+                ),
                 "emr-logs/j-TEST/node/i-1/applications/"
                 "hadoop-yarn-resourcemanager-rm.log": self.yarn_log_with_instance_type().encode(),
             }
@@ -211,6 +220,46 @@ def test_s3_rolling_event_log_segments_stay_grouped(self):
         self.assertEqual(24576.0, result.memory_mb_seconds)
         self.assertEqual({"m5.xlarge": 3.0}, result.instance_seconds_by_type)
 
+    def test_corrupt_zstd_event_log_is_retryable(self):
+        event_prefix = "spark-events/eventlog_v2_application_1_0001"
+        s3 = FakeS3Client(
+            {
+                f"{event_prefix}/events_1_application_1_0001.zstd": (
+                    b"not a zstandard frame"
+                )
+            }
+        )
+
+        result = calculate_emr_application_usage(
+            self.request(f"s3://test-bucket/{event_prefix}"),
+            emr_client=FakeEmrClient(),
+            s3_client=s3,
+        )
+
+        self.assertFalse(result.complete)
+        self.assertTrue(result.retryable)
+        self.assertIn("Unable to decompress", result.warnings[0])
+
+    def test_permanent_event_log_read_error_is_not_retryable(self):
+        event_log_dir = FIXTURE / "eventlog_v2_application_1_0001"
+        metadata = SimpleNamespace(
+            event_log_read_errors=["Event-log line exceeds the configured limit"],
+            event_log_read_retryable=False,
+        )
+        with mock.patch(
+            "yarn_job_cost_api.read_event_log_metadata",
+            return_value={"application_1_0001": metadata},
+        ):
+            result = calculate_emr_application_usage(
+                self.request(str(event_log_dir)),
+                emr_client=FakeEmrClient(),
+                s3_client=FakeS3Client({}),
+            )
+
+        self.assertFalse(result.complete)
+        self.assertFalse(result.retryable)
+        self.assertIn("configured limit", result.warnings[0])
+
     def test_s3_single_file_event_log_is_materialized_as_a_file(self):
         event_key = "spark-events/application_1_0001"
         event_log = (

yarn-resource-cost/yarn_job_cost_api.py

@@ -226,6 +226,12 @@ def calculate_emr_application_usage(
                 f"Application {request.application_id} is missing from the Spark event log",
                 retryable=True,
             )
+        if metadata.event_log_read_errors:
+            return _empty_result(
+                request,
+                " | ".join(metadata.event_log_read_errors),
+                retryable=metadata.event_log_read_retryable,
+            )
 
         yarn_root = root / "yarn"
         downloaded = _materialize_yarn_logs(emr_client, s3_client, request.cluster_id, yarn_root)