fix: stream random test data to disk to reduce memory footprint

This commit is contained in:
gmega 2024-12-20 13:57:28 -03:00
parent f3a66d9637
commit 87b9d43cd5
No known key found for this signature in database
GPG Key ID: 6290D34EAD824B18

View File

@ -47,20 +47,8 @@ class RandomTempData(ExperimentData[TInitialMetadata]):
self._context.__exit__(exc_type, exc_val, exc_tb)
@contextmanager
def temp_random_file(size: int, name: str = "data.bin"):
with tempfile.TemporaryDirectory() as temp_dir_str:
temp_dir = Path(temp_dir_str)
random_file = temp_dir / name
random_bytes = os.urandom(size)
with random_file.open("wb") as outfile:
outfile.write(random_bytes)
yield random_file
def await_predicate(
predicate: Callable[[], bool], timeout: float = 0, polling_interval: float = 0
predicate: Callable[[], bool], timeout: float = 0, polling_interval: float = 0
) -> bool:
current = time()
while (timeout == 0) or ((time() - current) <= timeout):
@ -88,3 +76,19 @@ def kilobytes(n: int) -> int:
def megabytes(n: int) -> int:
return kilobytes(n) * 1024
@contextmanager
def temp_random_file(size: int, name: str = "data.bin", batch_size: int = megabytes(50)):
with tempfile.TemporaryDirectory() as temp_dir_str:
temp_dir = Path(temp_dir_str)
random_file = temp_dir / name
with random_file.open("wb") as outfile:
while size > 0:
batch = min(size, batch_size)
random_bytes = os.urandom(batch)
outfile.write(random_bytes)
size -= batch
yield random_file