Merge pull request #610 from chrisburr/feat-zstd-sandbox-compression

feat: implement zstandard compression for sandbox files

Author: chrisburr

diracx-api/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
     "diracx-client",
     "diracx-core",
     "httpx",
+    "zstandard",
 ]
 dynamic = ["version"]
 

diracx-api/src/diracx/api/jobs.py

@@ -7,10 +7,12 @@
 import os
 import tarfile
 import tempfile
+from contextlib import contextmanager
 from pathlib import Path
-from typing import Literal
+from typing import BinaryIO, Literal
 
 import httpx
+import zstandard
 
 from diracx.client.aio import AsyncDiracClient
 from diracx.client.models import SandboxInfo
@@ -20,8 +22,29 @@
 logger = logging.getLogger(__name__)
 
 SANDBOX_CHECKSUM_ALGORITHM = "sha256"
-SANDBOX_COMPRESSION: Literal["bz2"] = "bz2"
-SANDBOX_OPEN_MODE: Literal["w|bz2"] = "w|bz2"
+SANDBOX_COMPRESSION: Literal["zst"] = "zst"
+
+
+@contextmanager
+def tarfile_open(fileobj: BinaryIO):
+    """Context manager to extend tarfile.open to support reading zstd compressed files.
+
+    This is only needed for Python <=3.13.
+    """
+    # Save current position and read magic bytes
+    current_pos = fileobj.tell()
+    magic = fileobj.read(4)
+    fileobj.seek(current_pos)
+
+    # Read magic bytes to determine compression format
+    if magic.startswith(b"\x28\xb5\x2f\xfd"):  # zstd magic number
+        dctx = zstandard.ZstdDecompressor()
+        with dctx.stream_reader(fileobj) as decompressor:
+            with tarfile.open(fileobj=decompressor, mode="r|") as tf:
+                yield tf
+    else:
+        with tarfile.open(fileobj=fileobj, mode="r") as tf:
+            yield tf
 
 
 @with_client
@@ -33,10 +56,18 @@ async def create_sandbox(paths: list[Path], *, client: AsyncDiracClient) -> str:
     be used to submit jobs.
     """
     with tempfile.TemporaryFile(mode="w+b") as tar_fh:
-        with tarfile.open(fileobj=tar_fh, mode=SANDBOX_OPEN_MODE) as tf:
-            for path in paths:
-                logger.debug("Adding %s to sandbox as %s", path.resolve(), path.name)
-                tf.add(path.resolve(), path.name, recursive=True)
+        # Create zstd compressed tar with level 18 and long matching enabled
+        compression_params = zstandard.ZstdCompressionParameters.from_level(
+            18, enable_ldm=1
+        )
+        cctx = zstandard.ZstdCompressor(compression_params=compression_params)
+        with cctx.stream_writer(tar_fh, closefd=False) as compressor:
+            with tarfile.open(fileobj=compressor, mode="w|") as tf:
+                for path in paths:
+                    logger.debug(
+                        "Adding %s to sandbox as %s", path.resolve(), path.name
+                    )
+                    tf.add(path.resolve(), path.name, recursive=True)
         tar_fh.seek(0)
 
         hasher = getattr(hashlib, SANDBOX_CHECKSUM_ALGORITHM)()
@@ -89,6 +120,6 @@ async def download_sandbox(pfn: str, destination: Path, *, client: AsyncDiracCli
         fh.seek(0)
         logger.debug("Sandbox downloaded for %s", pfn)
 
-        with tarfile.open(fileobj=fh) as tf:
+        with tarfile_open(fh) as tf:
             tf.extractall(path=destination, filter="data")
         logger.debug("Extracted %s to %s", pfn, destination)

diracx-client/src/diracx/client/_generated/models/_enums.py

@@ -38,6 +38,7 @@ class SandboxFormat(str, Enum, metaclass=CaseInsensitiveEnumMeta):
     """SandboxFormat."""
 
     TAR_BZ2 = "tar.bz2"
+    TAR_ZST = "tar.zst"
 
 
 class SandboxType(str, Enum, metaclass=CaseInsensitiveEnumMeta):

diracx-client/src/diracx/client/_generated/models/_models.py

@@ -930,7 +930,7 @@ class SandboxInfo(_serialization.Model):
     :vartype checksum: str
     :ivar size: Size. Required.
     :vartype size: int
-    :ivar format: SandboxFormat. Required. "tar.bz2"
+    :ivar format: SandboxFormat. Required. Known values are: "tar.bz2" and "tar.zst".
     :vartype format: str or ~_generated.models.SandboxFormat
     """
 
@@ -964,7 +964,7 @@ def __init__(
         :paramtype checksum: str
         :keyword size: Size. Required.
         :paramtype size: int
-        :keyword format: SandboxFormat. Required. "tar.bz2"
+        :keyword format: SandboxFormat. Required. Known values are: "tar.bz2" and "tar.zst".
         :paramtype format: str or ~_generated.models.SandboxFormat
         """
         super().__init__(**kwargs)

diracx-core/src/diracx/core/models.py

@@ -204,6 +204,7 @@ class ChecksumAlgorithm(StrEnum):
 
 class SandboxFormat(StrEnum):
     TAR_BZ2 = "tar.bz2"
+    TAR_ZST = "tar.zst"
 
 
 class SandboxInfo(BaseModel):

diracx-logic/tests/logic/test_sandboxes.py

@@ -72,14 +72,14 @@ async def test_upload_and_clean(
     """
     data = secrets.token_bytes(256)
     data_digest = hashlib.sha256(data).hexdigest()
-    key = f"fakevo/fake_group/fakeuser/sha256:{data_digest}.tar.bz2"
+    key = f"fakevo/fake_group/fakeuser/sha256:{data_digest}.tar.zst"
     expected_pfn = f"SB:SandboxSE|/S3/sandboxes/{key}"
 
     sandbox_info = SandboxInfo(
         checksum_algorithm=ChecksumAlgorithm.SHA256,
         checksum=data_digest,
         size=len(data),
-        format=SandboxFormat.TAR_BZ2,
+        format=SandboxFormat.TAR_ZST,
     )
 
     # Test with a new sandbox

extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_enums.py

@@ -38,6 +38,7 @@ class SandboxFormat(str, Enum, metaclass=CaseInsensitiveEnumMeta):
     """SandboxFormat."""
 
     TAR_BZ2 = "tar.bz2"
+    TAR_ZST = "tar.zst"
 
 
 class SandboxType(str, Enum, metaclass=CaseInsensitiveEnumMeta):

extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_models.py

@@ -951,7 +951,7 @@ class SandboxInfo(_serialization.Model):
     :vartype checksum: str
     :ivar size: Size. Required.
     :vartype size: int
-    :ivar format: SandboxFormat. Required. "tar.bz2"
+    :ivar format: SandboxFormat. Required. Known values are: "tar.bz2" and "tar.zst".
     :vartype format: str or ~_generated.models.SandboxFormat
     """
 
@@ -985,7 +985,7 @@ def __init__(
         :paramtype checksum: str
         :keyword size: Size. Required.
         :paramtype size: int
-        :keyword format: SandboxFormat. Required. "tar.bz2"
+        :keyword format: SandboxFormat. Required. Known values are: "tar.bz2" and "tar.zst".
         :paramtype format: str or ~_generated.models.SandboxFormat
         """
         super().__init__(**kwargs)