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 tensors and
get_tensor / get_sharded returned views onto uninitialized device memory.
Reported by @takeshi-yoshimura in #81 review.

  * BaseSafeTensorsFileLoader.get_keys() and get_shape() consult
    _tensor_filter at read time (set_tensor_filter can be called before
    or after add_filenames).
  * FilesBufferOnDevice takes an optional keep_tensor kwarg that excludes
    filtered keys from key_to_rank_lidx. copy_files_to_device() threads
    self._tensor_filter through. get_tensor / get_sharded then hit the
    existing _get_rank_lidx ValueError; fastsafe_open.keys() inherits the
    same narrowing.
  * FilesBufferOnDevice.get_filename now raises ValueError via
    _get_rank_lidx for unreachable tensors (previously returned ""). The
    behavior change also affects unknown names; the tgis_weight example
    is updated to probe key_to_rank_lidx directly.
  * ParallelLoader.iterate_weights() iterates fb.key_to_rank_lidx.keys(),
    so it skips filtered tensors automatically.

Adds two tests in test_fastsafetensors.py covering the API-surface
contract (so they run under the existing test-torch CI job).

Signed-off-by: git bisector <gitbisector@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

Author: gitbisector

examples/tgis_weight.py

@@ -151,13 +151,13 @@ def close(self):
         torch.cuda.empty_cache()
 
     def _get_alias(self, tensor_name: str) -> str:
-        if self._fb.get_filename(tensor_name) is None:
-            if tensor_name in self.aliases:
-                for alias in self.aliases[tensor_name]:
-                    if self._fb.get_filename(alias) is not None:
-                        return alias
-            raise RuntimeError(f"weight {tensor_name} does not exist")
-        return tensor_name
+        if tensor_name in self._fb.key_to_rank_lidx:
+            return tensor_name
+        if tensor_name in self.aliases:
+            for alias in self.aliases[tensor_name]:
+                if alias in self._fb.key_to_rank_lidx:
+                    return alias
+        raise RuntimeError(f"weight {tensor_name} does not exist")
 
     def get_shape(self, tensor_name: str) -> torch.Size:
         return torch.Size(self._fb.get_shape(self._get_alias(tensor_name)))

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
@@ -28,6 +28,11 @@ class FilesBufferOnDevice:
         rank_loaders (Dict<rank, list(LazyTensorFacotry)>): Tensor factories per rank, which hold device pointers for buffers.
         pg (ProcessGroupBase): process group for calling distributed ops.
         auto_mem_delete (bool): automatically release device buffers when all the tensors are shuffled.
+        keep_tensor (Callable[[str], bool], optional): If set, only tensors for
+            which ``keep_tensor(name)`` is True are registered in ``key_to_rank_lidx``;
+            others raise ``ValueError`` from ``get_tensor`` / ``get_filename`` /
+            ``get_shape``. Subclasses that reimplement the registration loop must
+            honor this.
 
     Examples:
         See examples/run_single.py and examples/run_parallel.py.
@@ -39,6 +44,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 +54,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"
@@ -69,9 +77,7 @@ def close(self):
         self.rank_loaders = {}
 
     def get_filename(self, tensor_name: str) -> str:
-        if tensor_name not in self.key_to_rank_lidx:
-            return ""
-        rank, lidx = self.key_to_rank_lidx[tensor_name]
+        rank, lidx = self._get_rank_lidx(tensor_name)
         return self.rank_loaders[rank][lidx].metadata.src
 
     def get_shape(self, tensor_name: str) -> List[int]:

fastsafetensors/loader.py

@@ -99,23 +99,27 @@ 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]:
+        if self._tensor_filter is not None and not self._tensor_filter(tensor_name):
+            raise ValueError(f"get_shape: key {tensor_name} is filtered out")
         return self.frames[tensor_name].shape
 
     def set_tensor_filter(self, keep_tensor: Optional[Callable[[str], bool]]) -> None:
         """Load only the tensors for which ``keep_tensor(name)`` is True.
 
-        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.
-        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.
+        The ``nogds`` and ``unified`` copiers skip reading bytes for filtered
+        tensors; other copiers load the full file. The filter narrows the
+        public API on every backend: ``get_keys()`` omits filtered tensors,
+        ``FilesBufferOnDevice`` does not register them, and ``get_tensor``,
+        ``get_filename``, and ``get_shape`` raise ``ValueError`` for them.
+        ``ParallelLoader.iterate_weights()`` skips them. ``None`` (the
+        default) loads every tensor. See
+        ``fastsafetensors.ep_slice.expert_parallel_filter``.
         """
         self._tensor_filter = keep_tensor
 
@@ -194,7 +198,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_fastsafetensors.py

@@ -7,9 +7,16 @@
 
 import pytest
 
-from fastsafetensors import SafeTensorsFileLoader, SafeTensorsMetadata, SingleGroup
+from fastsafetensors import (
+    ParallelLoader,
+    SafeTensorsFileLoader,
+    SafeTensorsMetadata,
+    SingleGroup,
+)
 from fastsafetensors import cpp as fstcpp
-from fastsafetensors import fastsafe_open
+from fastsafetensors import (
+    fastsafe_open,
+)
 from fastsafetensors.common import get_device_numa_node, is_gpu_found
 from fastsafetensors.copier.gds import GdsFileCopier
 from fastsafetensors.copier.nogds import NoGdsFileCopier
@@ -531,7 +538,8 @@ def test_SafeTensorsFileLoader(fstcpp_log, input_files, framework) -> None:
         assert bufs.get_filename(last_key) == input_files[0]
         assert bufs.get_shape(last_key) == last_shape
         assert loader.get_shape(last_key) == last_shape
-    assert bufs.get_filename("aaaaaaaaaaaaa") == ""
+    with pytest.raises(ValueError):
+        bufs.get_filename("aaaaaaaaaaaaa")
     bufs.close()
     loader.close()
     assert framework.get_mem_used() == 0
@@ -560,6 +568,62 @@ def test_SafeTensorsFileLoaderNoGds(fstcpp_log, input_files, framework) -> None:
     assert fstcpp.get_cpp_metrics().bounce_buffer_bytes == 0
 
 
+def test_tensor_filter_hides_skipped_tensors(fstcpp_log, input_files, framework):
+    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)
+    with pytest.raises(ValueError):
+        bufs.get_filename(skipped)
+    with pytest.raises(ValueError):
+        loader.get_shape(skipped)
+
+    bufs.close()
+    loader.close()
+
+
+def test_tensor_filter_iterate_weights_hides_skipped(
+    fstcpp_log, input_files, framework
+):
+    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
+
+    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,
+    )
+    yielded = {key for key, _t in loader.iterate_weights()}
+    assert yielded == kept
+
+    loader.close()
+    assert framework.get_mem_used() == 0
+
+
 def test_fastsafe_open(fstcpp_log, input_files, framework) -> None:
     device, _ = get_and_check_device(framework)