The validation logic here is incorrect. Using all() means an error is raised only if all offsets are out of bounds or all tensor sizes exceed the page size. This should be any(), so that an error is raised if any item is invalid. This bug could lead to out-of-bounds memory access or silent data corruption.

                any(offset < 0 or offset + size > self.size
                    for offset, size in zip(offsets, sizes)),
                any(size > self.bytes_per_page for size in sizes),

The gemini's comment seems to be correct. Can you check @ibifrost?

Using all() here was indeed a logic error, I've fix it

The observe method incorrectly attempts to iterate over num_failed_save and num_failed_load from transfer_stats_data. These values are integers, not iterables, which will cause a TypeError at runtime when metrics are being observed. The loop should be removed, and inc() should be called directly with the integer value.

            counter_obj[engine_idx].inc(transfer_stats_data[counter_item_key])

also fixed

This block of code, which skips the expert weight filter when EPLB is enabled, is a duplicate of the code block that immediately follows it. This duplication should be removed to improve code clarity and maintainability.

This block of code, which prepares the out_kwargs for the FlashInfer kernel, is a duplicate of the code block at lines 195-224. This redundant code should be removed to improve maintainability.

A naive question: why we need to define extra package-data here?

We rely on runtime compilation via
"
from torch.utils.cpp_extension import load
...
hf3fs_utils = load(
name="hf3fs_utils",
sources=[f"{root}/utils/hf3fs_utils.cpp"],
extra_include_paths=[cuda_include_path],
)
"
in hf3fs_client.py now, that's why the hf3fs_utils.cpp file should be included into pkg here.

@simon-mo Just wondering is it a common practice to introduce the c/c++ sources like this? Or there are some other common practices?

package data is fine. but we put those in setup.py

Sure, I've move the logic to setup.py as suggested.

The gemini's comment seems to be correct. Can you check @ibifrost?

Do we need to call free_buffer before self._fail_task?

Yes, we need to release the buffers in error paths. I've added a check before calling _fail_task to ensure resources are cleaned up.

Do we need to free the existing allocation_results[key] before overwriting it? I'm not so sure about the logic here.

Right, just to clarify the logic here:
When a key is already complete (key_meta.is_complete()), we reuse its existing page from metadata, so we do not need to free that existing page because it's still valid and in use.

However, since we pre-allocated a new page for every key at the start of the function, hitting this branch means we have a newly allocated page that we won't actually use. The original code failed to release this specific unused pre-allocated page, which could lead to a potential leak.

So here need to add logic to keep the existing mapping but release the unused allocated_pages, I'll add a test case to ensure the free page count stays consistent when reusing existing keys.

Do we need to do any check to prevent double-close or double-del?

You're right, I've moved the cleanup logic into a _release_resources method and added checks to ensures the close operation is idempotent.