From 5d54743d8ba5f377f4260d9573d532ecfcfaae93 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Harald=20Sch=C3=A4fer?= Date: Fri, 20 Feb 2026 15:19:39 -0800 Subject: [PATCH] safer model pkl chunking (#37283) * safer chunking * rm unchunked --- common/file_chunker.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/common/file_chunker.py b/common/file_chunker.py index f03d04a382..139a7dcacc 100644 --- a/common/file_chunker.py +++ b/common/file_chunker.py @@ -1,8 +1,9 @@ +import glob import math import os from pathlib import Path -CHUNK_SIZE = 49 * 1024 * 1024 # 49MB, under GitHub's 50MB limit +CHUNK_SIZE = 19 * 1024 * 1024 # 49MB, under GitHub's 50MB limit def get_chunk_name(name, idx, num_chunks): return f"{name}.chunk{idx+1:02d}of{num_chunks:02d}" @@ -12,6 +13,8 @@ def get_chunk_paths(path, file_size): return [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)] def chunk_file(path, num_chunks): + for old in glob.glob(f"{path}.chunk*"): + os.remove(old) with open(path, 'rb') as f: data = f.read() actual_num_chunks = max(1, math.ceil(len(data) / CHUNK_SIZE)) @@ -19,13 +22,15 @@ def chunk_file(path, num_chunks): for i in range(num_chunks): with open(get_chunk_name(path, i, num_chunks), 'wb') as f: f.write(data[i * CHUNK_SIZE:(i + 1) * CHUNK_SIZE]) + os.remove(path) def read_file_chunked(path): - for num_chunks in range(1, 100): - if os.path.isfile(get_chunk_name(path, 0, num_chunks)): - files = [get_chunk_name(path, i, num_chunks) for i in range(num_chunks)] - return b''.join(Path(f).read_bytes() for f in files) + chunks = sorted(glob.glob(f"{path}.chunk*")) + if chunks: + expected = [get_chunk_name(path, i, len(chunks)) for i in range(len(chunks))] + assert chunks == expected, f"Chunk mismatch: {chunks} != {expected}" + return b''.join(Path(f).read_bytes() for f in chunks) if os.path.isfile(path): return Path(path).read_bytes() raise FileNotFoundError(path)