[FEAT] add HDFS support for DynamicEmb checkpoint dump and load

corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py

@@ -38,6 +38,7 @@
     warning_for_cstm_score,
 )
 from dynamicemb.embedding_admission import MultiTableKVCounter
+from dynamicemb.filesystem import get_filesystem
 from dynamicemb.initializer import create_initializer_from_args
 from dynamicemb.key_value_table import (
     Cache,
@@ -90,6 +91,17 @@ def encode_counter_checkpoint_file_path(
     )
 
 
+def _strip_scheme_prefix(path: str) -> str:
+    """Strip ``hdfs://host:port`` prefix if present, returning ``/path``."""
+    if path.startswith("hdfs://"):
+        # hdfs://host:port/path  →  /path
+        without_scheme = path[len("hdfs://") :]
+        idx = without_scheme.find("/")
+        if idx != -1:
+            return without_scheme[idx:]
+    return path
+
+
 def find_files(root_path: str, table_name: str, suffix: str) -> Tuple[List[str], int]:
     suffix_to_encode_file_path_func = {
         "emb_keys": partial(encode_checkpoint_file_path, item="keys"),
@@ -105,10 +117,10 @@ def find_files(root_path: str, table_name: str, suffix: str) -> Tuple[List[str],
         raise RuntimeError(f"Invalid suffix: {suffix}")
     encode_file_path_func = suffix_to_encode_file_path_func[suffix]
 
-    import glob
-
     # v2 version
-    files = glob.glob(encode_file_path_func(root_path, table_name, "*", "*"))
+    files = get_filesystem(root_path).glob(
+        encode_file_path_func(root_path, table_name, "*", "*")
+    )
     if len(files) == 0:
         return [], 0
     files = sorted(files)
@@ -118,11 +130,23 @@ def find_files(root_path: str, table_name: str, suffix: str) -> Tuple[List[str],
             f"Checkpoints is corrupted. Found {len(files)} under path {root_path} for table {table_name}, but the number of checkpointed world size is {world_size}."
         )
 
+    # HDFS backends (fsspec/pyarrow) strip the scheme prefix from glob
+    # results.  Rewrite returned paths to include the original scheme so
+    # downstream code can use them directly.
+    if root_path.startswith("hdfs://"):
+        scheme = root_path[: root_path.index("/", len("hdfs://"))]  # hdfs://host:port
+        files = [f if f.startswith("hdfs://") else scheme + f for f in files]
+
+    # Validate that every expected rank file is present.
+    files_set = set(files)
     for i in range(world_size):
-        expected_file_path = encode_file_path_func(root_path, table_name, i, world_size)
-        if expected_file_path not in set(files):
+        expected = encode_file_path_func(root_path, table_name, i, world_size)
+        if (
+            expected not in files_set
+            and _strip_scheme_prefix(expected) not in files_set
+        ):
             raise RuntimeError(
-                f"Checkpoints is corrupted. Expected file path {expected_file_path} for table {table_name}, but it is not found."
+                f"Checkpoints is corrupted. Expected file path {expected} for table {table_name}, but it is not found."
             )
 
     return files, len(files)
@@ -134,7 +158,7 @@ def get_loading_files(
     rank: int,
     world_size: int,
 ) -> Tuple[List[str], List[str], List[str], List[str], int, int]:
-    if not os.path.exists(root_path):
+    if not get_filesystem(root_path).exists(root_path):
         raise RuntimeError(f"can't find path to load, path:", root_path)
 
     key_files, num_key_files = find_files(root_path, name, "emb_keys")

corelib/dynamicemb/dynamicemb/dump_load.py

@@ -22,6 +22,7 @@
 import torch
 import torch.distributed as dist
 from dynamicemb.batched_dynamicemb_tables import BatchedDynamicEmbeddingTablesV2
+from dynamicemb.filesystem import get_filesystem
 from torch import nn
 from torchrec.distributed.embedding import ShardedEmbeddingCollection
 from torchrec.distributed.embeddingbag import ShardedEmbeddingBagCollection
@@ -142,15 +143,16 @@ def DynamicEmbDump(
         torch.cuda.synchronize()
     # create path
     created_dirs = []
-    if not os.path.exists(path):
+    fs = get_filesystem(path)
+    if not fs.exists(path):
         created_dirs.append(path)
         try:
-            os.makedirs(path, exist_ok=True)
+            fs.makedirs(path, exist_ok=True)
         except Exception as e:
             raise Exception("can't build path:", path) from e
     else:
-        if os.path.isdir(path):
-            if not os.listdir(path):
+        if fs.isdir(path):
+            if not fs.ls(path):
                 pass
             elif not allow_overwrite:
                 raise Exception(
@@ -174,8 +176,9 @@ def DynamicEmbDump(
 
     for collection_path, _, _ in collections_list:
         full_collection_path = os.path.join(path, collection_path)
-        if not os.path.exists(full_collection_path):
-            os.makedirs(full_collection_path, exist_ok=True)
+        col_fs = get_filesystem(full_collection_path)
+        if not col_fs.exists(full_collection_path):
+            col_fs.makedirs(full_collection_path, exist_ok=True)
     dist.barrier(group=pg, device_ids=[torch.cuda.current_device()])
 
     for collection_path, _, current_collection_module in collections_list:
@@ -246,7 +249,7 @@ def DynamicEmbLoad(
     if torch.cuda.is_available():
         torch.cuda.synchronize()
 
-    if not os.path.exists(path):
+    if not get_filesystem(path).exists(path):
         raise Exception("can't find path to load, path:", path)
 
     collections_list: List[Tuple[str, str, nn.Module]] = find_sharded_modules(model, "")

corelib/dynamicemb/dynamicemb/exportable_tables.py

@@ -10,7 +10,6 @@
 """
 
 import itertools
-import os
 from typing import Dict, List, Optional
 
 import pynve.torch.nve_layers as nve_layers
@@ -24,6 +23,7 @@
     encode_meta_json_file_path,
     get_loading_files,
 )
+from dynamicemb.filesystem import get_filesystem
 from dynamicemb.key_value_table import _iter_batches_from_files, load_from_json
 from dynamicemb.scored_hashtable import ScorePolicy
 from dynamicemb_extensions import table_insert
@@ -380,7 +380,7 @@ def load_from_dynamicemb_file(
         save_dir: str,
         table_names: Optional[List[str]] = None,
     ) -> None:
-        if not os.path.exists(save_dir):
+        if not get_filesystem(save_dir).exists(save_dir):
             raise RuntimeError(f"Save directory does not exist: {save_dir}")
 
         if (
@@ -409,7 +409,7 @@ def load_from_dynamicemb_file(
                 continue
 
             meta_json_file = encode_meta_json_file_path(save_dir, table_name)
-            if os.path.exists(meta_json_file):
+            if get_filesystem(meta_json_file).exists(meta_json_file):
                 try:
                     _ = load_from_json(meta_json_file)
                 except Exception as e:

corelib/dynamicemb/dynamicemb/filesystem.py

@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Lightweight filesystem abstraction for DynamicEmb checkpoint I/O.
+
+Supports local POSIX filesystems natively and HDFS via the optional
+``fsspec[arrow]`` package.  The factory :func:`get_filesystem` selects
+the backend based on the path prefix (``hdfs://`` → Hdfs, otherwise
+local), so no public API signatures need to change.
+"""
+
+from __future__ import annotations
+
+import abc
+import glob as _glob
+import os
+from typing import IO, Any, List
+
+
+class AbstractFileSystem(abc.ABC):
+    """Minimal filesystem interface covering the operations used by
+    DynamicEmb dump / load.
+    """
+
+    @abc.abstractmethod
+    def open(self, path: str, mode: str = "rb") -> IO[Any]:
+        ...
+
+    @abc.abstractmethod
+    def exists(self, path: str) -> bool:
+        ...
+
+    @abc.abstractmethod
+    def isdir(self, path: str) -> bool:
+        ...
+
+    @abc.abstractmethod
+    def makedirs(self, path: str, exist_ok: bool = True) -> None:
+        ...
+
+    @abc.abstractmethod
+    def ls(self, path: str) -> List[str]:
+        ...
+
+    @abc.abstractmethod
+    def glob(self, pattern: str) -> List[str]:
+        ...
+
+    @abc.abstractmethod
+    def size(self, path: str) -> int:
+        ...
+
+
+# ---------------------------------------------------------------------------
+# Local filesystem (always available, zero extra dependencies)
+# ---------------------------------------------------------------------------
+
+
+class LocalFileSystem(AbstractFileSystem):
+    """Thin wrapper around Python's built-in file I/O and ``os.path``."""
+
+    def open(self, path: str, mode: str = "rb") -> IO[Any]:
+        return open(path, mode)
+
+    def exists(self, path: str) -> bool:
+        return os.path.exists(path)
+
+    def isdir(self, path: str) -> bool:
+        return os.path.isdir(path)
+
+    def makedirs(self, path: str, exist_ok: bool = True) -> None:
+        os.makedirs(path, exist_ok=exist_ok)
+
+    def ls(self, path: str) -> List[str]:
+        return os.listdir(path)
+
+    def glob(self, pattern: str) -> List[str]:
+        return _glob.glob(pattern)
+
+    def size(self, path: str) -> int:
+        return os.path.getsize(path)
+
+
+# ---------------------------------------------------------------------------
+# HDFS filesystem (requires ``fsspec[arrow]``)
+# ---------------------------------------------------------------------------
+
+
+class HdfsFileSystem(AbstractFileSystem):
+    """Filesystem backed by an HDFS cluster via *fsspec* / *pyarrow*.
+
+    The *fsspec* + *pyarrow* dependency is only imported when the first
+    I/O operation is performed.  Host and port are parsed from ``hdfs://``
+    URIs and passed explicitly to *pyarrow*, so the Hadoop configuration
+    files (``core-site.xml``) do **not** need to set ``fs.defaultFS``.
+    """
+
+    def __init__(self) -> None:
+        self._fs = None  # lazily initialised on first use
+        self._host: str | None = None
+        self._port: int = 0
+
+    # -- lazy init helper -----------------------------------------------
+
+    def _ensure_fs(self) -> None:
+        if self._fs is not None:
+            return
+        try:
+            import fsspec  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "HDFS support requires the 'fsspec[arrow]' package. "
+                "Install it with:  pip install fsspec[arrow]"
+            )
+        try:
+            import pyarrow.fs as _pafs
+        except ImportError:
+            raise ImportError(
+                "HDFS support requires the 'pyarrow' package (bundled with "
+                "fsspec[arrow]). Install it with:  pip install fsspec[arrow]"
+            )
+
+        # pyarrow HadoopFileSystem needs explicit host:port when the
+        # environment's core-site.xml does not set fs.defaultFS correctly.
+        # Otherwise libhdfs falls back to RawLocalFileSystem (file:///).
+        if self._host:
+            arrow_fs = _pafs.HadoopFileSystem(host=self._host, port=self._port)
+        else:
+            arrow_fs = _pafs.HadoopFileSystem()
+
+        from fsspec.implementations.arrow import ArrowFSWrapper
+
+        self._fs = ArrowFSWrapper(arrow_fs)
+
+    def _parse_uri(self, path: str) -> None:
+        """Extract host + port from an ``hdfs://host:port/...`` URI."""
+        if self._host is not None or not path.startswith("hdfs://"):
+            return
+        rest = path[len("hdfs://") :]
+        host_part = rest.split("/", 1)[0]
+        if ":" in host_part:
+            self._host, port_s = host_part.rsplit(":", 1)
+            self._port = int(port_s)
+        else:
+            self._host = host_part
+            self._port = 8020  # default HDFS NameNode port
+
+    # -- I/O methods ----------------------------------------------------
+
+    def open(self, path: str, mode: str = "rb") -> IO[Any]:
+        self._parse_uri(path)
+        self._ensure_fs()
+        return self._fs.open(path, mode)
+
+    def exists(self, path: str) -> bool:
+        self._parse_uri(path)
+        self._ensure_fs()
+        return self._fs.exists(path)
+
+    def isdir(self, path: str) -> bool:
+        self._parse_uri(path)
+        self._ensure_fs()
+        return self._fs.isdir(path)
+
+    def makedirs(self, path: str, exist_ok: bool = True) -> None:
+        self._parse_uri(path)
+        self._ensure_fs()
+        self._fs.makedirs(path, exist_ok=exist_ok)
+
+    def ls(self, path: str) -> List[str]:
+        self._parse_uri(path)
+        self._ensure_fs()
+        return self._fs.ls(path)
+
+    def glob(self, pattern: str) -> List[str]:
+        self._parse_uri(pattern)
+        self._ensure_fs()
+        return self._fs.glob(pattern)
+
+    def size(self, path: str) -> int:
+        self._parse_uri(path)
+        self._ensure_fs()
+        return self._fs.size(path)
+
+
+# ---------------------------------------------------------------------------
+# Factory
+# ---------------------------------------------------------------------------
+
+_HDFS: HdfsFileSystem | None = None
+_LOCAL: LocalFileSystem | None = None
+
+
+def get_filesystem(path: str) -> AbstractFileSystem:
+    """Return the appropriate filesystem for *path*.
+
+    Parameters
+    ----------
+    path : str
+        A filesystem path.  If it starts with ``"hdfs://"`` an
+        :class:`HdfsFileSystem` is returned; otherwise the default
+        :class:`LocalFileSystem` is returned.
+
+    Returns
+    -------
+    AbstractFileSystem
+    """
+    if path.startswith("hdfs://"):
+        global _HDFS
+        if _HDFS is None:
+            _HDFS = HdfsFileSystem()
+        return _HDFS
+    global _LOCAL
+    if _LOCAL is None:
+        _LOCAL = LocalFileSystem()
+    return _LOCAL