"""I/O utilities for ML components.
Provides functions for file operations, directory scanning, and data loading.
"""
import gzip
import hashlib
import io
import os
import pickle
from collections.abc import Callable
from functools import wraps
from os import PathLike
from pathlib import Path
from typing import Any, TextIO
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import zstandard as zstd
from atomworks.io.utils.io_utils import apply_sharding_pattern, build_sharding_pattern
from atomworks.ml.utils.misc import logger
[docs]
def open_file(filename: PathLike) -> TextIO:
"""Open a file, handling compressed files if necessary.
Args:
filename: The path to the file to open.
Returns:
A file-like object for reading.
Raises:
AssertionError: If the file does not exist.
"""
filename = Path(filename)
# ...assert that the file exists
assert filename.exists(), f"File {filename} does not exist"
# ...open the file for reading, accepting gzipped, zstd, or plaintext files
if filename.suffix == ".gz":
return gzip.open(filename, "rt")
elif filename.suffix == ".zst":
# Open zstd file and wrap in TextIOWrapper for text mode
# Note: The file handle is managed by the TextIOWrapper/stream_reader
dctx = zstd.ZstdDecompressor()
fh = open(filename, "rb") # noqa: SIM115
reader = dctx.stream_reader(fh)
return io.TextIOWrapper(reader, encoding="utf-8")
return filename.open("r")
[docs]
def scan_directory(dir_path: PathLike, max_depth: int) -> list[str]:
"""Fast, order-independent directory scan for files up to max_depth levels deep.
Args:
dir_path: The root directory to scan.
max_depth: The maximum depth to scan. A max_depth of 1 means only the top-level directory.
Returns:
A list of file paths found within the specified directory and depth.
"""
file_paths = []
for root, dirs, files in os.walk(dir_path):
current_depth = len(Path(root).relative_to(dir_path).parts)
if current_depth >= max_depth:
dirs.clear()
continue
for file in files:
file_path = os.path.join(root, file)
file_paths.append(file_path)
return file_paths
[docs]
def cache_based_on_subset_of_args(cache_keys: list[str], maxsize: int | None = None) -> Callable:
"""Decorator to cache function results based on a subset of its keyword arguments.
Most helpful when some arguments may be unhashable types (e.g., dictionaries, AtomArray).
If the value of any of the cache keys is None, the function is executed and the result is not cached.
Note:
The wrapped function must use keyword arguments for those specified in cache_keys.
Positional arguments are not supported for cache key extraction.
Args:
cache_keys: The names of the keyword arguments to use as the cache key.
maxsize: The maximum number of entries to store in the cache.
If None, the cache size is unlimited.
Returns:
A decorator that caches the function results based on the specified keyword arguments.
Example:
.. code-block:: python
@cache_based_on_subset_of_args(["arg1"], maxsize=2)
def function(*, arg1, arg2):
return arg1 + arg2
result1 = function(arg1=1, arg2=2) # Caches with key 1
result2 = function(arg1=1, arg2=3) # Retrieves from cache
"""
def decorator(func: Callable) -> Callable:
cache = {}
cache_order: list[tuple[Any, ...]] = [] # To track the order of keys for eviction
@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
# Extract the cache key values from kwargs
key_values = tuple(kwargs.get(key) for key in cache_keys)
# Check if any of the key values are None
if None in key_values:
# Bypass caching if any key value is None
return func(*args, **kwargs)
# Use the key values to form a unique cache key
cache_key = tuple(key_values)
if cache_key not in cache:
# Evict the oldest entry if the cache is full
if maxsize is not None and len(cache) >= maxsize:
oldest_key = cache_order.pop(0)
del cache[oldest_key]
# Cache the result
cache[cache_key] = func(*args, **kwargs)
cache_order.append(cache_key)
return cache[cache_key]
return wrapper
return decorator
[docs]
def cache_to_disk_as_pickle(
cache_dir: PathLike | None = None, use_gzip: bool = True, directory_depth: int = 2
) -> Callable:
"""
A decorator to cache the results of a function to disk as a pickle file.
Creates a unique cached pickle file for each set of function arguments using an MD5 hash.
If the cache file exists, the result is loaded from the file. Otherwise, the
function is called, and the result is saved to the cache file.
If `cache_dir` is `None`, caching is disabled and the function is always executed.
Args:
cache_dir (PathLike or None): The directory where cache files will be stored, or
`None` to disable caching.
use_gzip (bool): Whether to use gzip compression for the cache files.
directory_depth (int): The depth of the directory structure for sharding cache files.
Returns:
function: The wrapped function with optional disk caching enabled.
"""
def decorator(func: Callable) -> Callable:
@wraps(func)
def wrapper(*args: Any, **kwargs: Any) -> Any:
if cache_dir is None:
# If caching is disabled, always execute the function
return func(*args, **kwargs)
# ... create cache directory if it doesn't exist
cache_dir_path = Path(cache_dir)
cache_dir_path.mkdir(parents=True, exist_ok=True)
# ... create a unique cache file path based on the MD5 hash of function arguments
args_repr = f"{args}_{kwargs}"
hash_hex = hashlib.md5(args_repr.encode()).hexdigest()
file_extension = ".pkl.gz" if use_gzip else ".pkl"
sharding_pattern = build_sharding_pattern(depth=directory_depth, chars_per_dir=2)
sharded_path = apply_sharding_pattern(hash_hex, sharding_pattern)
cache_file = Path(cache_dir) / sharded_path.with_suffix(file_extension)
# ... check if cache file exists
open_func = gzip.open if use_gzip else open
if cache_file.exists():
try:
# ... try to load the result from cache file
with open_func(cache_file, "rb") as f:
result = pickle.load(f)
return result
except Exception as e:
# (Fallback to executing the function, with a warning)
logger.error(f"Error loading cache file {cache_file}: {e}")
# If cache file doesn't exist, execute the function
result = func(*args, **kwargs)
# ... save the result to cache file, creating directories if necessary
cache_file.parent.mkdir(parents=True, exist_ok=True)
with open_func(cache_file, "wb") as f:
pickle.dump(result, f)
return result
return wrapper
return decorator
