loaderx 0.3.0__tar.gz → 0.4.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
loaderx-0.4.1/PKG-INFO ADDED
@@ -0,0 +1,98 @@
1
+ Metadata-Version: 2.4
2
+ Name: loaderx
3
+ Version: 0.4.1
4
+ Summary: A record-based data runtime, focused on delivering extreme throughput and low latency
5
+ Author-email: Ben0i0d <ben0i0d@foxmail.com>
6
+ License: MIT License
7
+
8
+ Copyright (c) 2025 EOELAB AI Research
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Homepage, https://codeberg.org/eoelab/loaderx
29
+ Project-URL: Documentation, https://codeberg.org/eoelab/loaderx
30
+ Project-URL: Source, https://codeberg.org/eoelab/loaderx
31
+ Project-URL: Bug Tracker, https://codeberg.org/eoelab/loaderx
32
+ Keywords: flax,python,dataloader
33
+ Classifier: Development Status :: 3 - Alpha
34
+ Classifier: Intended Audience :: Developers
35
+ Classifier: License :: OSI Approved :: MIT License
36
+ Classifier: Programming Language :: Python :: 3
37
+ Requires-Python: ==3.13
38
+ Description-Content-Type: text/markdown
39
+ License-File: LICENSE
40
+ Requires-Dist: numpy
41
+ Requires-Dist: blosc2
42
+ Requires-Dist: cffi
43
+ Dynamic: license-file
44
+
45
+ # Loaderx
46
+ A compact and high-performance single-machine data loader designed for JAX/Flax.
47
+
48
+ **Only Linux_amd64**
49
+ ## Design Philosophy
50
+
51
+ loaderx is built around several core principles:
52
+
53
+ 1. A pragmatic approach that prioritizes minimal memory overhead and minimal dependencies.
54
+ 2. A strong focus on single-machine training workflows.
55
+ 3. We implement based on NumPy semantics, supporting Blosc2 backends.
56
+ 4. An **immortal (endless) step-based data loader**, rather than the traditional epoch-based design—better aligned with modern ML training practices.
57
+
58
+ ## Zsampler
59
+ Index Generator: a high-performance sampler implemented in Zig
60
+
61
+ 1. Sequential generation: indices are produced by traversing the index space in order.
62
+ * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
63
+ 2. Random generation: indices are sampled randomly from the index space.
64
+ * Global random: a set of samples is drawn randomly from the entire index space.
65
+
66
+ ## Convert a NumPy tensor to Blosc2
67
+ ```
68
+ import blosc2
69
+ import numpy as np
70
+
71
+ np_arr = np.load('arr.npy',mmap_mode='r')
72
+ b2_arr = blosc2.asarray(np_arr, urlpath="arr.b2nd", mode="w")
73
+ ```
74
+ Then, you can use `b2_arr` or load as `b2_arr = blosc2.open("arr.b2nd")`
75
+
76
+ ## Current Limitations
77
+ Currently, loaderx only supports single-host environments and does not yet support multi-host training.
78
+
79
+ ## Quick Start
80
+ ```
81
+ from loaderx import Dataset, DataLoader
82
+
83
+ dataset = Dataset('train_data.b2nd')
84
+ labelset = Dataset('train_label.b2nd')
85
+
86
+ loader = DataLoader(dataset, labelset)
87
+
88
+ for i, batch in enumerate(loader):
89
+ if i >= 256:
90
+ break
91
+
92
+ print(batch['data'].shape)
93
+ print(batch['label'].shape)
94
+ ```
95
+
96
+ ### Integrating with JAX/Flax
97
+
98
+ For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
@@ -0,0 +1,54 @@
1
+ # Loaderx
2
+ A compact and high-performance single-machine data loader designed for JAX/Flax.
3
+
4
+ **Only Linux_amd64**
5
+ ## Design Philosophy
6
+
7
+ loaderx is built around several core principles:
8
+
9
+ 1. A pragmatic approach that prioritizes minimal memory overhead and minimal dependencies.
10
+ 2. A strong focus on single-machine training workflows.
11
+ 3. We implement based on NumPy semantics, supporting Blosc2 backends.
12
+ 4. An **immortal (endless) step-based data loader**, rather than the traditional epoch-based design—better aligned with modern ML training practices.
13
+
14
+ ## Zsampler
15
+ Index Generator: a high-performance sampler implemented in Zig
16
+
17
+ 1. Sequential generation: indices are produced by traversing the index space in order.
18
+ * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
19
+ 2. Random generation: indices are sampled randomly from the index space.
20
+ * Global random: a set of samples is drawn randomly from the entire index space.
21
+
22
+ ## Convert a NumPy tensor to Blosc2
23
+ ```
24
+ import blosc2
25
+ import numpy as np
26
+
27
+ np_arr = np.load('arr.npy',mmap_mode='r')
28
+ b2_arr = blosc2.asarray(np_arr, urlpath="arr.b2nd", mode="w")
29
+ ```
30
+ Then, you can use `b2_arr` or load as `b2_arr = blosc2.open("arr.b2nd")`
31
+
32
+ ## Current Limitations
33
+ Currently, loaderx only supports single-host environments and does not yet support multi-host training.
34
+
35
+ ## Quick Start
36
+ ```
37
+ from loaderx import Dataset, DataLoader
38
+
39
+ dataset = Dataset('train_data.b2nd')
40
+ labelset = Dataset('train_label.b2nd')
41
+
42
+ loader = DataLoader(dataset, labelset)
43
+
44
+ for i, batch in enumerate(loader):
45
+ if i >= 256:
46
+ break
47
+
48
+ print(batch['data'].shape)
49
+ print(batch['label'].shape)
50
+ ```
51
+
52
+ ### Integrating with JAX/Flax
53
+
54
+ For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
@@ -0,0 +1,4 @@
1
+ from .dataset import Dataset
2
+ from .dataloader import DataLoader
3
+
4
+ __version__ = "0.4.1"
@@ -1,10 +1,28 @@
1
- import os
2
1
  import time
3
- import ctypes
2
+ from pathlib import Path
3
+
4
4
  import numpy as np
5
- from typing import Union
5
+ from cffi import FFI
6
+
7
+ ffi = FFI()
8
+
9
+ ffi.cdef("""
10
+ typedef struct Sampler Sampler;
11
+
12
+ Sampler* cinit(
13
+ uint64_t length,
14
+ uint64_t batch_size,
15
+ uint8_t mode,
16
+ uint64_t seed
17
+ );
6
18
 
7
- from lib import libsampler
19
+ void cdeinit(Sampler* s);
20
+
21
+ void cnext(Sampler* s, uint64_t* indices_ptr, size_t indices_len);
22
+ """)
23
+
24
+ lib_path = Path(__file__).resolve().parent / "lib" / "libsampler.so"
25
+ lib = ffi.dlopen(str(lib_path))
8
26
 
9
27
  class Sampler:
10
28
  class Mode:
@@ -13,20 +31,30 @@ class Sampler:
13
31
 
14
32
  def __init__(self, length: int, batch_size: int, mode: int, seed: int = 42):
15
33
  # sampler
16
- self.sampler = libsampler.Sampler.init(length, batch_size, mode, seed)
34
+ self.sampler = lib.cinit(length, batch_size, mode, seed)
35
+ if self.sampler == ffi.NULL:
36
+ raise MemoryError("Failed to create sampler")
17
37
  # indices
18
38
  self.indices = np.zeros(batch_size, dtype=np.uint64)
19
39
 
20
40
  # step
21
41
  def next(self):
22
- self.sampler.next(self.indices)
42
+ ptr = ffi.cast("uint64_t *", self.indices.ctypes.data)
43
+ lib.cnext(self.sampler, ptr, self.indices.size)
44
+ # return copy to avoid loss
45
+ return self.indices.copy()
23
46
 
24
47
  # iterator
25
48
  def __iter__(self):
26
49
  return self
27
50
  def __next__(self):
28
- self.next()
29
- return self.indices
51
+ return self.next()
52
+
53
+ # clean
54
+ def __del__(self):
55
+ if hasattr(self, "sampler") and self.sampler != ffi.NULL:
56
+ lib.cdeinit(self.sampler)
57
+ self.sampler = ffi.NULL
30
58
 
31
59
  def run():
32
60
  from itertools import islice
@@ -44,7 +72,7 @@ def run():
44
72
  for indices in islice(sampler, n_steps):
45
73
  print(indices)
46
74
 
47
- # Benchmark, zig-sampler speeder 2.2x
75
+ # Benchmark, zig-sampler speeder 3.66x
48
76
  def bench():
49
77
  length = 1_000_000
50
78
  batch_size = 8192
@@ -47,8 +47,7 @@ class DataLoader:
47
47
  Sample indices from the dataset and put them into the index queue.
48
48
  """
49
49
  while not self.stop_signal.is_set():
50
- self.sampler.next()
51
- self.indices.put(self.sampler.indices.copy())
50
+ self.indices.put(self.sampler.next())
52
51
 
53
52
  def _prefetch(self, transform):
54
53
  """
@@ -94,9 +93,6 @@ class DataLoader:
94
93
 
95
94
  for thread in self.threads:
96
95
  thread.join()
97
-
98
- self.dataset.close()
99
- self.labelset.close()
100
96
 
101
97
  def __del__(self):
102
98
  self.close()
@@ -0,0 +1,39 @@
1
+ import blosc2
2
+ import numpy as np
3
+
4
+ class Dataset:
5
+ def __init__(self, path):
6
+ """
7
+ Initialize a dataset.
8
+ """
9
+ self._array = blosc2.open(path)
10
+ def __getitem__(self, idx):
11
+ """
12
+ Get the item at index idx from the dataset.
13
+
14
+ Args:
15
+ idx (int): Index of the item to retrieve.
16
+
17
+ Returns:
18
+ np.ndarray: The item at index idx from the dataset.
19
+ """
20
+ return self._array[idx]
21
+ def __len__(self):
22
+ """
23
+ Return the number of samples in the dataset.
24
+
25
+ Returns:
26
+ int: The number of samples in the dataset.
27
+ """
28
+ return len(self._array)
29
+ def __getitems__(self, idxs):
30
+ """
31
+ Get the items at index idxs from the dataset.
32
+
33
+ Args:
34
+ idxs (list of int): indices of the items to retrieve.
35
+
36
+ Returns:
37
+ np.ndarray: The items at index idxs from the dataset.
38
+ """
39
+ return self._array[idxs]
Binary file
@@ -0,0 +1,98 @@
1
+ Metadata-Version: 2.4
2
+ Name: loaderx
3
+ Version: 0.4.1
4
+ Summary: A record-based data runtime, focused on delivering extreme throughput and low latency
5
+ Author-email: Ben0i0d <ben0i0d@foxmail.com>
6
+ License: MIT License
7
+
8
+ Copyright (c) 2025 EOELAB AI Research
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Homepage, https://codeberg.org/eoelab/loaderx
29
+ Project-URL: Documentation, https://codeberg.org/eoelab/loaderx
30
+ Project-URL: Source, https://codeberg.org/eoelab/loaderx
31
+ Project-URL: Bug Tracker, https://codeberg.org/eoelab/loaderx
32
+ Keywords: flax,python,dataloader
33
+ Classifier: Development Status :: 3 - Alpha
34
+ Classifier: Intended Audience :: Developers
35
+ Classifier: License :: OSI Approved :: MIT License
36
+ Classifier: Programming Language :: Python :: 3
37
+ Requires-Python: ==3.13
38
+ Description-Content-Type: text/markdown
39
+ License-File: LICENSE
40
+ Requires-Dist: numpy
41
+ Requires-Dist: blosc2
42
+ Requires-Dist: cffi
43
+ Dynamic: license-file
44
+
45
+ # Loaderx
46
+ A compact and high-performance single-machine data loader designed for JAX/Flax.
47
+
48
+ **Only Linux_amd64**
49
+ ## Design Philosophy
50
+
51
+ loaderx is built around several core principles:
52
+
53
+ 1. A pragmatic approach that prioritizes minimal memory overhead and minimal dependencies.
54
+ 2. A strong focus on single-machine training workflows.
55
+ 3. We implement based on NumPy semantics, supporting Blosc2 backends.
56
+ 4. An **immortal (endless) step-based data loader**, rather than the traditional epoch-based design—better aligned with modern ML training practices.
57
+
58
+ ## Zsampler
59
+ Index Generator: a high-performance sampler implemented in Zig
60
+
61
+ 1. Sequential generation: indices are produced by traversing the index space in order.
62
+ * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
63
+ 2. Random generation: indices are sampled randomly from the index space.
64
+ * Global random: a set of samples is drawn randomly from the entire index space.
65
+
66
+ ## Convert a NumPy tensor to Blosc2
67
+ ```
68
+ import blosc2
69
+ import numpy as np
70
+
71
+ np_arr = np.load('arr.npy',mmap_mode='r')
72
+ b2_arr = blosc2.asarray(np_arr, urlpath="arr.b2nd", mode="w")
73
+ ```
74
+ Then, you can use `b2_arr` or load as `b2_arr = blosc2.open("arr.b2nd")`
75
+
76
+ ## Current Limitations
77
+ Currently, loaderx only supports single-host environments and does not yet support multi-host training.
78
+
79
+ ## Quick Start
80
+ ```
81
+ from loaderx import Dataset, DataLoader
82
+
83
+ dataset = Dataset('train_data.b2nd')
84
+ labelset = Dataset('train_label.b2nd')
85
+
86
+ loader = DataLoader(dataset, labelset)
87
+
88
+ for i, batch in enumerate(loader):
89
+ if i >= 256:
90
+ break
91
+
92
+ print(batch['data'].shape)
93
+ print(batch['label'].shape)
94
+ ```
95
+
96
+ ### Integrating with JAX/Flax
97
+
98
+ For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
@@ -5,7 +5,6 @@ loaderx/__init__.py
5
5
  loaderx/_sampler.py
6
6
  loaderx/dataloader.py
7
7
  loaderx/dataset.py
8
- loaderx/utils.py
9
8
  loaderx.egg-info/PKG-INFO
10
9
  loaderx.egg-info/SOURCES.txt
11
10
  loaderx.egg-info/dependency_links.txt
@@ -0,0 +1,3 @@
1
+ numpy
2
+ blosc2
3
+ cffi
@@ -19,7 +19,7 @@ classifiers = [
19
19
  "License :: OSI Approved :: MIT License",
20
20
  "Programming Language :: Python :: 3",
21
21
  ]
22
- dependencies = ["numpy", "array-record", "tqdm"]
22
+ dependencies = ["numpy", "blosc2", "cffi"]
23
23
 
24
24
  [tool.setuptools]
25
25
  packages = ["loaderx"]
loaderx-0.3.0/PKG-INFO DELETED
@@ -1,260 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: loaderx
3
- Version: 0.3.0
4
- Summary: A record-based data runtime, focused on delivering extreme throughput and low latency
5
- Author-email: Ben0i0d <ben0i0d@foxmail.com>
6
- License: MIT License
7
-
8
- Copyright (c) 2025 EOELAB AI Research
9
-
10
- Permission is hereby granted, free of charge, to any person obtaining a copy
11
- of this software and associated documentation files (the "Software"), to deal
12
- in the Software without restriction, including without limitation the rights
13
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
- copies of the Software, and to permit persons to whom the Software is
15
- furnished to do so, subject to the following conditions:
16
-
17
- The above copyright notice and this permission notice shall be included in all
18
- copies or substantial portions of the Software.
19
-
20
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
- SOFTWARE.
27
-
28
- Project-URL: Homepage, https://codeberg.org/eoelab/loaderx
29
- Project-URL: Documentation, https://codeberg.org/eoelab/loaderx
30
- Project-URL: Source, https://codeberg.org/eoelab/loaderx
31
- Project-URL: Bug Tracker, https://codeberg.org/eoelab/loaderx
32
- Keywords: flax,python,dataloader
33
- Classifier: Development Status :: 3 - Alpha
34
- Classifier: Intended Audience :: Developers
35
- Classifier: License :: OSI Approved :: MIT License
36
- Classifier: Programming Language :: Python :: 3
37
- Requires-Python: ==3.13
38
- Description-Content-Type: text/markdown
39
- License-File: LICENSE
40
- Requires-Dist: numpy
41
- Requires-Dist: array-record
42
- Requires-Dist: tqdm
43
- Dynamic: license-file
44
-
45
- # loaderx
46
- A record-based data runtime, focused on delivering extreme throughput and low latency
47
-
48
- **Only Python3.13_Linux_amd64**
49
-
50
- ## Sampler
51
- a high-performance sampler implemented in Zig
52
- ### Index Generator
53
- 1. Sequential generation: indices are produced by traversing the index space in order.
54
- * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
55
- 2. Random generation: indices are sampled randomly from the index space.
56
- * Global random: a set of samples is drawn randomly from the entire index space.
57
-
58
- ## zrecord
59
- 基于Zig实现一个并发友好、实现更简单、单机性能更优的 record 存储系统
60
-
61
- 1. ZRecord基于 record 思想设计,我们不假定record间存在顺序关系,对于顺序语义应在record内部表达或建立timestamp
62
- 2. record是相互独立的逻辑记录,仅在 Python 侧提供与 NumPy 兼容的 array 接口投影
63
- 3. zrecord是无类型的,数据只是Bytes,类型解析由上层完成
64
- 4. chunk只是存储管理单元,提高并发性能,不具备语义
65
-
66
- ### 存储形式
67
- 1. ZRecord 仅包含一个多流数据集,长度为 N(N 个 record)
68
- * 索引空间为 [0, N),所有索引与切块语义等价于 gather 操作
69
- * 多流,类似 `set['A'], set['B']`, 但必须等长
70
- ```
71
- Dataset:
72
- record_id = 0..N-1
73
-
74
- Field A: bytes / tensor / audio / text
75
- Field B: bytes / label / meta
76
- ```
77
- 2. 为降低存储体积,支持粒度为record级的透明压缩,以下是可选的压缩方法
78
- ```
79
- | 方法 | 算法 |
80
- | ----- | ------- |
81
- | raw | 无压缩 |
82
- | flate | Deflate |
83
- ```
84
- 3. 通过offset支持高效随机访问
85
- * offset表项是一个三元组 [[chunk_id : u32, offset : u64, physical_length : u32], ...],长度为N。将外部索引空间(outside[0, N])指向一个实际存储地址
86
- * chunk_id为u32,表示具体chunk
87
- * offset为u64, 表示具体chunk内的偏移
88
- * physical_length为u32,表示record数据大小
89
- * offset表每项的字节大小最好为 8 的整数倍,提升对齐和缓存效率
90
- * offset表是mmap访问模式,由于等长,直接将 idx 转换为 ptr + 16*idx
91
- * 不同位宽表示的数值范围
92
- ```
93
- u16 → 2^16 - 1 (65535)
94
- u32 → 2^32 - 1 (4.29e9)
95
- u64 → 2^64 - 1 (1.84e19)
96
- ```
97
- * 不同位宽表示的数据大小
98
- ```
99
- u16 → 64 KiB
100
- u32 → 4 GiB
101
- u64 → 16 EiB
102
- ```
103
- ```
104
- outside idx (global)
105
-
106
- ↓offset[idx]
107
-
108
- chunk_id, offset, physical_length
109
- ```
110
- 4. ZRecord的存储是元数据+分块的格式
111
- 1. 元数据
112
- * meta.json:描述全局元数据,包括全局参数与Field参数
113
- ```
114
- {
115
- "length": 65536,
116
- "chunks": { "num": 12, "size": 8192, "unfull": {[8, write_pos, record_count], ... } },
117
- "fields": [ {"name": "A", "dtype": "f32", "compress": "flate"}, {"name": "B", "dtype": "i32", "compress": "raw"} ]
118
- }
119
- ```
120
- * Field_offset.zr:索引表
121
- ```
122
- 0 0 255
123
- 0 256 511
124
- ...
125
- ```
126
- 2. 分块数据
127
- * chunk/x.zr: 分块的chunk数据,chunk不区分数据所属
128
-
129
- 存储格式如下
130
- ```
131
- dataset
132
- ├── meta.json
133
- ├── A.offset
134
- ├── B.offset
135
- ├── chunk
136
- │ ├── 0.zr
137
- └──└── 1.zr
138
- ```
139
-
140
- ### 执行器
141
-
142
- 负责维护任务队列,工作线程,handle
143
-
144
- 1. 任务队列:
145
- * write_task:写任务队列,压入[ops, record],等待writer弹出
146
- * read_task:读任务队列,预分配batch并压入[batch, pos, chunk_id, offset, physical_length],等待reader弹出
147
- * batch: 具体读取请求下创建的与indices等长的二元组
148
- * pos:在batch中的位置,保证等序返回
149
- * gc_task:垃圾回收任务队列,压入batch,等待cleaner弹出
150
-
151
- 2. 工作线程:执行具体任务的工作线程
152
- * writer:线程池执行写任务,但每个线程固定对应一个chunk
153
- * writer只会向自己的chunk写入(写满时新建chunk提交给执行器完成)
154
- * num_writer推荐:1-4
155
- * reader:线程池执行读任务
156
- * num_reader推荐:8-32
157
- * cleaner:线程池执行垃圾回收任务
158
- * num_cleaner推荐:2-4
159
-
160
- 3. handle:
161
- 1. mmap:维护全局offset、chunk文件的handle
162
- * 初始化阶段注册,关闭前释放
163
- * 对于full的chunk,使用mmap+close(只读),对于unfull的chunk使用mmap(读写)。一旦写满(writer提交),新建chunk并返回给writer,原mmap close。(避免过多占用fd)
164
- * 对于chunk/offset文件,动态完成分块增长,基于ftruncate+mremap实现
165
- * 扩容offset表,一次增长 1M项(16MiB)
166
- * 扩容chunk文件,一次增长 1GiB
167
- 2. batch:维护每次返回batch的handle,也就是[ptr, logical_length]中ptr的合规性
168
- * 工作线程完成后注册
169
- * 必须调用函数来手动释放,释放时batch压入GC队列
170
-
171
- #### 在线任务
172
-
173
- 1. 写任务:写入只允许chunk-level append-only,其余操作基于offset重定向,缺失的field构造为全为0的项(physical_length为0代表无)
174
- 1. 追加:在chunk当前写入位置追加一个record,offset表更新length+1条目,全局长度增加
175
- * 写满:写入完成时检查,如果record_count达到chunk_size, 新建chunk
176
- * 写入越界:捕捉错误,并对chunk扩容
177
- * offset[ length+1 ]不存在:捕捉错误,并对offset扩容
178
- 2. 修改:追加一个record,并修改offset项
179
- 3. 删除:删除offset项(将目标项替换为最后一项),length减一(越界访问属于UB行为)
180
-
181
- 2. 读任务:用一个indices(索引数组),从原数组里访问指定位置的元素(gather)
182
- 1. 对于无压缩的数据,将直接返回文件内地址实现zero-copy
183
-
184
- 读取流程
185
- ```
186
- [idx]
187
-
188
- [batch, pos, idx]
189
- ↓field_offset
190
- [batch, pos, chunk_id, offset, physical_length]
191
-
192
- [ptr, logical_length]
193
- ```
194
-
195
- 3. 垃圾回收任务:cleaner从GC队列中释放对应batch
196
- 1. 对于无压缩的数据,不释放内存,只释放 batch 本身
197
- 2. 对于有压缩的数据,释放内存
198
- 3. GC后访问ptr属于UB行为
199
-
200
- #### 离线任务
201
-
202
- 1. 重平衡(Rebalance):zrecord经过大量写操作后,会出现chunk访问不平衡,通过重平衡来恢复chunk的可访问性
203
- * 实现:根据offset进行重写入,并在完成后替换
204
- * Rebalance只允许手动运行,确保用户知情,提供一个利用率(sum(length)/sum(chunk_size*num_chunks))
205
-
206
- ### 分布式扩展
207
-
208
- **该部分暂时不会实现,只是提前架构**
209
-
210
- 1. 分布式下,我们会得到一个更高的层次-cluster,单机将变为cluster内的node
211
- 2. node持有一个shard,包含若干个chunk
212
- 3. 添加一个indirection表,将全局路径映射到具体node的表,indirection也可以使用hash实现,从而无锁,索引路径变为
213
- ```
214
- outside idx (global)
215
-
216
- ↓indirection[idx]
217
-
218
- node, idx
219
-
220
- ↓offset[idx]
221
-
222
- chunk_id, offset, physical_length
223
- ```
224
-
225
- ## Convert a NumPy tensor to Array_record
226
-
227
- *This will create a directory containing file shards, which helps improve I/O performance.*
228
-
229
- ```
230
- import numpy as np
231
- from loaderx import converter
232
-
233
- train_data = np.load('train_data.npy',mmap_mode='r')
234
- converter(train_data, 'train_data')
235
- ```
236
-
237
- ## Current Limitations
238
- Currently, loaderx only supports single-host environments and does not yet support multi-host training.
239
-
240
- ## Quick Start
241
- ```
242
- import numpy as np
243
- from loaderx import NPDataset, ARDataset, DataLoader
244
-
245
- dataset = ARDataset('train_data')
246
- labelset = NPDataset('xsub/train_label.npy')
247
-
248
- loader = DataLoader(dataset, labelset)
249
-
250
- for i, batch in enumerate(loader):
251
- if i >= 256:
252
- break
253
-
254
- print(batch['data'].shape)
255
- print(batch['label'].shape)
256
- ```
257
-
258
- ### Integrating with JAX/Flax
259
-
260
- For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
loaderx-0.3.0/README.md DELETED
@@ -1,216 +0,0 @@
1
- # loaderx
2
- A record-based data runtime, focused on delivering extreme throughput and low latency
3
-
4
- **Only Python3.13_Linux_amd64**
5
-
6
- ## Sampler
7
- a high-performance sampler implemented in Zig
8
- ### Index Generator
9
- 1. Sequential generation: indices are produced by traversing the index space in order.
10
- * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
11
- 2. Random generation: indices are sampled randomly from the index space.
12
- * Global random: a set of samples is drawn randomly from the entire index space.
13
-
14
- ## zrecord
15
- 基于Zig实现一个并发友好、实现更简单、单机性能更优的 record 存储系统
16
-
17
- 1. ZRecord基于 record 思想设计,我们不假定record间存在顺序关系,对于顺序语义应在record内部表达或建立timestamp
18
- 2. record是相互独立的逻辑记录,仅在 Python 侧提供与 NumPy 兼容的 array 接口投影
19
- 3. zrecord是无类型的,数据只是Bytes,类型解析由上层完成
20
- 4. chunk只是存储管理单元,提高并发性能,不具备语义
21
-
22
- ### 存储形式
23
- 1. ZRecord 仅包含一个多流数据集,长度为 N(N 个 record)
24
- * 索引空间为 [0, N),所有索引与切块语义等价于 gather 操作
25
- * 多流,类似 `set['A'], set['B']`, 但必须等长
26
- ```
27
- Dataset:
28
- record_id = 0..N-1
29
-
30
- Field A: bytes / tensor / audio / text
31
- Field B: bytes / label / meta
32
- ```
33
- 2. 为降低存储体积,支持粒度为record级的透明压缩,以下是可选的压缩方法
34
- ```
35
- | 方法 | 算法 |
36
- | ----- | ------- |
37
- | raw | 无压缩 |
38
- | flate | Deflate |
39
- ```
40
- 3. 通过offset支持高效随机访问
41
- * offset表项是一个三元组 [[chunk_id : u32, offset : u64, physical_length : u32], ...],长度为N。将外部索引空间(outside[0, N])指向一个实际存储地址
42
- * chunk_id为u32,表示具体chunk
43
- * offset为u64, 表示具体chunk内的偏移
44
- * physical_length为u32,表示record数据大小
45
- * offset表每项的字节大小最好为 8 的整数倍,提升对齐和缓存效率
46
- * offset表是mmap访问模式,由于等长,直接将 idx 转换为 ptr + 16*idx
47
- * 不同位宽表示的数值范围
48
- ```
49
- u16 → 2^16 - 1 (65535)
50
- u32 → 2^32 - 1 (4.29e9)
51
- u64 → 2^64 - 1 (1.84e19)
52
- ```
53
- * 不同位宽表示的数据大小
54
- ```
55
- u16 → 64 KiB
56
- u32 → 4 GiB
57
- u64 → 16 EiB
58
- ```
59
- ```
60
- outside idx (global)
61
-
62
- ↓offset[idx]
63
-
64
- chunk_id, offset, physical_length
65
- ```
66
- 4. ZRecord的存储是元数据+分块的格式
67
- 1. 元数据
68
- * meta.json:描述全局元数据,包括全局参数与Field参数
69
- ```
70
- {
71
- "length": 65536,
72
- "chunks": { "num": 12, "size": 8192, "unfull": {[8, write_pos, record_count], ... } },
73
- "fields": [ {"name": "A", "dtype": "f32", "compress": "flate"}, {"name": "B", "dtype": "i32", "compress": "raw"} ]
74
- }
75
- ```
76
- * Field_offset.zr:索引表
77
- ```
78
- 0 0 255
79
- 0 256 511
80
- ...
81
- ```
82
- 2. 分块数据
83
- * chunk/x.zr: 分块的chunk数据,chunk不区分数据所属
84
-
85
- 存储格式如下
86
- ```
87
- dataset
88
- ├── meta.json
89
- ├── A.offset
90
- ├── B.offset
91
- ├── chunk
92
- │ ├── 0.zr
93
- └──└── 1.zr
94
- ```
95
-
96
- ### 执行器
97
-
98
- 负责维护任务队列,工作线程,handle
99
-
100
- 1. 任务队列:
101
- * write_task:写任务队列,压入[ops, record],等待writer弹出
102
- * read_task:读任务队列,预分配batch并压入[batch, pos, chunk_id, offset, physical_length],等待reader弹出
103
- * batch: 具体读取请求下创建的与indices等长的二元组
104
- * pos:在batch中的位置,保证等序返回
105
- * gc_task:垃圾回收任务队列,压入batch,等待cleaner弹出
106
-
107
- 2. 工作线程:执行具体任务的工作线程
108
- * writer:线程池执行写任务,但每个线程固定对应一个chunk
109
- * writer只会向自己的chunk写入(写满时新建chunk提交给执行器完成)
110
- * num_writer推荐:1-4
111
- * reader:线程池执行读任务
112
- * num_reader推荐:8-32
113
- * cleaner:线程池执行垃圾回收任务
114
- * num_cleaner推荐:2-4
115
-
116
- 3. handle:
117
- 1. mmap:维护全局offset、chunk文件的handle
118
- * 初始化阶段注册,关闭前释放
119
- * 对于full的chunk,使用mmap+close(只读),对于unfull的chunk使用mmap(读写)。一旦写满(writer提交),新建chunk并返回给writer,原mmap close。(避免过多占用fd)
120
- * 对于chunk/offset文件,动态完成分块增长,基于ftruncate+mremap实现
121
- * 扩容offset表,一次增长 1M项(16MiB)
122
- * 扩容chunk文件,一次增长 1GiB
123
- 2. batch:维护每次返回batch的handle,也就是[ptr, logical_length]中ptr的合规性
124
- * 工作线程完成后注册
125
- * 必须调用函数来手动释放,释放时batch压入GC队列
126
-
127
- #### 在线任务
128
-
129
- 1. 写任务:写入只允许chunk-level append-only,其余操作基于offset重定向,缺失的field构造为全为0的项(physical_length为0代表无)
130
- 1. 追加:在chunk当前写入位置追加一个record,offset表更新length+1条目,全局长度增加
131
- * 写满:写入完成时检查,如果record_count达到chunk_size, 新建chunk
132
- * 写入越界:捕捉错误,并对chunk扩容
133
- * offset[ length+1 ]不存在:捕捉错误,并对offset扩容
134
- 2. 修改:追加一个record,并修改offset项
135
- 3. 删除:删除offset项(将目标项替换为最后一项),length减一(越界访问属于UB行为)
136
-
137
- 2. 读任务:用一个indices(索引数组),从原数组里访问指定位置的元素(gather)
138
- 1. 对于无压缩的数据,将直接返回文件内地址实现zero-copy
139
-
140
- 读取流程
141
- ```
142
- [idx]
143
-
144
- [batch, pos, idx]
145
- ↓field_offset
146
- [batch, pos, chunk_id, offset, physical_length]
147
-
148
- [ptr, logical_length]
149
- ```
150
-
151
- 3. 垃圾回收任务:cleaner从GC队列中释放对应batch
152
- 1. 对于无压缩的数据,不释放内存,只释放 batch 本身
153
- 2. 对于有压缩的数据,释放内存
154
- 3. GC后访问ptr属于UB行为
155
-
156
- #### 离线任务
157
-
158
- 1. 重平衡(Rebalance):zrecord经过大量写操作后,会出现chunk访问不平衡,通过重平衡来恢复chunk的可访问性
159
- * 实现:根据offset进行重写入,并在完成后替换
160
- * Rebalance只允许手动运行,确保用户知情,提供一个利用率(sum(length)/sum(chunk_size*num_chunks))
161
-
162
- ### 分布式扩展
163
-
164
- **该部分暂时不会实现,只是提前架构**
165
-
166
- 1. 分布式下,我们会得到一个更高的层次-cluster,单机将变为cluster内的node
167
- 2. node持有一个shard,包含若干个chunk
168
- 3. 添加一个indirection表,将全局路径映射到具体node的表,indirection也可以使用hash实现,从而无锁,索引路径变为
169
- ```
170
- outside idx (global)
171
-
172
- ↓indirection[idx]
173
-
174
- node, idx
175
-
176
- ↓offset[idx]
177
-
178
- chunk_id, offset, physical_length
179
- ```
180
-
181
- ## Convert a NumPy tensor to Array_record
182
-
183
- *This will create a directory containing file shards, which helps improve I/O performance.*
184
-
185
- ```
186
- import numpy as np
187
- from loaderx import converter
188
-
189
- train_data = np.load('train_data.npy',mmap_mode='r')
190
- converter(train_data, 'train_data')
191
- ```
192
-
193
- ## Current Limitations
194
- Currently, loaderx only supports single-host environments and does not yet support multi-host training.
195
-
196
- ## Quick Start
197
- ```
198
- import numpy as np
199
- from loaderx import NPDataset, ARDataset, DataLoader
200
-
201
- dataset = ARDataset('train_data')
202
- labelset = NPDataset('xsub/train_label.npy')
203
-
204
- loader = DataLoader(dataset, labelset)
205
-
206
- for i, batch in enumerate(loader):
207
- if i >= 256:
208
- break
209
-
210
- print(batch['data'].shape)
211
- print(batch['label'].shape)
212
- ```
213
-
214
- ### Integrating with JAX/Flax
215
-
216
- For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
@@ -1,5 +0,0 @@
1
- from .dataset import BaseDataset, NPDataset, ARDataset
2
- from .dataloader import DataLoader
3
- from .utils import converter
4
-
5
- __version__ = "0.3.0"
@@ -1,93 +0,0 @@
1
- import json
2
- import numpy as np
3
- from pathlib import Path
4
- from array_record.python.array_record_data_source import ArrayRecordDataSource
5
-
6
- class BaseDataset:
7
- def __init__(self, path):
8
- """
9
- Initialize a dataset.
10
- """
11
- raise NotImplementedError
12
- def __getitem__(self, idx):
13
- """
14
- Get the item at index idx from the dataset.
15
-
16
- Args:
17
- idx (int): Index of the item to retrieve.
18
-
19
- Returns:
20
- np.ndarray: The item at index idx from the dataset.
21
- """
22
- raise NotImplementedError
23
- def __len__(self):
24
- """
25
- Return the number of samples in the dataset.
26
-
27
- Returns:
28
- int: The number of samples in the dataset.
29
- """
30
- raise NotImplementedError
31
- def __getitems__(self, idxs):
32
- """
33
- Get the items at index idxs from the dataset.
34
-
35
- Args:
36
- idxs (list of int): indices of the items to retrieve.
37
-
38
- Returns:
39
- np.ndarray: The items at index idxs from the dataset.
40
- """
41
- raise NotImplementedError
42
- def close(self):
43
- """
44
- Close the dataset.
45
-
46
- If the dataset is stored in a memory-mapped file,
47
- close the memory map to free up system resources.
48
-
49
- If the dataset is stored in an ArrayRecord file,
50
- close the file to free up system resources.
51
- """
52
- pass
53
-
54
- class NPDataset(BaseDataset):
55
- def __init__(self, path):
56
- self._array = np.load(path, mmap_mode='r')
57
-
58
- def __getitem__(self, idx):
59
- return self._array[idx]
60
-
61
- def __getitems__(self, idxs):
62
- return self._array[idxs]
63
-
64
- def __len__(self):
65
- return len(self._array)
66
-
67
- def close(self):
68
- mmap_obj = getattr(self._array, "_mmap", None)
69
- if mmap_obj is not None:
70
- mmap_obj.close()
71
-
72
- class ARDataset(BaseDataset):
73
- def __init__(self, path):
74
- path = Path(path)
75
- with open(path / "meta.json", "r", encoding="utf-8") as f:
76
- meta = json.load(f)
77
-
78
- self._length = meta["length"]
79
- self._shape = meta["shape"]
80
- self._dtype = meta['dtype']
81
- self._array = ArrayRecordDataSource([str(path / f) for f in meta["blocks"]])
82
-
83
- def __getitem__(self, idx):
84
- return np.frombuffer(self._array[idx], dtype=self._dtype).reshape(self._shape)
85
-
86
- def __getitems__(self, idxs):
87
- return np.frombuffer(b"".join(self._array.__getitems__(idxs)), dtype=self._dtype).reshape(len(idxs), *self._shape)
88
-
89
- def __len__(self):
90
- return self._length
91
-
92
- def close(self):
93
- self._array.__exit__(None, None, None)
Binary file
@@ -1,51 +0,0 @@
1
- import json
2
- import shutil
3
- from tqdm import tqdm
4
- from pathlib import Path
5
- from array_record.python.array_record_module import ArrayRecordWriter
6
-
7
- def converter(data, output_dir, replace=False, blocksize=4096, ar_options= "group_size:1,zstd"):
8
- """
9
- Convert a NumPy array to an ArrayRecord format.
10
-
11
- Args:
12
- data : numpy.ndarray
13
- The input data to be converted.
14
- output_dir : str
15
- The output directory where the ArrayRecord files will be saved.
16
- replace : bool, optional
17
- If True, the output directory will be deleted if it already exists.
18
- Defaults to False.
19
- blocksize : int, optional
20
- The number of records to write to each ArrayRecord file.
21
- Defaults to 4096.
22
- ar_options : str, optional
23
- The options to be passed to the ArrayRecordWriter.
24
- Defaults to "group_size:1,zstd".
25
- """
26
- output_dir = Path(output_dir)
27
- if output_dir.exists():
28
- if replace:
29
- shutil.rmtree(output_dir)
30
- output_dir.mkdir(parents=True, exist_ok=True)
31
- elif any(output_dir.iterdir()):
32
- raise ValueError(f"output_dir is not empty: {output_dir}")
33
- else:
34
- output_dir.mkdir(parents=True, exist_ok=True)
35
-
36
- meta = {"length": data.shape[0], "shape": data.shape[1:], "dtype": str(data.dtype), "blocks":[]}
37
-
38
- bidx = 0
39
- writer = None
40
- for i,tmp in enumerate(tqdm(data)):
41
- if i % blocksize == 0:
42
- if writer is not None:
43
- writer.close()
44
- writer = ArrayRecordWriter(str(output_dir / f"{bidx}.ar"), options=ar_options)
45
- meta["blocks"].append(f"{bidx}.ar")
46
- bidx += 1
47
- writer.write(tmp.tobytes())
48
- writer.close()
49
-
50
- with open(output_dir / "meta.json", 'w', encoding='utf-8') as f:
51
- json.dump(meta, f)
@@ -1,260 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: loaderx
3
- Version: 0.3.0
4
- Summary: A record-based data runtime, focused on delivering extreme throughput and low latency
5
- Author-email: Ben0i0d <ben0i0d@foxmail.com>
6
- License: MIT License
7
-
8
- Copyright (c) 2025 EOELAB AI Research
9
-
10
- Permission is hereby granted, free of charge, to any person obtaining a copy
11
- of this software and associated documentation files (the "Software"), to deal
12
- in the Software without restriction, including without limitation the rights
13
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
- copies of the Software, and to permit persons to whom the Software is
15
- furnished to do so, subject to the following conditions:
16
-
17
- The above copyright notice and this permission notice shall be included in all
18
- copies or substantial portions of the Software.
19
-
20
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
- SOFTWARE.
27
-
28
- Project-URL: Homepage, https://codeberg.org/eoelab/loaderx
29
- Project-URL: Documentation, https://codeberg.org/eoelab/loaderx
30
- Project-URL: Source, https://codeberg.org/eoelab/loaderx
31
- Project-URL: Bug Tracker, https://codeberg.org/eoelab/loaderx
32
- Keywords: flax,python,dataloader
33
- Classifier: Development Status :: 3 - Alpha
34
- Classifier: Intended Audience :: Developers
35
- Classifier: License :: OSI Approved :: MIT License
36
- Classifier: Programming Language :: Python :: 3
37
- Requires-Python: ==3.13
38
- Description-Content-Type: text/markdown
39
- License-File: LICENSE
40
- Requires-Dist: numpy
41
- Requires-Dist: array-record
42
- Requires-Dist: tqdm
43
- Dynamic: license-file
44
-
45
- # loaderx
46
- A record-based data runtime, focused on delivering extreme throughput and low latency
47
-
48
- **Only Python3.13_Linux_amd64**
49
-
50
- ## Sampler
51
- a high-performance sampler implemented in Zig
52
- ### Index Generator
53
- 1. Sequential generation: indices are produced by traversing the index space in order.
54
- * Sliding traversal: indices are obtained using a fixed-size sliding window. Note that in this case, the index space is treated as a circular queue to avoid truncation at the tail.
55
- 2. Random generation: indices are sampled randomly from the index space.
56
- * Global random: a set of samples is drawn randomly from the entire index space.
57
-
58
- ## zrecord
59
- 基于Zig实现一个并发友好、实现更简单、单机性能更优的 record 存储系统
60
-
61
- 1. ZRecord基于 record 思想设计,我们不假定record间存在顺序关系,对于顺序语义应在record内部表达或建立timestamp
62
- 2. record是相互独立的逻辑记录,仅在 Python 侧提供与 NumPy 兼容的 array 接口投影
63
- 3. zrecord是无类型的,数据只是Bytes,类型解析由上层完成
64
- 4. chunk只是存储管理单元,提高并发性能,不具备语义
65
-
66
- ### 存储形式
67
- 1. ZRecord 仅包含一个多流数据集,长度为 N(N 个 record)
68
- * 索引空间为 [0, N),所有索引与切块语义等价于 gather 操作
69
- * 多流,类似 `set['A'], set['B']`, 但必须等长
70
- ```
71
- Dataset:
72
- record_id = 0..N-1
73
-
74
- Field A: bytes / tensor / audio / text
75
- Field B: bytes / label / meta
76
- ```
77
- 2. 为降低存储体积,支持粒度为record级的透明压缩,以下是可选的压缩方法
78
- ```
79
- | 方法 | 算法 |
80
- | ----- | ------- |
81
- | raw | 无压缩 |
82
- | flate | Deflate |
83
- ```
84
- 3. 通过offset支持高效随机访问
85
- * offset表项是一个三元组 [[chunk_id : u32, offset : u64, physical_length : u32], ...],长度为N。将外部索引空间(outside[0, N])指向一个实际存储地址
86
- * chunk_id为u32,表示具体chunk
87
- * offset为u64, 表示具体chunk内的偏移
88
- * physical_length为u32,表示record数据大小
89
- * offset表每项的字节大小最好为 8 的整数倍,提升对齐和缓存效率
90
- * offset表是mmap访问模式,由于等长,直接将 idx 转换为 ptr + 16*idx
91
- * 不同位宽表示的数值范围
92
- ```
93
- u16 → 2^16 - 1 (65535)
94
- u32 → 2^32 - 1 (4.29e9)
95
- u64 → 2^64 - 1 (1.84e19)
96
- ```
97
- * 不同位宽表示的数据大小
98
- ```
99
- u16 → 64 KiB
100
- u32 → 4 GiB
101
- u64 → 16 EiB
102
- ```
103
- ```
104
- outside idx (global)
105
-
106
- ↓offset[idx]
107
-
108
- chunk_id, offset, physical_length
109
- ```
110
- 4. ZRecord的存储是元数据+分块的格式
111
- 1. 元数据
112
- * meta.json:描述全局元数据,包括全局参数与Field参数
113
- ```
114
- {
115
- "length": 65536,
116
- "chunks": { "num": 12, "size": 8192, "unfull": {[8, write_pos, record_count], ... } },
117
- "fields": [ {"name": "A", "dtype": "f32", "compress": "flate"}, {"name": "B", "dtype": "i32", "compress": "raw"} ]
118
- }
119
- ```
120
- * Field_offset.zr:索引表
121
- ```
122
- 0 0 255
123
- 0 256 511
124
- ...
125
- ```
126
- 2. 分块数据
127
- * chunk/x.zr: 分块的chunk数据,chunk不区分数据所属
128
-
129
- 存储格式如下
130
- ```
131
- dataset
132
- ├── meta.json
133
- ├── A.offset
134
- ├── B.offset
135
- ├── chunk
136
- │ ├── 0.zr
137
- └──└── 1.zr
138
- ```
139
-
140
- ### 执行器
141
-
142
- 负责维护任务队列,工作线程,handle
143
-
144
- 1. 任务队列:
145
- * write_task:写任务队列,压入[ops, record],等待writer弹出
146
- * read_task:读任务队列,预分配batch并压入[batch, pos, chunk_id, offset, physical_length],等待reader弹出
147
- * batch: 具体读取请求下创建的与indices等长的二元组
148
- * pos:在batch中的位置,保证等序返回
149
- * gc_task:垃圾回收任务队列,压入batch,等待cleaner弹出
150
-
151
- 2. 工作线程:执行具体任务的工作线程
152
- * writer:线程池执行写任务,但每个线程固定对应一个chunk
153
- * writer只会向自己的chunk写入(写满时新建chunk提交给执行器完成)
154
- * num_writer推荐:1-4
155
- * reader:线程池执行读任务
156
- * num_reader推荐:8-32
157
- * cleaner:线程池执行垃圾回收任务
158
- * num_cleaner推荐:2-4
159
-
160
- 3. handle:
161
- 1. mmap:维护全局offset、chunk文件的handle
162
- * 初始化阶段注册,关闭前释放
163
- * 对于full的chunk,使用mmap+close(只读),对于unfull的chunk使用mmap(读写)。一旦写满(writer提交),新建chunk并返回给writer,原mmap close。(避免过多占用fd)
164
- * 对于chunk/offset文件,动态完成分块增长,基于ftruncate+mremap实现
165
- * 扩容offset表,一次增长 1M项(16MiB)
166
- * 扩容chunk文件,一次增长 1GiB
167
- 2. batch:维护每次返回batch的handle,也就是[ptr, logical_length]中ptr的合规性
168
- * 工作线程完成后注册
169
- * 必须调用函数来手动释放,释放时batch压入GC队列
170
-
171
- #### 在线任务
172
-
173
- 1. 写任务:写入只允许chunk-level append-only,其余操作基于offset重定向,缺失的field构造为全为0的项(physical_length为0代表无)
174
- 1. 追加:在chunk当前写入位置追加一个record,offset表更新length+1条目,全局长度增加
175
- * 写满:写入完成时检查,如果record_count达到chunk_size, 新建chunk
176
- * 写入越界:捕捉错误,并对chunk扩容
177
- * offset[ length+1 ]不存在:捕捉错误,并对offset扩容
178
- 2. 修改:追加一个record,并修改offset项
179
- 3. 删除:删除offset项(将目标项替换为最后一项),length减一(越界访问属于UB行为)
180
-
181
- 2. 读任务:用一个indices(索引数组),从原数组里访问指定位置的元素(gather)
182
- 1. 对于无压缩的数据,将直接返回文件内地址实现zero-copy
183
-
184
- 读取流程
185
- ```
186
- [idx]
187
-
188
- [batch, pos, idx]
189
- ↓field_offset
190
- [batch, pos, chunk_id, offset, physical_length]
191
-
192
- [ptr, logical_length]
193
- ```
194
-
195
- 3. 垃圾回收任务:cleaner从GC队列中释放对应batch
196
- 1. 对于无压缩的数据,不释放内存,只释放 batch 本身
197
- 2. 对于有压缩的数据,释放内存
198
- 3. GC后访问ptr属于UB行为
199
-
200
- #### 离线任务
201
-
202
- 1. 重平衡(Rebalance):zrecord经过大量写操作后,会出现chunk访问不平衡,通过重平衡来恢复chunk的可访问性
203
- * 实现:根据offset进行重写入,并在完成后替换
204
- * Rebalance只允许手动运行,确保用户知情,提供一个利用率(sum(length)/sum(chunk_size*num_chunks))
205
-
206
- ### 分布式扩展
207
-
208
- **该部分暂时不会实现,只是提前架构**
209
-
210
- 1. 分布式下,我们会得到一个更高的层次-cluster,单机将变为cluster内的node
211
- 2. node持有一个shard,包含若干个chunk
212
- 3. 添加一个indirection表,将全局路径映射到具体node的表,indirection也可以使用hash实现,从而无锁,索引路径变为
213
- ```
214
- outside idx (global)
215
-
216
- ↓indirection[idx]
217
-
218
- node, idx
219
-
220
- ↓offset[idx]
221
-
222
- chunk_id, offset, physical_length
223
- ```
224
-
225
- ## Convert a NumPy tensor to Array_record
226
-
227
- *This will create a directory containing file shards, which helps improve I/O performance.*
228
-
229
- ```
230
- import numpy as np
231
- from loaderx import converter
232
-
233
- train_data = np.load('train_data.npy',mmap_mode='r')
234
- converter(train_data, 'train_data')
235
- ```
236
-
237
- ## Current Limitations
238
- Currently, loaderx only supports single-host environments and does not yet support multi-host training.
239
-
240
- ## Quick Start
241
- ```
242
- import numpy as np
243
- from loaderx import NPDataset, ARDataset, DataLoader
244
-
245
- dataset = ARDataset('train_data')
246
- labelset = NPDataset('xsub/train_label.npy')
247
-
248
- loader = DataLoader(dataset, labelset)
249
-
250
- for i, batch in enumerate(loader):
251
- if i >= 256:
252
- break
253
-
254
- print(batch['data'].shape)
255
- print(batch['label'].shape)
256
- ```
257
-
258
- ### Integrating with JAX/Flax
259
-
260
- For practical integration examples, please refer to the **[Data2Latent](https://codeberg.org/eoelab/Data2Latent)** repository
@@ -1,3 +0,0 @@
1
- numpy
2
- array-record
3
- tqdm
File without changes
File without changes