# CRA5Prepare

**Repository Path**: AI4EarthLab/CRA5Prepare

## Basic Information

- **Project Name**: CRA5Prepare
- **Description**: No description available
- **Primary Language**: Python
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-07-03
- **Last Updated**: 2026-07-03

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# CRA5Prepare

把 CRA5 压缩的 `.bin` 文件批量解码为 latent representation（`.npy`），方便组内同学直接 `np.load()` 使用。

## 输出格式

- 路径：`latents/<YYYY>/<YYYY-MM-DDTHH:MM:SS>.npy`
- dtype：`float32`
- shape：`(256, 72, 144)`
- 单文件大小：约 **10.1 MB**
- 全量（2001–2023，6-hourly，约 33,608 个时间戳）：约 **340 GB**

时间锚点为 6-hourly UTC（00/06/12/18），可调整 `bin2latent.py --cadence-hours`。

## 下游加载示例

```python
import numpy as np
z = np.load("latents/2001/2001-01-01T00:00:00.npy")  # (256, 72, 144) float32
assert np.isfinite(z).all()
```

## 环境与依赖

- **Conda 环境**：`xichen_py39`（Python 3.9.23 + PyTorch 2.1.0 + CUDA 11.8）
  - 环境定义文件：`./xichen_py39.yml`
  - 缺省 prefix：`/fs1/home/yangjh15/miniconda3/envs/xichen_py39`
  - 其他用户可通过 `--export=CONDA_ROOT=/your/conda/prefix` 覆盖
- **CRA5 子模块**：作为 git submodule 在 `./CRA5/` 下，不走 `pip install`（避免污染 conda 环境）
- **CRA5 ckpt**：默认读取 `/fs1/home/yangjinhui/public_data/CRA5/models/cra5_268v_300k.pth`（HPC4 共享路径）；可通过 `bin2latent.py --ckpt-path /your/path.pth` 覆盖
- **flash-attn**（可选加速 attention）：见下方 "flash-attn 手动安装" 段。预编译 wheel 在 GitHub releases，需手动 `wget` + `pip install` 本地 wheel。
- **compressai**（CRA5 间接依赖，必须）：`pip install compressai`（环境不自带；CRA5 vendored entropy_models 依赖外部 `compressai._CXX` C++ 扩展和 `compressai.ops/ans` 等高层 API）

### flash-attn 手动安装（推荐）

flash-attn 的 PyPI sdist 需要本地编译（10+ 分钟），预编译 wheel 在 [Dao-AILab/flash-attention GitHub releases](https://github.com/Dao-AILab/flash-attention/releases/tag/v2.5.9.post1)。步骤：

```bash
# 1. 确认 PyTorch 的 cxx11 ABI（conda 默认 PyTorch 通常是 False）
source /fs1/home/yangjh15/miniconda3/etc/profile.d/conda.sh
conda activate xichen_py39
python -c "import torch; print(torch._C._GLIBCXX_USE_CXX11_ABI)"
# 输出 True 选 cxx11abiTRUE wheel；输出 False 选 cxx11abiFALSE wheel

# 2. 下载 wheel
mkdir -p ~/wheels && cd ~/wheels
# ABI=False（conda PyTorch 大多数情况）：
wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.9.post1/flash_attn-2.5.9.post1%2Bcu118torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl
# ABI=True 时改下载 TRUE 变体：
# wget https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.9.post1/flash_attn-2.5.9.post1%2Bcu118torch2.1cxx11abiTRUE-cp39-cp39-linux_x86_64.whl

# 3. pip install 本地 wheel
pip install ~/wheels/flash_attn-2.5.9.post1+cu118torch2.1cxx11abiFALSE-cp39-cp39-linux_x86_64.whl

# 4. 验证
python -c "import flash_attn; print(flash_attn.__version__)"
```

ABI 选错会报 `undefined symbol` 错误——换另一个 ABI 即可。flash-attn ≥ 2.7 要求 torch ≥ 2.2 不兼容；2.5.9.post1 是兼容 torch 2.1 的最新稳定版。

## 一次性 Bootstrap（HPC4 登录节点）

```bash
git clone <repo-url> && cd CRA5Prepare
git submodule update --init --recursive

# 激活项目 conda 环境（CONDA_ROOT 可覆盖）
: "${CONDA_ROOT:=/fs1/home/yangjh15/miniconda3}"
source "$CONDA_ROOT/etc/profile.d/conda.sh"
conda activate xichen_py39

# 必须：装 CRA5 间接依赖 compressai（CRA5 vendored entropy_models 用到 C++ RANS 扩展）
pip install compressai

# 可选：手动装 flash-attn（推荐 wget wheel 而非 pip install，避免本地编译 10+ 分钟）
# 见 README "flash-attn 手动安装" 段；CUDA 11.8 + PyTorch 2.1 + Python 3.9 wheel 在 GitHub releases。

# 确认 CRA5 ckpt 可读
ls -la /fs1/home/yangjinhui/public_data/CRA5/models/cra5_268v_300k.pth
```

## 单 bin 烟雾测试

```bash
bash verify_single.sh <path/to/some.bin> /tmp/test.npy
```

预期输出：
```
[verify] OK  shape=(256, 72, 144) dtype=float32 min=... max=... mean=... std=...
```

## 全量提交（HPC4 SLURM）

```bash
BIN_ROOT=/path/to/CRA5 \
OUT_ROOT=/path/to/latents \
./submit_all.sh
```

默认参数：
- partition：`gpu`（A100×2）；可用 `PARTITION=gpu4` 改成 8 卡节点
- 每个 job 跑一年（`--array=0-22`），23 个 task 并行
- 时间预算：单年 ≤ 2 小时（`--time=02:00:00`）
- 资源：`--gres=gpu:1 --cpus-per-task=4 --mem=16G`

### 监控

```bash
squeue -u $USER            # 列出当前用户的作业
sacct -j <jobid>           # 历史作业查询
tail -f logs/<jobname>_<jobid>_<taskid>.log
```

### 中断与续跑

```bash
scancel <jobid>                            # 取消整个 job array
scancel <jobid>_<taskid>                   # 取消单个 task
# 续跑：`bin2latent.py --resume` 会跳过已存在且 mtime 较新的 .npy
sbatch --export=BIN_ROOT=$BIN_ROOT,OUT_ROOT=$OUT_ROOT bin2latent.sbatch
```

## 文件清单

```
CRA5Prepare/
├── .gitmodules                # CRA5 子模块配置（HTTPS 公开 URL）
├── .gitignore                 # 排除 latents/、logs/、*.npy 等
├── README.md                  # 本文件
├── xichen_py39.yml            # conda 环境定义
├── CRA5/                      # git submodule（codelib, 不修改）
├── bin2latent.py              # ⭐ 核心解码脚本（单文件，~180 行）
├── bin2latent.sbatch          # SLURM 提交模板（--array=0-22，按年并行）
├── submit_all.sh              # 一键提交全量作业
└── verify_single.sh           # 单 bin 烟雾测试
```

## 实现要点

- **不 `pip install CRA5`**：`bin2latent.py` 顶部 `sys.path.insert(0, ./CRA5)` 把子模块目录当包 import，避免污染 conda 环境。
- **直接调用 `cra5_api.bin_to_latent()`**：CRA5 自带的解码入口完整覆盖 bin 头解析 + 熵解码；不做包装类。
- **本地 ckpt**：`bin2latent.py` 通过 monkey-patch `cra5.models.compressai.zoo.image.model_urls` 把默认 S3 URL 替换为 `file://` 本地路径。
- **Atomic write**：先写 `.npy.tmp` 再 rename，避免被中断留下损坏文件。
- **幂等**：`--resume` 跳过 `dst.exists() and dst.mtime >= src.mtime` 的目标；中断后可重跑。
- **退出码**：0 = 全成功；2 = 部分失败（可 `--resume` 重跑）；3 = `--strict` 模式任一失败。

## 性能

- 单年（1460 个 bin）约 25 min（1× A100）
- 全量（23 年，8-GPU 节点并行）约 3–4 小时

## 不在本项目范围

- ❌ decode → 反归一化 → ERA5 NC 重建路径（项目只产出 latent）
- ❌ CRA5 模型重训练/微调
- ❌ 单元测试框架

## 故障排查

| 症状 | 排查 |
|---|---|
| `ModuleNotFoundError: cra5` | `git submodule status` 应显示 CRA5 已 clone；否则 `git submodule update --init --recursive` |
| `FileNotFoundError: cra5_268v_300k.pth` | 确认 ckpt 在 `/fs1/home/yangjinhui/public_data/CRA5/models/` 或用 `--ckpt-path` 指定其他位置 |
| `RuntimeError: CUDA out of memory` | 单 bin latent 仅 10 MB，理论上不可能 OOM；检查 `--gres=gpu:1` 是否生效 |
| Latent 全为 0 / NaN | 多数情况是 ckpt 加载失败；`print(self.net)` 在 cra5_api.__init__ 末尾调试 |
| job OOM kill (`--mem=16G`) | 单 task 内存 < 4 GB；如频繁被 kill 提升 `--mem` 到 32G |