refactor: sharding operations

refactor: sharding operations

Author: nscorley

.ipd/.env

@@ -25,6 +25,8 @@ CCD_MIRROR_PATH=/projects/ml/frozen_pdb_copies/2025_07_13_ccd
 # --- Local MSA directories ---
 LOCAL_MSA_DIRS=/projects/msa/hhblits,/projects/msa/mmseqs_gpu,/projects/msa/lab
 
+# --- Cache directories ---
+RESIDUE_CACHE_DIR=/net/tukwila/lschaaf/datahub/MACE-Egret-3-noH/mace_embeddings
 
 # --- External tools ---
 
@@ -54,5 +56,3 @@ COLABFOLD_LOCAL_DB_PATH_CPU=/local/databases/colabfold/
 # Network access (fallback; may cause IO-related issues)
 COLABFOLD_NET_DB_PATH_GPU=/net/databases/colabfold/gpu
 COLABFOLD_NET_DB_PATH_CPU=/net/databases/colabfold/
-
-

src/atomworks/io/parser.py

@@ -42,7 +42,13 @@
 from atomworks.io.utils.bonds import get_struct_conn_dict_from_atom_array
 from atomworks.io.utils.ccd import check_ccd_codes_are_available
 from atomworks.io.utils.chain import create_chain_id_generator
-from atomworks.io.utils.io_utils import get_structure, infer_pdb_file_type, read_any
+from atomworks.io.utils.io_utils import (
+    apply_sharding_pattern,
+    build_sharding_pattern,
+    get_structure,
+    infer_pdb_file_type,
+    read_any,
+)
 from atomworks.io.utils.non_rcsb import (
     get_identity_assembly_gen_category,
     get_identity_op_expr_category,
@@ -53,20 +59,31 @@
 
 __all__ = ["parse"]
 
-DEFAULT_PARSE_KWARGS = {
+STANDARD_PARSER_ARGS = {
     "add_missing_atoms": True,
     "add_id_and_entity_annotations": True,
     "add_bond_types_from_struct_conn": ["covale"],
     "remove_ccds": CRYSTALLIZATION_AIDS,
     "remove_waters": True,
     "fix_ligands_at_symmetry_centers": True,
-    "hydrogen_policy": "keep",
     "fix_arginines": True,
     "fix_formal_charges": True,
-    "convert_mse_to_met": True,
-    "build_assembly": "all",
+    "fix_bond_types": True,
+    "convert_mse_to_met": True,  # Changed from False to True vs. atomworks.io.parser.parse default
+    "hydrogen_policy": "keep",
+    "model": None,  # all models
 }
-"""Some fairly standard parsing arguments that can be imported for convenience."""
+"""Common cif parser arguments for `atomworks.io.parse` for many biomolecular use cases.
+
+Similar to the defaults below but additionally converts selenomethionine (MSE) residues to methionine (MET) residues,
+which is desirable for many practical applications but would not be appropriate as a universal default.
+
+This dictionary exists to provide a convenient import for the standard parameters.
+"""
+
+# Cache sharding configuration (internal, not exposed to parse() to avoid complexity)
+_CACHE_SHARDING_DEPTH = 2  # Use 2-level sharding by default (e.g., ab/cd/abcdef123456/)
+_CACHE_SHARDING_CHARS_PER_DIR = 2  # Number of characters per directory level
 
 
 def _get_atomworks_version() -> str:
@@ -79,6 +96,32 @@ def _get_atomworks_version() -> str:
         return "unknown"
 
 
+def _parse_args_to_hash(parse_arguments: dict[str, Any], truncate: int = 8) -> str:
+    """Compute hash from parse arguments with sorted keys."""
+    args_string = ",".join(str(parse_arguments[k]) for k in sorted(parse_arguments.keys()))
+    return string_to_md5_hash(args_string, truncate=truncate)
+
+
+def _build_cache_file_path(
+    cache_dir: Path,
+    args_hash: str,
+    filename: os.PathLike,
+    assembly_info: str,
+) -> Path:
+    """Build sharded cache file path for parsed structure."""
+    structure_id = Path(filename).stem
+
+    # Pad structure ID to minimum required length for sharding
+    min_length = _CACHE_SHARDING_DEPTH * _CACHE_SHARDING_CHARS_PER_DIR
+    structure_id_padded = structure_id.ljust(min_length, "_")
+
+    # Build sharded path
+    sharding_pattern = build_sharding_pattern(depth=_CACHE_SHARDING_DEPTH, chars_per_dir=_CACHE_SHARDING_CHARS_PER_DIR)
+    sharded_path = apply_sharding_pattern(structure_id_padded, sharding_pattern)
+
+    return cache_dir / args_hash / sharded_path / f"{structure_id}_assembly_{assembly_info}.pkl.gz"
+
+
 def parse(
     filename: os.PathLike | io.StringIO | io.BytesIO,
     *,
@@ -182,12 +225,11 @@ def parse(
                 Should typically not be used directly.
 
     """
-    # CCD mirror
     if ccd_mirror_path and not os.path.exists(ccd_mirror_path):
-        logger.warning(
-            f"Local mirror of the Chemical Component Dictionary does not exist: {ccd_mirror_path}. Falling back to Biotite's built-in CCD."
+        raise FileNotFoundError(
+            f"Local mirror of the Chemical Component Dictionary does not exist: {ccd_mirror_path}. "
+            "To use Biotite's built-in CCD, set `ccd_mirror_path` to None."
         )
-        ccd_mirror_path = None
 
     # Set default value for remove_ccds if None
     if remove_ccds is None:
@@ -232,15 +274,13 @@ def parse(
             "convert_mse_to_met": convert_mse_to_met,
             "hydrogen_policy": hydrogen_policy,
         }
-        # Compose args_string from parse_arguments values (in order)
-        args_string = ",".join(str(parse_arguments[k]) for k in parse_arguments)
-        args_hash = string_to_md5_hash(args_string, truncate=8)
+        args_hash = _parse_args_to_hash(parse_arguments)
 
         # ... generate assembly info
         assembly_info = ",".join(build_assembly) if isinstance(build_assembly, list | tuple) else build_assembly
 
-        # ... construct the full cache file path
-        cache_file_path = cache_dir / args_hash / f"{Path(filename).stem}_assembly_{assembly_info}.pkl.gz"
+        # ... construct the full cache file path with sharding
+        cache_file_path = _build_cache_file_path(cache_dir, args_hash, filename, assembly_info)
 
         # If we are loading from cache, try to load the result from the cache
         if load_from_cache:
@@ -268,8 +308,7 @@ def parse(
                     result["assemblies"] = assemblies
                     return result
             except Exception as e:
-                # Log an error, and continue to parse from CIF
-                logger.error(f"Error loading from cache: {e}")
+                raise RuntimeError(f"Error loading from cache: {e}, tried path: {cache_file_path}") from e
 
     if file_type == "pdb":
         result = _parse_from_pdb(
@@ -315,9 +354,9 @@ def parse(
 
     if not is_buffer and save_to_cache and cache_dir and (not cache_file_path.exists()):
         # We want our cache to include:
-        #   (1) All keys in `result` excep the assemblies and
-        #   (2) The information needed to rebuild the assembly(s), which is stored in `result["extra_info"]`
-        #   (3) The parse_arguments and atomworks.io version
+        #   (1) All keys in `result` except the assemblies; and,
+        #   (2) The information needed to rebuild the assembly(s), which is stored in `result["extra_info"]`; and,
+        #   (3) The parse_arguments and atomworks version
 
         # Add parse_arguments and version to metadata before saving
         result.setdefault("metadata", {}).update(

src/atomworks/io/tools/inference.py

@@ -25,7 +25,7 @@
 )
 from atomworks.enums import ChainType, ChainTypeInfo
 from atomworks.io import parse
-from atomworks.io.parser import DEFAULT_PARSE_KWARGS
+from atomworks.io.parser import STANDARD_PARSER_ARGS
 from atomworks.io.template import build_template_atom_array
 from atomworks.io.tools.fasta import one_letter_to_ccd_code, split_generalized_fasta_sequence
 from atomworks.io.utils.bonds import (
@@ -257,7 +257,7 @@ def _parse_standard_pdb_or_cif(self) -> None:
             self.custom_parse_kwargs = {}
 
         # We add missing atoms later to the fully-concatenated inference AtomArray
-        parse_kwargs = {**DEFAULT_PARSE_KWARGS, "add_missing_atoms": False} | self.custom_parse_kwargs
+        parse_kwargs = {**STANDARD_PARSER_ARGS, "add_missing_atoms": False} | self.custom_parse_kwargs
 
         if parse_kwargs["add_missing_atoms"]:
             logger.warning(

src/atomworks/io/utils/io_utils.py

@@ -1,8 +1,11 @@
 """General utility functions for working with CIF files in Biotite."""
 
 __all__ = [
+    "apply_sharding_pattern",
+    "build_sharding_pattern",
     "get_structure",
     "load_any",
+    "parse_sharding_pattern",
     "read_any",
     "suppress_logging_messages",
     "to_cif_buffer",
@@ -14,6 +17,7 @@
 import io
 import logging
 import os
+import re
 import warnings
 from collections.abc import Generator
 from contextlib import contextmanager
@@ -934,3 +938,100 @@ def find_files_by_extension(input_dir: Path, extension: str) -> list[Path]:
         raise FileNotFoundError(f"No files with extension {extension} found in {input_dir}")
 
     return files
+
+
+def build_sharding_pattern(depth: int, chars_per_dir: int = 2) -> str:
+    """Build a sharding pattern string from depth and characters per directory.
+
+    Args:
+        depth: Number of directory levels.
+        chars_per_dir: Number of characters to use for each directory level.
+
+    Returns:
+        Sharding pattern string.
+
+    Examples:
+        >>> build_sharding_pattern(2, 2)
+        '/0:2/2:4/'
+        >>> build_sharding_pattern(3, 1)
+        '/0:1/1:2/2:3/'
+    """
+    if depth == 0:
+        return ""
+
+    parts = []
+    for i in range(depth):
+        start = i * chars_per_dir
+        end = start + chars_per_dir
+        parts.append(f"/{start}:{end}")
+
+    return "".join(parts) + "/"
+
+
+def parse_sharding_pattern(sharding_pattern: str) -> list[tuple[int, int]]:
+    """Parse a sharding pattern string into directory levels.
+
+    Args:
+        sharding_pattern: String like ``"/1:2/0:2/"`` where each ``/start:end/`` defines a directory level.
+            ``start:end`` defines the character range to use for that directory level.
+
+    Returns:
+        List of (start, end) tuples for each directory level.
+
+    Examples:
+        >>> parse_sharding_pattern("/1:2/0:2/")
+        [(1, 2), (0, 2)]
+    """
+    # Find all patterns like /start:end/ using a non-consuming lookahead
+    pattern = r"/(\d+):(\d+)(?=/)"
+    matches = []
+    for match in re.finditer(pattern, sharding_pattern):
+        matches.append((int(match.group(1)), int(match.group(2))))
+
+    if not matches:
+        raise ValueError(f"Invalid sharding pattern format: {sharding_pattern}. Expected format like '/1:2/0:2/'")
+
+    return matches
+
+
+def apply_sharding_pattern(path: os.PathLike, sharding_pattern: str | None = None) -> Path:
+    """Apply a sharding pattern to construct a file path.
+
+    Args:
+        path: The base path or identifier (e.g., PDB ID).
+        sharding_pattern: Pattern for organizing files in subdirectories. Examples:
+            - ``"/0:2/"``: Use first two characters for first directory level
+            - ``"/0:2/2:4/"``: Use chars 0-2 for first dir, then chars 2-4 for second dir
+            - ``None``: No sharding (default)
+
+    Returns:
+        The constructed file path with sharding applied.
+
+    Examples:
+        >>> apply_sharding_pattern("12as", "/0:2/1:3/")
+        Path("12/2a/12as")
+    """
+    path_str = str(path)
+    assert path_str and path_str != ".", "Path cannot be empty"
+
+    if not sharding_pattern:
+        return Path(path_str)
+
+    if not sharding_pattern.startswith("/"):
+        raise ValueError(f"Sharding pattern must start with '/': {sharding_pattern}")
+
+    try:
+        shard_ranges = parse_sharding_pattern(sharding_pattern)
+    except ValueError as e:
+        raise ValueError(f"Invalid sharding pattern '{sharding_pattern}': {e}") from e
+
+    # Validate all ranges before building path
+    for start, end in shard_ranges:
+        if end > len(path_str):
+            raise ValueError(f"Sharding range {start}:{end} exceeds path length {len(path_str)} for '{path_str}'")
+
+    # Build directory components from sharding ranges
+    directory_parts = [path_str[start:end] for start, end in shard_ranges]
+
+    # Construct final path: directories + filename
+    return Path(*directory_parts, path_str)

src/atomworks/ml/datasets/loaders.py

@@ -11,8 +11,8 @@
 import pandas as pd
 from toolz import keyfilter
 
-from atomworks.io.parser import parse
-from atomworks.ml.utils.io import apply_sharding_pattern
+from atomworks.io.parser import STANDARD_PARSER_ARGS, parse
+from atomworks.io.utils.io_utils import apply_sharding_pattern
 
 
 def _construct_metadata_hierarchy(row: pd.Series, attrs: dict | None = None) -> dict[str, Any]:
@@ -70,11 +70,13 @@ def _construct_structure_path(
 
 
 def _load_structure_from_path(path: Path, assembly_id: str, parser_args: dict | None = None) -> dict[str, Any]:
-    """Load structure from file path using the CIF parser."""
+    """Load structure from file path using the CIF parser, merging with STANDARD_PARSER_ARGS."""
+    # Merge STANDARD_PARSER_ARGS with parser_args (parser_args takes precedence)
+    merged_args = {**STANDARD_PARSER_ARGS, **(parser_args or {})}
     result_dict = parse(
         filename=path,
         build_assembly=(assembly_id,),
-        **(parser_args or {}),
+        **merged_args,
     )
     return result_dict
 

src/atomworks/ml/datasets/parsers/base.py

@@ -4,26 +4,8 @@
 
 import pandas as pd
 
-from atomworks.constants import CRYSTALLIZATION_AIDS
 from atomworks.io import parse
-
-DEFAULT_PARSER_ARGS = {
-    "add_missing_atoms": True,
-    "add_id_and_entity_annotations": True,
-    "add_bond_types_from_struct_conn": ["covale"],
-    "remove_ccds": CRYSTALLIZATION_AIDS,
-    "remove_waters": True,
-    "fix_ligands_at_symmetry_centers": True,
-    "fix_arginines": True,
-    "fix_formal_charges": True,
-    "fix_bond_types": True,
-    "convert_mse_to_met": True,
-    "hydrogen_policy": "remove",
-    "model": None,  # all models
-}
-"""Default cif parser arguments for `atomworks.io.parse`.
-This dictionary exists to provide a convenient import for the default parameters.
-"""
+from atomworks.io.parser import STANDARD_PARSER_ARGS
 
 
 class MetadataRowParser(ABC):
@@ -123,7 +105,7 @@ def load_example_from_metadata_row(
         cif_parser_args.setdefault("save_to_cache", True)
 
     # Merge DEFAULT_CIF_PARSER_ARGS with cif_parser_args, overriding with the keys present in cif_parser_args
-    merged_cif_parser_args = {**DEFAULT_PARSER_ARGS, **cif_parser_args}
+    merged_cif_parser_args = {**STANDARD_PARSER_ARGS, **cif_parser_args}
 
     # Use the parse function with the merged CIF parser arguments
     result_dict = parse(

src/atomworks/ml/preprocessing/msa/organizing.py

@@ -16,9 +16,9 @@
 
 from atomworks.enums import MSAFileExtension
 from atomworks.io.utils.compression import transfer_with_compression
-from atomworks.io.utils.io_utils import find_files_by_extension
+from atomworks.io.utils.io_utils import apply_sharding_pattern, find_files_by_extension
 from atomworks.ml.preprocessing.msa.finding import sequence_has_msa
-from atomworks.ml.utils.io import apply_sharding_pattern, open_file
+from atomworks.ml.utils.io import open_file
 from atomworks.ml.utils.misc import hash_sequence
 
 logger = logging.getLogger(__name__)