fix: 修复预处理流水线 4 个致命问题
- pipeline: 单条数据异常不再崩溃整条流水线, 改 log warning 后跳过 - pipeline: _align_bucket 统一用 len(ids) 填充, 修复多输出模式下长度错配 - writer: BinWriter/H5Writer 写入失败自动清理残留文件并记录详细错误 - packing: BFDPacking 真正将序列打包进 bin 而非仅重排, 减少碎片
This commit is contained in:
@@ -6,7 +6,9 @@ List[Tensor]}`` dict and delegates the write to the writer selected
|
||||
by ``output.storage_format``.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List
|
||||
|
||||
@@ -15,6 +17,8 @@ import torch
|
||||
from astrai.dataset.storage import save_bin, save_h5
|
||||
from astrai.factory import BaseFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class StoreWriter(ABC):
|
||||
"""Write pre-tokenized tensors to disk in a format-specific way."""
|
||||
@@ -37,11 +41,35 @@ class StoreWriterFactory(BaseFactory["StoreWriter"]):
|
||||
class BinWriter(StoreWriter):
|
||||
def save(self, output_dir, domain, shard_idx, tensors):
|
||||
shard_path = os.path.join(output_dir, domain, f"shard_{shard_idx:04d}")
|
||||
save_bin(shard_path, tensors)
|
||||
try:
|
||||
save_bin(shard_path, tensors)
|
||||
except Exception:
|
||||
if os.path.exists(shard_path):
|
||||
shutil.rmtree(shard_path, ignore_errors=True)
|
||||
logger.error(
|
||||
"Failed to write shard %s/%s_%04d, cleaned up partial output",
|
||||
domain,
|
||||
"shard",
|
||||
shard_idx,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
@StoreWriterFactory.register("h5")
|
||||
class H5Writer(StoreWriter):
|
||||
def save(self, output_dir, domain, shard_idx, tensors):
|
||||
chunk_dir = os.path.join(output_dir, domain)
|
||||
save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
|
||||
file_path = os.path.join(chunk_dir, f"data_{shard_idx:04d}.h5")
|
||||
try:
|
||||
save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
|
||||
except Exception:
|
||||
if os.path.exists(file_path):
|
||||
os.remove(file_path)
|
||||
logger.error(
|
||||
"Failed to write shard %s/data_%04d.h5, cleaned up partial output",
|
||||
domain,
|
||||
shard_idx,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
Reference in New Issue
Block a user