5.5 Datos a escala: leer desde almacenamiento de objetos en la nube

Todos los conjuntos de datos de este libro hasta ahora cabían en memoria. Este cuaderno trabaja contra uno que no cabe: el análisis MUR SST — temperatura superficial del mar global diaria a 0.01° (~1 km) de resolución, 2002–2020, publicado como un almacén Zarr en el registro de datos abiertos de AWS (s3://mur-sst/zarr-v1, región us-west-2, acceso anónimo). Decodificado, el almacén describe unos 106 TiB de arreglos. Vamos a abrirlo, tomar un subconjunto y calcular sobre él desde una laptop, con una pregunta siempre al frente: ¿cuántos bytes se mueven en realidad?

Esta es la versión ejecutable del patrón de «leer en el lugar» de la sección 5.4.

Crédito de los datos: JPL MUR MEaSUREs Project (2015), GHRSST Level 4 MUR Global Foundation Sea Surface Temperature Analysis v4.1, NASA PO.DAAC, NASA/JPL (2015); Chin, Vazquez-Cuervo & Armstrong (2017), Remote Sensing of Environment 200.

🖥️ Diapositivas — Sesión 27 (vie 4 dic)

Alcanzar el almacén — y fallar rápido si no se puede¶

Un objeto pequeño, .zmetadata, contiene los metadatos consolidados del almacén: la forma, el tipo de dato, el fragmentado y la compresión de cada arreglo, en una sola lectura de ~10 KiB. Lo revisamos primero con tiempos de espera cortos, para que una red bloqueada produzca un error claro en segundos en vez de un cuelgue silencioso de varios minutos.

import time

import numpy as np
import s3fs
import xarray as xr

STORE = "mur-sst/zarr-v1"
S3_OPTS = {
    "anon": True,  # public bucket, no credentials
    "config_kwargs": {"connect_timeout": 10, "read_timeout": 30, "retries": {"max_attempts": 2}},
}

fs = s3fs.S3FileSystem(**S3_OPTS)
try:
    meta = fs.info(f"{STORE}/.zmetadata")
except Exception as err:
    raise RuntimeError(
        f"Cannot reach s3://{STORE} (anonymous, region us-west-2). "
        "Check your connection. On restricted networks, outbound anonymous S3 is often "
        "blocked entirely -- see the admonition at the end of this notebook and the "
        "'Restricted environments' section of 5.4."
    ) from err


def fmt(nbytes):
    """Human-readable byte count."""
    for unit in ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]:
        if nbytes < 1024:
            return f"{nbytes:,.1f} {unit}"
        nbytes /= 1024


print(f"store reachable; consolidated metadata: {fmt(meta['size'])} -- one GET describes everything below")

store reachable; consolidated metadata: 10.3 KiB -- one GET describes everything below

Abrir 106 TiB de forma perezosa¶

xr.open_dataset(..., engine="zarr", chunks={}) lee únicamente ese objeto de metadatos. chunks={} significa «dame un arreglo dask cuyos fragmentos coincidan con los fragmentos en disco» — ningún valor de datos cruza la red hasta que se lo pidamos.

t0 = time.time()
ds = xr.open_dataset(
    f"s3://{STORE}",
    engine="zarr",
    chunks={},  # lazy: dask chunks = storage chunks
    backend_kwargs={"storage_options": S3_OPTS},
)
print(f"opened in {time.time() - t0:.1f} s: {fmt(ds.nbytes)} of arrays described, ~{fmt(meta['size'])} actually read")
ds

opened in 12.9 s: 106.3 TiB of arrays described, ~10.3 KiB actually read

Los fragmentos: la unidad de todo¶

Un fragmento (chunk) es la unidad atómica de E/S y de compresión del almacén: cada lectura trae y descomprime fragmentos completos, nunca valores sueltos, así que su patrón de acceso paga en monedas del tamaño del fragmento. Zarr v3 agrega los shards (agrupaciones de almacenamiento) — muchos fragmentos pequeños empacados en un solo objeto de almacenamiento, para que un almacén de objetos no se ahogue en millones de archivos diminutos; este almacén es anterior a ellos (disposición de Zarr v2), así que cada fragmento es un objeto de S3 y la unidad de transferencia por red.

analysed_sst se almacena como int16 (un par escala/desplazamiento reconstruye los kelvin) en fragmentos de 5 días × 1799 lat × 3600 lon. Pesemos uno.

sst = ds["analysed_sst"]
storage_chunks = sst.encoding["chunks"]
chunk_of = dict(zip(sst.dims, storage_chunks))
print("shape:", sst.shape, "-> storage chunks:", storage_chunks)
print("stored dtype:", sst.encoding["dtype"], "| decoded dtype:", sst.dtype)

nchunks = tuple(int(np.ceil(s / c)) for s, c in zip(sst.shape, storage_chunks))
print("chunk grid:", nchunks, "=", f"{int(np.prod(nchunks)):,}", "chunk objects for this variable")

one = fs.info(f"{STORE}/analysed_sst/1200.7.1")  # one arbitrary chunk object
decoded_chunk = int(np.prod(storage_chunks)) * sst.dtype.itemsize
print(f"one chunk: {fmt(one['size'])} compressed on S3 -> {fmt(decoded_chunk)} decoded in RAM")
print(f"variable:  ~{fmt(one['size'] * np.prod(nchunks))} compressed (est. from that chunk) -> {fmt(sst.nbytes)} decoded")

# Rechunk on read: dask chunks must be multiples of the storage chunks, or every
# dask task re-reads partial storage chunks. This merges pairs of 5-day chunks:
sst10 = sst.chunk({"time": 10})  # equivalently: chunks={"time": 10, ...} in open_dataset
print("dask chunks after rechunk-on-read:", sst10.data.chunksize)

shape: (6443, 17999, 36000) -> storage chunks: (5, 1799, 3600)
stored dtype: int16 | decoded dtype: float64
chunk grid: (1289, 11, 10) = 141,790 chunk objects for this variable
one chunk: 20.0 MiB compressed on S3 -> 247.1 MiB decoded in RAM
variable:  ~2.7 TiB compressed (est. from that chunk) -> 30.4 TiB decoded

dask chunks after rechunk-on-read: (10, 1799, 3600)

Tomar el subconjunto de forma perezosa y luego calcular¶

Diez días sobre la plataforma del Pacífico Noroeste: 46–49°N, 126–122°O, junio de 2019. Rebanar un arreglo perezoso edita el grafo de tareas — sigue moviendo cero bytes.

TIME = slice("2019-06-01", "2019-06-10")
LAT, LON = slice(46, 49), slice(-126, -122)

sub = sst.sel(time=TIME, lat=LAT, lon=LON)
print(f"subset {sub.shape}: {fmt(sub.data.nbytes)} across {sub.data.npartitions} chunks -- 0 bytes moved so far")

subset (10, 301, 401): 9.2 MiB across 3 chunks -- 0 bytes moved so far

.compute() es donde llega la cuenta de la red. Importan dos conteos de bytes distintos, y medimos ambos:

bytes en RAM — el nbytes de dask sobre el resultado calculado: lo que su respuesta cuesta en memoria;
bytes movidos — la red transfiere objetos de fragmento comprimidos completos, así que enumeramos exactamente qué objetos de fragmento toca el subconjunto y sumamos sus tamaños en S3.

t0 = time.time()
box = sub.compute()
elapsed = time.time() - t0

# Which chunk objects did that read touch? Convert the label slices to integer
# index ranges, divide by the chunk size along each dimension, and list the keys.
ranges = {}
for dim, sel in {"time": TIME, "lat": LAT, "lon": LON}.items():
    start, stop = sst.get_index(dim).slice_locs(sel.start, sel.stop)
    ranges[dim] = range(start // chunk_of[dim], (stop - 1) // chunk_of[dim] + 1)
keys = [
    f"{STORE}/analysed_sst/{t}.{y}.{x}"
    for t in ranges["time"] for y in ranges["lat"] for x in ranges["lon"]
]
moved = sum(fs.info(k)["size"] for k in keys)

print(f"chunk objects fetched : {len(keys)} (= {sub.data.npartitions} graph partitions)")
print(f"bytes moved           : {fmt(moved)} compressed, in {elapsed:.1f} s")
print(f"bytes in RAM (result) : {fmt(box.nbytes)}")
print(f"bytes stored (decoded): {fmt(sst.nbytes)}")
print(f"moved : stored        = 1 : {sst.nbytes / moved:,.0f}")
print(f"mean SST in the box   : {float(box.mean()):.2f} K")

chunk objects fetched : 3 (= 3 graph partitions)
bytes moved           : 60.6 MiB compressed, in 20.4 s
bytes in RAM (result) : 9.2 MiB
bytes stored (decoded): 30.4 TiB
moved : stored        = 1 : 525,483
mean SST in the box   : 285.94 K

Esa razón es el argumento completo a favor de los formatos optimizados para la nube. El almacén contiene 30.4 TiB (decodificados) de analysed_sst; responder nuestra pregunta movió 60.6 MiB — una parte en ~525 000 — porque el fragmentado abarata las lecturas parciales y la pereza implica que solo pedimos los fragmentos que toca nuestra rebanada. Los 60.6 MiB son además ~6× más grandes que los 9.2 MiB que aterrizaron en RAM: la granularidad del fragmento es el impuesto que se paga por el almacenamiento de objetos, y la razón por la que la forma del fragmento debería parecerse a su patrón de acceso.

El momento en que muere el trabajo en memoria¶

Haga la aritmética antes de llamar .compute() sobre algo grande. Un solo paso temporal global de analysed_sst se decodifica a 17 999 × 36 000 × 8 bytes ≈ 4.8 GiB. Tres días agotan una laptop de 16 GiB. La variable completa son 30.4 TiB — unas dos mil laptops de RAM — y sst.values (o .compute() sobre el arreglo sin subconjuntar) intentaría alegremente materializarla. No vamos a correr eso, y usted tampoco debería: la solución nunca es un .compute() más grande, es un algoritmo que retenga solo unos pocos fragmentos a la vez.

Transmisión por flujo: la agregación que sí funciona¶

Una media mensual sobre la misma caja necesita 30 días globales — ~145 GiB decodificados si cargara pasos temporales enteros, pero aun así solo 7 objetos de fragmento para nuestra rebanada. mean("time") sobre el arreglo perezoso construye un grafo de medias parciales por fragmento que dask combina: cada trabajador trae un objeto de fragmento de 20 MiB, lo descomprime (62 MiB de int16, 247 MiB una vez decodificado a float64), toma su media parcial sobre la rebanada y lo descarta. La memoria máxima se queda en un puñado de fragmentos por más largo que sea el eje temporal.

june = sst.sel(time=slice("2019-06-01", "2019-06-30"), lat=LAT, lon=LON)
june_mean = june.mean("time")  # lazy: per-chunk partial means, combined at the end

print(f"graph input : {fmt(june.data.nbytes)} across {june.data.npartitions} chunks")
print(f"graph output: {fmt(june_mean.data.nbytes)}")

t0 = time.time()
june_map = june_mean.compute()
print(f"streamed in {time.time() - t0:.1f} s; peak memory ~ a few chunks, not {fmt(sst.nbytes)}")

graph input : 27.6 MiB across 7 chunks
graph output: 943.0 KiB

streamed in 40.8 s; peak memory ~ a few chunks, not 30.4 TiB

import matplotlib.pyplot as plt

june_map.plot(figsize=(7, 4), cmap="viridis", cbar_kwargs={"label": "mean SST (K)"})
plt.title("MUR SST, June 2019 mean -- Pacific Northwest coast")
plt.tight_layout()
plt.show()

Las lecturas anónimas de S3 fallan en redes restringidas

En muchas redes institucionales — laboratorios nacionales, enclaves de agencias, algunos nodos de cómputo de HPC universitario — el acceso anónimo de salida a S3 está bloqueado. El síntoma rara vez es un error limpio: la apertura se cuelga lo que dure el tiempo de espera y luego lanza ConnectTimeoutError/EndpointConnectionError, o un proxy responde con un 403 engañoso. (La primera celda de este cuaderno usa tiempos de espera de 10 segundos y un mensaje de error en lenguaje llano precisamente para que la falla sea rápida y legible.) La alternativa no es pelear con el cortafuegos: use un espejo institucional del archivo si existe, o prepare de antemano un subconjunto — corra las celdas de subconjuntado en una máquina conectada, guarde el resultado con box.to_netcdf(...) (9.2 MiB, no 30.4 TiB), muévalo por su puerta de transferencia aprobada y fije su hash con pooch (sección 1.6). El lado del entorno del mismo problema se cubre en 5.4, «Entornos restringidos».

Laptop → HPC → nube: la decisión en números¶

Este cuaderno respondió preguntas sobre una variable de 30.4 TiB moviendo ~200 MiB en total — holgadamente un trabajo de laptop, porque los fragmentos que tocamos caben en nuestro ancho de banda y nuestra RAM. Esa es la regla general:

Los fragmentos que toca ≪ su RAM y su paciencia → laptop, exactamente como aquí. Apertura perezosa, subconjunto, flujo.
Los fragmentos que toca ≈ el almacén completo (una climatología global de todo el registro; entrenamiento de aprendizaje automático sobre cada fragmento) → lleve el cómputo a los datos, no los datos al cómputo: una instancia en us-west-2 lee este bucket a varios GB/s sin cargo de salida, y un clúster de HPC funciona cuando los datos (o un espejo) ya están en su sistema de archivos.
En cualquier caso, la consulta viaja hacia los datos y solo las respuestas viajan de vuelta — el patrón de «leer en el lugar» de la sección 5.4, que además cubre la escalera de cómputo en sí: clúster departamental, asignaciones nacionales (NSF ACCESS en EE. UU.; LANCAD, SNCAD, NLHPC o la RES en el ámbito iberoamericano), nube comercial, y la disciplina de costos que hace sobrevivible a esta última.

References¶

NASA/JPL. (2015). GHRSST Level 4 MUR Global Foundation Sea Surface Temperature Analysis (v4.1). NASA Physical Oceanography Distributed Active Archive Center. 10.5067/GHGMR-4FJ04