Support len for SPDL dataloader (facebookresearch#839)

Summary:

Dataloader length may be useful, in particular to estimate the number of steps in an epoch. We add support for dataloader length for SPDL HiveDataLoaders that use the SPDL datasource.

Reviewed By: ynonaolga

Differential Revision: D78133462

Author: Victor Bourgin

src/spdl/source/utils.py

@@ -13,7 +13,7 @@
 import random
 import sys
 import time
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Iterable, Iterator, Sequence, Sized
 from typing import TypeVar
 
 from ._type import IterableWithShuffle
@@ -220,6 +220,14 @@ def __iter__(self) -> Iterator[T]:
         if self._shuffle_last:
             self._shuffle()
 
+    def __len__(self) -> int:
+        if isinstance(self.src, Sized):
+            return len(self.src)
+        else:
+            raise TypeError(
+                f"Source iterator of type {type(self.src)} does not support length"
+            )
+
 
 def embed_shuffle(
     src: IterableWithShuffle[T], /, *, shuffle_last: bool = False, epoch: int = 0