Strict type-checking + tests for foundry.utils (instantiators, ddp, logging); remove dead rf3.data.paired_msa

models/rf3/configs/datasets/pdb_and_distillation.yaml

@@ -9,7 +9,6 @@ defaults:
     - monomer_distillation
     - na_complex_distillation
     - disorder_distillation
-    # - domain_distillation
     # - rna_monomer_distillation
   - val/af3_validation@val.af3_validation
   - val/af3_validation@val.quick_af3_validation_with_templating
@@ -28,8 +27,6 @@ train:
     probability: 0.02
   disorder_distillation:
     probability: 0.02
-  # multidomain_distillation:
-  #   probability: 0.06
   # rna_monomer_distillation:
   #   probability: 0.04
 

pyproject.toml

@@ -258,10 +258,7 @@ module = [
 ignore_errors = true
 
 # NOTE: the rf3 enablement ratchet (0014) is fully cleared — there is no rf3 module-level
-# mypy exemption here. The one module that cannot type-check, `rf3.data.paired_msa`
-# (broken against the installed atomworks; needs a `PandasDataset`-API refactor), carries
-# a file-level `# mypy: ignore-errors` directive in the module itself, so the suppression
-# is visible where the code is and the module stays inside mypy's `files` scope.
+# mypy exemption here, and every rf3 module type-checks with no in-file suppressions.
 
 # Per-module strictness ratchet (direction (b)). The global baseline above leaves
 # disallow_untyped_defs / check_untyped_defs off; fully-annotated modules opt into strict
@@ -274,6 +271,10 @@ module = [
     "foundry.utils.alignment",
     "foundry.utils.weights",
     "foundry.utils.rotation_augmentation",
+    "foundry.utils.instantiators",
+    "foundry.utils.ddp",
+    "foundry.utils.squashfs",
+    "foundry.utils.logging",
 ]
 disallow_untyped_defs = true
 check_untyped_defs = true

src/foundry/utils/ddp.py

@@ -79,7 +79,7 @@ def __init__(
         self.rank_zero_only = rank_zero_only
 
     def log(  # type: ignore[override]  # deliberately extends LoggerAdapter.log with a `rank` parameter
-        self, level: int, msg: str, rank: int | None = None, *args, **kwargs
+        self, level: int, msg: str, rank: int | None = None, *args: Any, **kwargs: Any
     ) -> None:
         """
         Delegate a log call to the underlying logger, after prefixing its message with the rank

src/foundry/utils/logging.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 
 import pandas as pd
-from beartype.typing import Any
+from beartype.typing import Any, Iterator
 from lightning.fabric.utilities import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 from rich.console import Console
@@ -20,14 +20,14 @@
 class CachedDataFilter(logging.Filter):
     """Filter to suppress atomworks cached data logging messages."""
 
-    def filter(self, record):
+    def filter(self, record: logging.LogRecord) -> bool:
         # Filter out "Cached data not found" messages
         if "Cached data not found" in record.getMessage():
             return False
         return True
 
 
-def silence_warnings():
+def silence_warnings() -> None:
     """Silence common warnings that appear during foundry execution."""
     warnings.filterwarnings(
         "ignore", message="All-NaN slice encountered", category=RuntimeWarning
@@ -67,7 +67,7 @@ def silence_warnings():
 
 
 @contextmanager
-def suppress_warnings(is_inference: bool = False):
+def suppress_warnings(is_inference: bool = False) -> Iterator[None]:
     """Context manager to suppress specific warnings within its scope.
 
     Args:
@@ -178,7 +178,7 @@ def print_model_parameters(model: nn.Module, name: str = "") -> None:
 
 def log_hyperparameters_with_all_loggers(
     trainer: Any, cfg: dict | DictConfig, model: Any
-):
+) -> None:
     """Logs hyperparameters using all loggers in the trainer.
 
     Args:
@@ -260,7 +260,7 @@ def table_from_df(df: pd.DataFrame, title: str) -> Table:
     return table
 
 
-def safe_print(obj: Any, console_width=100, logger: Any | None = None) -> None:
+def safe_print(obj: Any, console_width: int = 100, logger: Any | None = None) -> None:
     """Print a Rich object in a console- and logger-safe manner."""
     console = Console(force_terminal=False, color_system=None, width=console_width)