Apply tensor filter to public API surface, not just I/O

set_tensor_filter() previously only narrowed which bytes the nogds/unified
copiers read; the public API still advertised filtered-out tensors and
get_tensor / get_sharded returned views onto uninitialized device memory
rather than raising. Reported by @takeshi-yoshimura in foundation-model-stack#81 review.

  * SafeTensorsFileLoader.get_keys() honors _tensor_filter at read time
    so set_tensor_filter() may still be called before or after
    add_filenames().
  * FilesBufferOnDevice gains an optional keep_tensor kwarg and excludes
    filtered keys from key_to_rank_lidx at construction time.
    SafeTensorsFileLoader.copy_files_to_device() threads its own
    _tensor_filter through. get_tensor / get_sharded therefore hit the
    existing _get_rank_lidx ValueError instead of returning uninitialized
    data; fastsafe_open.keys() inherits this because it iterates
    key_to_rank_lidx.
  * ParallelLoader.iterate_weights() already iterates
    list(fb.key_to_rank_lidx.keys()), so it skips filtered-out tensors
    with no further change.

Adds two tests to tests/unit/test_ep_slice.py covering the new contract:
the exact test from the PR review (get_keys / key_to_rank_lidx /
get_tensor) plus a ParallelLoader.iterate_weights variant.

Author: gitbisector

fastsafetensors/file_buffer.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections import OrderedDict
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from .common import init_logger
 from .frameworks import FrameworkOpBase, ProcessGroupBase, TensorBase
@@ -39,6 +39,7 @@ def __init__(
         pg: ProcessGroupBase,
         framework: FrameworkOpBase,
         auto_mem_delete: bool = True,
+        keep_tensor: Optional[Callable[[str], bool]] = None,
     ):
         self.framework = framework
         self.rank_loaders: Dict[int, List[LazyTensorFactory]] = rank_loaders
@@ -48,6 +49,8 @@ def __init__(
             self.instantiated[rank] = {}
             for lidx, loader in enumerate(loaders):
                 for key in loader.metadata.tensors.keys():
+                    if keep_tensor is not None and not keep_tensor(key):
+                        continue
                     if key in self.key_to_rank_lidx:
                         raise Exception(
                             f"FilesBufferOnDevice: key {key} must be unique among files"

fastsafetensors/loader.py

@@ -99,7 +99,10 @@ def close(self):
         del self.copier_constructor
 
     def get_keys(self) -> List[str]:
-        return list(self.frames.keys())
+        if self._tensor_filter is None:
+            return list(self.frames.keys())
+        keep = self._tensor_filter
+        return [k for k in self.frames.keys() if keep(k)]
 
     def get_shape(self, tensor_name: str) -> List[int]:
         return self.frames[tensor_name].shape
@@ -109,13 +112,20 @@ def set_tensor_filter(self, keep_tensor: Optional[Callable[[str], bool]]) -> Non
 
         When set, each owned file is read partially: bytes belonging to
         filtered-out tensors are skipped and their device-buffer regions are
-        left uninitialized, so those tensors must not be requested afterwards.
+        left uninitialized. Filtered-out tensors are also hidden from the
+        public API: ``get_keys()`` omits them, ``FilesBufferOnDevice`` does
+        not register them in ``key_to_rank_lidx``, and any attempt to fetch
+        one via ``get_tensor`` / ``get_sharded`` raises ``ValueError`` rather
+        than returning uninitialized data. ``ParallelLoader.iterate_weights``
+        likewise skips them.
+
         Useful e.g. under expert parallelism to read only this rank's owned
         experts -- see ``fastsafetensors.ep_slice.expert_parallel_filter``.
         ``None`` (the default) loads every tensor.
 
         Partial reads are implemented by the ``nogds`` and ``unified`` copiers;
-        other copiers ignore the filter and load the full file.
+        other copiers ignore the filter and load the full file, but the API
+        surface is still narrowed for consistency regardless of backend.
         """
         self._tensor_filter = keep_tensor
 
@@ -194,7 +204,12 @@ def copy_files_to_device(
             lidx += 1
         for factory in need_wait:
             factory.wait_io(dtype=dtype, noalign=False)
-        return FilesBufferOnDevice(factories, pg=self.pg, framework=self.framework)
+        return FilesBufferOnDevice(
+            factories,
+            pg=self.pg,
+            framework=self.framework,
+            keep_tensor=self._tensor_filter,
+        )
 
 
 class SafeTensorsFileLoader(BaseSafeTensorsFileLoader):

tests/unit/test_ep_slice.py

@@ -9,7 +9,12 @@
 import pytest
 import torch
 
-from fastsafetensors import SafeTensorsMetadata
+from fastsafetensors import (
+    ParallelLoader,
+    SafeTensorsFileLoader,
+    SafeTensorsMetadata,
+    SingleGroup,
+)
 from fastsafetensors import cpp as fstcpp
 from fastsafetensors.copier.nogds import NoGdsFileCopier
 from fastsafetensors.copier.unified import new_unified_copier
@@ -196,3 +201,66 @@ def test_unified_full_read_unchanged(fstcpp_log, input_files, framework):
     framework.free_tensor_memory(gbuf, device)
     del copier
     assert framework.get_mem_used() == 0
+
+
+# ---- filter must narrow the public API, not just I/O ----
+
+
+def test_tensor_filter_hides_skipped_tensors(fstcpp_log, input_files, framework):
+    """A skipped tensor must not appear in get_keys / key_to_rank_lidx, and
+    asking for it via get_tensor must raise ValueError rather than returning
+    uninitialized device memory."""
+    device, _ = get_and_check_device(framework)
+    meta = SafeTensorsMetadata.from_file(input_files[0], framework)
+
+    kept = set(sorted(meta.tensors.keys())[::2])
+    keep = lambda name: name in kept  # noqa: E731
+    skipped = next(name for name in meta.tensors if name not in kept)
+
+    loader = SafeTensorsFileLoader(
+        pg=SingleGroup(),
+        device=device.as_str(),
+        framework=framework.get_name(),
+        nogds=True,
+    )
+    loader.set_tensor_filter(keep)
+    loader.add_filenames({0: [input_files[0]]})
+    bufs = loader.copy_files_to_device()
+
+    assert set(loader.get_keys()) == kept
+    assert skipped not in bufs.key_to_rank_lidx
+    with pytest.raises(ValueError):
+        bufs.get_tensor(skipped)
+
+    bufs.close()
+    loader.close()
+
+
+def test_tensor_filter_iterate_weights_hides_skipped(
+    fstcpp_log, input_files, framework
+):
+    """ParallelLoader.iterate_weights must yield only filter-kept tensors;
+    skipped ones (which would return uninitialized data) must never appear."""
+    device, _ = get_and_check_device(framework)
+    meta = SafeTensorsMetadata.from_file(input_files[0], framework)
+
+    kept = set(sorted(meta.tensors.keys())[::2])
+    keep = lambda name: name in kept  # noqa: E731
+    assert kept and any(
+        n not in kept for n in meta.tensors
+    ), "fixture should split into kept and skipped"
+
+    loader = ParallelLoader(
+        pg=SingleGroup(),
+        hf_weights_files=[input_files[0]],
+        device=device.as_str(),
+        nogds=True,
+        framework=framework.get_name(),
+        tensor_filter=keep,
+        all_local=True,  # required when tensor_filter drops tensors
+    )
+    yielded = {key for key, _tensor in loader.iterate_weights()}
+    assert yielded == kept
+
+    loader.close()
+    assert framework.get_mem_used() == 0