kaxil opened a new pull request, #76108: URL: https://github.com/apache/airflow/pull/76108
Offloading a value through the common.io object-storage state store held more copies of it than it needed. `serialize_task_state_store_to_ref` and `serialize_asset_state_store_to_ref` encoded the JSON string to bytes once, only to compare its length with `state_store_objectstorage_threshold`. Then `_write_to_object_storage` encoded it again and passed the whole bytes object to one `write()`. On buffered filesystems such as s3fs, that write is copied once more into the upload buffer. The threshold check now uses the string's length. `json.dumps` escapes non-ASCII characters, so the length is the byte count, and the inline-or-offload decision is the same at the same threshold. A new boundary test pins this for task and asset scope, with ASCII and non-ASCII values, below, at and above the threshold. The write now encodes and writes the string in 1 MiB slices. Stored bytes and refs are unchanged, so nothing needs migrating. Peak memory measured with tracemalloc on a 24 MB JSON payload with the threshold at 0. The "buffered" rows use a filesystem built on fsspec's own `AbstractBufferedFile` write path, with 5 MB blocks, that discards the uploaded bytes: | Offload path | Before | After | |---|---|---| | Buffered filesystem, no compression | 72.0 MB | **48.0 MB** | | Buffered filesystem, gzip | 58.6 MB | **48.0 MB** | | Local file, gzip | 58.6 MB | **48.0 MB** | | Local file, no compression | 48.0 MB | 48.0 MB | The remaining 48 MB comes from `json.dumps`, which peaks at twice its output size while it joins its pieces. The read path is unchanged. `json.load(fh)` is `json.loads(fh.read())`, and the decoded string is the input `json.loads` needs anyway. The measured read peak was the parsed object plus one copy of the payload. #### Follow-ups - Streaming the encode, with `json.dump` writing straight into the file handle, would bring the write peak down to a few MB for list and dict values. It uses the pure-Python encoder, though: it was 2–3x slower in my measurements and no better for a single large string. The threshold decision would also have to be made mid-stream. - The common.io XCom backend (`xcom/backend.py`) has the same encode-then-write-everything pattern. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
