Performance imporve: mmap and file handler cache

[yezhengmao1]

🚀 Feature Request

Optimize Streaming mds reader I/O Performance with mmap and File Handle Caching

This proposal suggests implementing a file handle cache (using an LRU policy) and memory-mapped (mmap) I/O to resolve a significant performance bottleneck in the streaming mds reader.

Motivation

I'm frustrated by the poor performance of our streaming data pipeline. Profiling has revealed a critical I/O bottleneck within the mds reader component. The current implementation is inefficient, and a large portion of execution time is spent on I/O overhead rather than useful work.

Our analysis indicates the time is spent as follows:

• Frequent File Handle Operations: Approximately $1/3$ of the execution time is consumed by repeatedly opening and closing file handles (open()/close() system calls). This is a major source of overhead.
• Inefficient I/O Reads: Another $1/3$ of the time is spent on the actual data transfer during I/O read operations.

This inefficient I/O pattern leads to high system call overhead, slow data loading, and underutilization of compute resources, significantly hindering the throughput of our training and inference workloads.

[Optional] Implementation

To address this, I propose a two-part solution:

1. File Handle LRU Cache

   • Implement a Least Recently Used (LRU) cache to manage open file handles.
   • When the reader needs to access a file, it will first check the cache for a valid, open file handle. If one exists, it will be reused, completely avoiding the expensive open() system call.
   • If the cache is full, the least recently used file handle will be closed to free up system resources.
   • The size of this cache should be configurable to allow users to tune it based on their system's file handle limits and specific data access patterns.

2. Memory-Mapped I/O (mmap)

   • Instead of using traditional read() system calls, we should use mmap to map file contents directly into the process's virtual address space.
   • Benefits:
     • Reduces System Calls: After the initial mmap call, subsequent data access does not require further system calls.
     • Eliminates Data Copies: mmap avoids the extra memory copy from the kernel's page cache to the user-space buffer that occurs with read().
     • Leverages OS Optimizations: It allows the operating system's virtual memory manager to handle paging and pre-fetching in a much more efficient manner.

These two strategies work synergistically: the LRU cache provides a persistent file descriptor, and mmap uses that descriptor to establish a highly efficient "zero-copy" read pathway.

Additional context

By implementing these changes, we can expect to drastically cut down the I/O-related overhead, which currently accounts for roughly $2/3$ of the execution time. This will lead to a substantial improvement in data loading throughput.

I can provide detailed profiling reports (RP) to clearly demonstrate the current bottleneck if needed.

[XiaohanZhangCMU]

@yezhengmao1 this is a very interesting proposal! Are you interested in creating a PR?

streaming is currently in KTLO mode since we don't have resource to do development. But we welcome contributions from oss community for any performance improvement like this.

Two side notes to consider:

• some platforms provide fused apis which sort of implement the cache and mmap logics you proposed and we do observe nice performance improvement (~50%)
• hopefully this is rare but in the cases where cache miss is high (for example small nvme + complete shuffle), the proposed approach may lead to a throughput drop.

[XiaohanZhangCMU]

Can you try #965 and see if the performance improved?