[FEAT] add HDFS support for DynamicEmb checkpoint dump and load#423
[FEAT] add HDFS support for DynamicEmb checkpoint dump and load#423PWZER wants to merge 1 commit into
Conversation
![greptile-apps[bot]](https://avatars.github.com/in/867647?s=60&v=4){kind=small}
| fopt_states = None | ||
| if include_optim: | ||
| fopt_states = fs.open(opt_file_path, mode) | ||
|
|
||
| for keys, embeddings, opt_states_batch, scores in export_keys_values_iter( | ||
| state, device=device, table_id=table_id | ||
| ): | ||
| fkey.write(keys.cpu().numpy().tobytes()) | ||
| fembedding.write(embeddings.cpu().numpy().tobytes()) | ||
| if state.evict_strategy == EvictStrategy.KLru: | ||
| scores = timestamp - scores | ||
| fscore.write(scores.cpu().numpy().tobytes()) | ||
| if fopt_states and opt_states_batch is not None: | ||
| to_write = truncate_optimizer_states_for_checkpoint( | ||
| state.optimizer, | ||
| state.table_emb_dims_cpu[table_id], | ||
| opt_states_batch, | ||
| ) | ||
| fopt_states.write(to_write.cpu().numpy().tobytes()) | ||
|
|
||
| fkey.close() | ||
| fembedding.close() | ||
| if fscore: | ||
| fscore.close() | ||
| if fopt_states: | ||
| fopt_states.close() | ||
| if fopt_states: | ||
| fopt_states.close() |
There was a problem hiding this comment.
fopt_states file handle not closed on exception
fopt_states is opened with fs.open(opt_file_path, mode) inside the with block but is not part of the context-manager chain. If export_keys_values_iter raises (e.g., CUDA OOM or a network error writing to HDFS), execution exits the with block without hitting if fopt_states: fopt_states.close(), leaving the file handle open and the partially-written optimizer file unflushed/untruncated. The three sibling files (fkey, fembedding, fscore) are correctly protected by the with statement; fopt_states should be treated consistently — either added to the with chain or wrapped in its own with block or try/finally.
| _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 |
There was a problem hiding this comment.
_HDFS singleton silently ignores a second HDFS host
get_filesystem() returns the same HdfsFileSystem instance for every hdfs:// URI. _parse_uri guards with if self._host is not None: return, so the host/port captured from the very first URI is permanent. If a caller later uses a path on a different NameNode (e.g., hdfs://standby:8020/...), the singleton silently continues to talk to the original host. A production job that mixes two HDFS clusters — or that switches between a primary and DR NameNode — would silently read/write to the wrong cluster with no error.
| 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) |
There was a problem hiding this comment.
os.path.join used to build HDFS sub-paths
os.path.join(path, collection_path) is called twice (lines 178 and 185) after the HDFS path refactor. On Linux this happens to produce the correct hdfs://host:port/parent/child string, but os.path.join is OS-aware path manipulation, not URI manipulation. If collection_path ever begins with /, os.path.join silently discards the base, producing a broken path like /collection instead of hdfs://host:port/checkpoints/model_v1/collection. Using f"{path.rstrip('/')}/{collection_path}" or a filesystem-aware join would be unambiguous.
|
Hi @PWZER , can you also update it to the documents, for example: using |
jiashuy
left a comment
jiashuy
left a comment
There was a problem hiding this comment.
LGTM, I will explore how to test HDFS in a single node with multi GPUs and then add relevant tests to the current CI tests
3 participants
Add support for reading/writing DynamicEmb checkpoints via HDFS using
fsspec[arrow], enabling distributed training jobs to store checkpointson Hadoop clusters without any API changes.
What changed
New:
filesystem.py— lightweight filesystem abstractionTwo backends, selected automatically by path prefix:
LocalFileSystemHdfsFileSystemhdfs://host:port/...fsspec[arrow]>=2023.0.0(optional)HdfsFileSystemlazily initialisespyarrow.HadoopFileSystemon firstI/O, passing host/port explicitly so
core-site.xmldoes not needto set
fs.defaultFS.Refactor: route all checkpoint I/O through
get_filesystem(path)Replaced direct
os.path.exists,os.makedirs,os.listdir,os.path.getsize,open, andglob.globcalls in:batched_dynamicemb_tables.py—find_files,get_loading_filesdump_load.py—DynamicEmbDump,DynamicEmbLoadexportable_tables.py—InferenceEmbeddingCollection.savekey_value_table.py—_dump_table,_iter_batches_from_files,_validate_load_meta,save_to_json,load_from_jsonPath handling for HDFS glob results
find_filesnow handles thehdfs://scheme prefix: stripped for globpattern matching, then re-attached to returned paths so downstream
consumers can use them directly, regardless of backend behaviour.
Extras dependency in
setup.pyUsage (no API change)
Checklist