(
checker_func: callable,
)
| 178 | |
| 179 | |
| 180 | def run_with_files( |
| 181 | checker_func: callable, |
| 182 | ): |
| 183 | rank = int(os.getenv("RANK")) |
| 184 | ctx = get_context("spawn") |
| 185 | queue = ctx.Queue() |
| 186 | _device_uuid = _get_physical_gpu_id(device_manager, rank) |
| 187 | ps = ParameterServer(auto_pg=True) |
| 188 | _device_uuid = _get_physical_gpu_id(ps.device_manager, rank) |
| 189 | named_tensors = dict(gen_test_tensors(rank)) |
| 190 | |
| 191 | # Save 1/3 tensors to /dev/shm/ as .safetensors files |
| 192 | # Save 1/3 tensors to ./tmp (disk) as .safetensors files |
| 193 | # Keep 1/3 tensors in memory |
| 194 | import safetensors.torch |
| 195 | |
| 196 | files = [] |
| 197 | dev_shm_dir = "/dev/shm/checkpoint_engine_tests" # noqa: S108 |
| 198 | disk_dir = "/tmp/checkpoint_engine_tests" # noqa: S108 |
| 199 | os.makedirs(dev_shm_dir, exist_ok=True) |
| 200 | os.makedirs(disk_dir, exist_ok=True) |
| 201 | tensors_items = list(named_tensors.items()) |
| 202 | tensors_in_dev_shm = dict(tensors_items[: len(tensors_items) // 2]) |
| 203 | tensors_in_disk = dict(tensors_items[len(tensors_items) // 3 : 2 * len(tensors_items) // 3]) |
| 204 | tensors_in_memory = dict(tensors_items[1 * len(tensors_items) // 2 :]) |
| 205 | disk_files = [ |
| 206 | os.path.join(disk_dir, f"rank{_rank}_checkpoint.safetensors") |
| 207 | for _rank in range(get_world_size()) |
| 208 | ] |
| 209 | safetensors.torch.save_file(tensors_in_disk, disk_files[rank]) |
| 210 | time.sleep(1) |
| 211 | files.append(disk_files[rank]) |
| 212 | dev_shm_files = [ |
| 213 | os.path.join(dev_shm_dir, f"rank{rank}_checkpoint.safetensors") |
| 214 | for _ in range(get_world_size()) |
| 215 | ] |
| 216 | safetensors.torch.save_file(tensors_in_dev_shm, dev_shm_files[rank]) |
| 217 | time.sleep(1) |
| 218 | files.append(dev_shm_files[rank]) |
| 219 | |
| 220 | checkpoint_name = "test_with_files" |
| 221 | proc = ctx.Process(target=checker_func, args=(rank, _device_uuid, named_tensors, queue)) |
| 222 | proc.start() |
| 223 | ps.register_checkpoint(checkpoint_name, named_tensors=tensors_in_memory, files=files) |
| 224 | ps.gather_metas(checkpoint_name) |
| 225 | ps.update(checkpoint_name, queue.put, ranks=[]) |
| 226 | # sleep 3s to wait process group is destroyed |
| 227 | time.sleep(3) |
| 228 | ps.unregister_checkpoint(checkpoint_name) |
| 229 | queue.put(None) |
| 230 | proc.join() |
| 231 | if rank == 0: |
| 232 | import shutil |
| 233 | |
| 234 | shutil.rmtree(dev_shm_dir) |
| 235 | shutil.rmtree(disk_dir) |
| 236 | assert proc.exitcode == 0 |
| 237 |
no test coverage detected