Merge branch 'main' into dion/enzyme-demo

Author: dionhaefner

benchmarks/test_array_encoding.py

@@ -31,13 +31,24 @@ class ArrayModel(BaseModel):
     data: Array[(None,), Float64]
 
 
-ENCODINGS = ["json", "base64", "binref"]
+ENCODINGS = ["json", "base64", "binref", "base64+lz4", "binref+lz4"]
 
 # Maps short encoding name to the format string used by output_to_bytes
 _ENCODING_TO_FORMAT: dict[str, supported_format_type] = {
     "json": "json",
     "base64": "json+base64",
     "binref": "json+binref",
+    "base64+lz4": "json+base64",
+    "binref+lz4": "json+binref",
+}
+
+# Maps short encoding name to extra kwargs passed to output_to_bytes
+_ENCODING_TO_KWARGS: dict[str, dict] = {
+    "json": {},
+    "base64": {},
+    "binref": {},
+    "base64+lz4": {"compression": "lz4"},
+    "binref+lz4": {"compression": "lz4"},
 }
 
 
@@ -78,44 +89,50 @@ def test_encoding(benchmark, encoding_and_size):
     encoding, size = encoding_and_size
     model = ArrayModel(data=create_test_array(size))
     fmt = _ENCODING_TO_FORMAT[encoding]
+    extra_kwargs = _ENCODING_TO_KWARGS[encoding]
+    uses_binref = "binref" in encoding
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        if encoding == "binref":
+        if uses_binref:
 
             def setup():
                 _clear_dir(tmpdir)
 
             benchmark.pedantic(
                 output_to_bytes,
                 args=(model, fmt),
-                kwargs={"base_dir": tmpdir},
+                kwargs={"base_dir": tmpdir, **extra_kwargs},
                 setup=setup,
                 rounds=_binref_rounds(size),
             )
         else:
-            benchmark(output_to_bytes, model, fmt)
+            benchmark(output_to_bytes, model, fmt, **extra_kwargs)
 
 
 def test_decoding(benchmark, encoding_and_size):
     encoding, size = encoding_and_size
     model = ArrayModel(data=create_test_array(size))
     fmt = _ENCODING_TO_FORMAT[encoding]
+    extra_kwargs = _ENCODING_TO_KWARGS[encoding]
+    uses_binref = "binref" in encoding
 
     with tempfile.TemporaryDirectory() as tmpdir:
         ctx: dict[str, str] = {}
-        if encoding == "binref":
+        if uses_binref:
             ctx["base_dir"] = tmpdir
 
-        encoded = output_to_bytes(model, fmt, base_dir=tmpdir)
+        encoded = output_to_bytes(model, fmt, base_dir=tmpdir, **extra_kwargs)
 
-        if encoding == "binref":
+        if uses_binref:
             # binref filenames are random UUIDs, so we must re-encode in setup
             # and pass the fresh payload to the decode call via a mutable wrapper.
             payload = [encoded]
 
             def setup():
                 _clear_dir(tmpdir)
-                payload[0] = output_to_bytes(model, fmt, base_dir=tmpdir)
+                payload[0] = output_to_bytes(
+                    model, fmt, base_dir=tmpdir, **extra_kwargs
+                )
 
             def decode():
                 ArrayModel.model_validate_json(payload[0], context=ctx)
@@ -129,17 +146,19 @@ def test_roundtrip(benchmark, encoding_and_size):
     encoding, size = encoding_and_size
     model = ArrayModel(data=create_test_array(size))
     fmt = _ENCODING_TO_FORMAT[encoding]
+    extra_kwargs = _ENCODING_TO_KWARGS[encoding]
+    uses_binref = "binref" in encoding
 
     with tempfile.TemporaryDirectory() as tmpdir:
         ctx: dict[str, str] = {}
-        if encoding == "binref":
+        if uses_binref:
             ctx["base_dir"] = tmpdir
 
         def roundtrip():
-            enc = output_to_bytes(model, fmt, base_dir=tmpdir)
+            enc = output_to_bytes(model, fmt, base_dir=tmpdir, **extra_kwargs)
             ArrayModel.model_validate_json(enc, context=ctx)
 
-        if encoding == "binref":
+        if uses_binref:
 
             def setup():
                 _clear_dir(tmpdir)

docs/content/using-tesseracts/array-encodings.md

@@ -140,3 +140,19 @@ $ curl \
 The `.bin` file references are relative to the `--output-path`.
 :::
 ::::
+
+### lz4 compression
+
+Set the `TESSERACT_COMPRESSION=lz4` config variable to compress arrays in the output. This applies to both `json+binref` and `json+base64` formats. For binref, each array is compressed individually, preserving offset-based random access, with the compressed size embedded in the buffer path (`<file>:<offset>:<compressed_size>`).
+
+`TESSERACT_COMPRESSION` only affects how a Tesseract encodes its _output_; inputs are always decoded according to their own per-array `compression` field, so no configuration is needed to read compressed data.
+
+Because it is a runtime config variable, it must be forwarded into the container with `-e` when using `tesseract run` (a plain shell prefix would stay on the host):
+
+```bash
+$ tesseract run -e TESSERACT_COMPRESSION=lz4 vectoradd apply -f "json+binref" -o /tmp/output @examples/vectoradd/example_inputs.json
+$ cat /tmp/output/results.json
+{"result":{"object_type":"array","shape":[3],"dtype":"float64","data":{"buffer":"....bin:0:35","encoding":"binref","compression":"lz4"}}}
+```
+
+When calling `tesseract-runtime` directly (e.g. inside a container), the plain `TESSERACT_COMPRESSION=lz4 tesseract-runtime ...` prefix works as usual.