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,7 +75,7 @@ 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.
 
 Incomplete archived logs return `complete=False` and indicate whether a later
 retry can help. Missing summaries, allocations, terminal transitions, or node

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,36 @@ 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_sum_successful_task_duration_and_require_all_segments(self):
         with tempfile.TemporaryDirectory() as directory:
@@ -1170,7 +1202,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 +1216,18 @@ def test_portable_eventlog_reader_handles_literal_lz4_block(self):
             payload, EVENTLOG.lz4_decompress_block(compressed, len(payload))
         )
 
+    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_handles_tar_bundle(self):
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
@@ -1195,11 +1243,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_yarn_job_cost_api.py

@@ -10,6 +10,8 @@
 import unittest
 from pathlib import Path
 
+import zstandard
+
 from yarn_job_cost_api import (
     EmrApplicationUsageRequest,
     calculate_emr_application_usage,
@@ -170,7 +172,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 +192,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(),
             }

yarn-resource-cost/yarn_job_cost_discovery.py

@@ -250,20 +250,22 @@ def integer_property(properties: dict, name: str, default: int | None) -> int |
 def read_event_log_metadata(path: Path) -> dict[str, EventLogApplication]:
     applications: dict[str, EventLogApplication] = {}
     by_directory: dict[str, EventLogApplication] = {}
-    for app_dir, member_name, stream, compressed in iter_eventlog_streams(path):
+    for app_dir, member_name, stream, codec in iter_eventlog_streams(path):
         current = by_directory.setdefault(
             app_dir, EventLogApplication(application_id="")
         )
         segment_match = EVENT_SEGMENT_RE.search(member_name)
         if segment_match:
             current.event_segments.add(int(segment_match.group("segment")))
-        for line in iter_text_lines(stream, compressed):
+        for line in iter_text_lines(stream, codec):
             if not line:
                 continue
             try:
                 event = json.loads(line)
             except json.JSONDecodeError:
                 continue
+            if not isinstance(event, dict):
+                continue
             event_name = event.get("Event")
             if event_name == "SparkListenerLogStart":
                 current.spark_version = str(event.get("Spark Version") or "")

yarn-resource-cost/yarn_job_cost_eventlog.py

@@ -8,12 +8,18 @@
 
 import tarfile
 from pathlib import Path
-from typing import BinaryIO, Iterable
+from typing import BinaryIO, Iterable, Literal
+
+import zstandard
 
 
 LZ4_BLOCK_MAGIC = b"LZ4Block"
 RAW_BLOCK = 0x10
 LZ4_COMPRESSED_BLOCK = 0x20
+LZ4_CODEC = "lz4"
+ZSTD_CODEC = "zstd"
+EventLogCodec = Literal["lz4", "zstd"]
+READ_SIZE = 1024 * 1024
 
 
 def lz4_decompress_block(src: bytes, expected_len: int) -> bytes:
@@ -82,9 +88,26 @@ def iter_lz4block_chunks(stream: BinaryIO) -> Iterable[bytes]:
             raise ValueError(f"Unsupported LZ4Block token {token:#x}")
 
 
-def iter_text_lines(stream: BinaryIO, compressed: bool) -> Iterable[str]:
+def iter_zstd_chunks(stream: BinaryIO) -> Iterable[bytes]:
+    decompressor = zstandard.ZstdDecompressor()
+    with decompressor.stream_reader(
+        stream, read_across_frames=True, closefd=False
+    ) as reader:
+        yield from iter(lambda: reader.read(READ_SIZE), b"")
+
+
+def iter_text_lines(
+    stream: BinaryIO, codec: EventLogCodec | None
+) -> Iterable[str]:
     pending = ""
-    chunks = iter_lz4block_chunks(stream) if compressed else iter(lambda: stream.read(1024 * 1024), b"")
+    if codec == LZ4_CODEC:
+        chunks = iter_lz4block_chunks(stream)
+    elif codec == ZSTD_CODEC:
+        chunks = iter_zstd_chunks(stream)
+    elif codec is None:
+        chunks = iter(lambda: stream.read(READ_SIZE), b"")
+    else:
+        raise ValueError(f"Unsupported event-log compression codec: {codec}")
     for chunk in chunks:
         text = pending + chunk.decode("utf-8", errors="replace")
         lines = text.splitlines(keepends=True)
@@ -104,7 +127,18 @@ def normalized_member_name(name: str) -> str:
     return name
 
 
-def iter_eventlog_streams(path: Path) -> Iterable[tuple[str, str, BinaryIO, bool]]:
+def eventlog_codec(name: str) -> EventLogCodec | None:
+    suffix = Path(name).suffix.lower()
+    if suffix == ".lz4":
+        return LZ4_CODEC
+    if suffix == ".zstd":
+        return ZSTD_CODEC
+    return None
+
+
+def iter_eventlog_streams(
+    path: Path,
+) -> Iterable[tuple[str, str, BinaryIO, EventLogCodec | None]]:
     if path.is_file() and tarfile.is_tarfile(path):
         with tarfile.open(path, mode="r:*") as tar:
             members = sorted(
@@ -120,7 +154,7 @@ def iter_eventlog_streams(path: Path) -> Iterable[tuple[str, str, BinaryIO, bool
                     continue
                 app_dir = name.split("/", 1)[0]
                 with extracted:
-                    yield app_dir, name, extracted, name.endswith(".lz4")
+                    yield app_dir, name, extracted, eventlog_codec(name)
         return
 
     if path.is_dir():
@@ -131,4 +165,4 @@ def iter_eventlog_streams(path: Path) -> Iterable[tuple[str, str, BinaryIO, bool
         rel = file_path.as_posix()
         app_dir = file_path.parent.name if file_path.parent.name.startswith("eventlog_") else file_path.stem
         with file_path.open("rb") as handle:
-            yield app_dir, rel, handle, file_path.suffix == ".lz4"
+            yield app_dir, rel, handle, eventlog_codec(file_path.name)