如何在使用xarray.to_zarr上传GCP时显示上传进度？

可以通过Dask的诊断工具或tqdm来跟踪大型xarray数据集上传GCP的进度，核心是先将数据集分块，让Dask把上传任务拆分为多个子任务，再用进度工具监控这些子任务的完成情况。以下是具体实现方案：

方法1：使用Dask自带的ProgressBar #

先给数据集添加分块配置，再用ProgressBar上下文管理器包裹任务的执行过程：

import os
import xarray as xr
import numpy as np
import gcsfs
from dask.diagnostics import ProgressBar

if __name__ == '__main__':
    # 配置GCP凭证
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "service-account.json"

    # 创建带分块的xarray数据集
    data_arr = np.random.rand(5000, 100, 100)
    # 设置分块，可根据数据大小调整块尺寸
    data_xarr = xr.DataArray(data_arr, dims=["x", "y", "z"], chunks={"x": 500})
    ds = data_xarr.to_dataset(name="data")

    # 定义GCP存储
    gcp_blob_uri = "gs://gprlib/test.zarr"
    gcs = gcsfs.GCSFileSystem()
    store = gcs.get_mapper(gcp_blob_uri)

    # 生成延迟上传任务并执行，同时显示进度
    write_job = ds.to_zarr(store, mode="w", compute=False)
    with ProgressBar():
        write_job.compute()

说明：

• 给DataArray添加chunks参数后，xarray会自动用Dask处理数据，将大任务拆分为多个分块上传任务。
• ProgressBar上下文管理器会实时显示每个分块的上传进度、速度和剩余时间。

方法2：使用tqdm配合Dask #

如果更喜欢tqdm的进度条样式，可以用dask.diagnostics.tqdm：

import os
import xarray as xr
import numpy as np
import gcsfs
from dask.diagnostics.tqdm import TqdmProgressBar

if __name__ == '__main__':
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "service-account.json"

    data_arr = np.random.rand(5000, 100, 100)
    data_xarr = xr.DataArray(data_arr, dims=["x", "y", "z"], chunks={"x": 500})
    ds = data_xarr.to_dataset(name="data")

    gcp_blob_uri = "gs://gprlib/test.zarr"
    gcs = gcsfs.GCSFileSystem()
    store = gcs.get_mapper(gcp_blob_uri)

    write_job = ds.to_zarr(store, mode="w", compute=False)
    with TqdmProgressBar():
        write_job.compute()

说明：

• TqdmProgressBar会生成动态进度条，支持显示已完成任务数、总任务数和预估剩余时间。

关键注意点 #

• 分块尺寸要合理：块太小会导致任务过多增加开销，块太大则进度更新不频繁，可根据数据集大小和网络情况调整。
• 必须用compute=False生成延迟任务，再用上下文管理器包裹compute()调用，这样进度工具才能捕获任务执行过程。

内容的提问来源于stack exchange，提问作者Henry Ruiz