求助,ray.data.DataSet 数据级联 map 如何限制处理进度,以降低内存开销

【Ray使用环境】测试 Python 3.11
【Ray版本和类库】2.39.0
【使用现场】
运行环境:单机8卡4090 512GB 内存
处理逻辑如下所示:
先有一个路径列表,可能有几万条;我全部转换成 ds: ray.data.DataSet。
再声明他的执行流程,例如以下的例子:
ds

  • Load 这个时候会读取这个路径的文件,这个文件可能很大,有几百MB 或者几个GB
  • Transform1(num_gpus=1,concurrency=2)
  • Transform2(num_gpus=0.25,concurrency=8) 以此举例一些 GPU 操作
  • Store 这个时候会把处理好的文件放回其他二级存储

我单机只有 512G 的内存,肯定没法把这些文件全部加载到内存。所以是否有一些机制,监测到内存不够的时候(例如已经使用了 400G/512G),就自动停止 Load 的步骤,等到 Store 步骤结束之后,有了空余内存,再继续 Load。