TensorStore pour un stockage de baie évolutif hautes performances


De nombreuses applications contemporaines passionnantes de l’informatique et de l’apprentissage automatique (ML) manipulent des ensembles de données multidimensionnels qui couvrent un seul grand système de coordonnées, par exemple, modélisation météorologique à partir de mesures atmosphériques sur une grille spatiale ou l’imagerie médicale prédictions à partir des valeurs d’intensité d’image multicanaux dans un balayage 2d ou 3d. Dans ces paramètres, même un seul ensemble de données peut nécessiter des téraoctets ou des pétaoctets de stockage de données. De tels ensembles de données sont également difficiles à utiliser car les utilisateurs peuvent lire et écrire des données à des intervalles irréguliers et à des échelles variables, et sont souvent intéressés par la réalisation d’analyses à l’aide de nombreuses machines fonctionnant en parallèle.

Aujourd’hui, nous vous présentons TensorStoreune bibliothèque logicielle open-source C++ et Python conçue pour le stockage et la manipulation de n-données dimensionnelles qui :

  • Fournit une API uniforme pour la lecture et l’écriture de plusieurs formats de tableau, y compris et N5.
  • Prise en charge native plusieurs systèmes de stockagey compris Stockage en nuage Google, systèmes de fichiers locaux et réseau, Serveurs HTTPet stockage en mémoire.
  • Prend en charge la mise en cache et les transactions en lecture/écriture différée, avec une forte atomicité, isolation, cohérence et durabilité (ACIDE) garanties.
  • Prend en charge un accès sûr et efficace à partir de plusieurs processus et machines via une simultanéité optimiste.
  • Offre une API asynchrone pour permettre un accès à haut débit même au stockage distant à latence élevée.
  • Fournit des fonctionnalités avancées et entièrement composables indexage opérations et vues virtuelles.

TensorStore a déjà été utilisé pour résoudre des défis d’ingénierie clés dans le calcul scientifique (par exemple, la gestion et le traitement de grands ensembles de données en neurosciences, tels que données de microscopie électronique 3d à l’échelle péta et « 4d » vidéos d’activité neuronale). TensorStore a également été utilisé dans la création de modèles d’apprentissage automatique à grande échelle tels que Palmier en abordant le problème de la gestion des paramètres du modèle (points de contrôle) lors de l’apprentissage distribué.

API familière pour l’accès et la manipulation des données
TensorStore fournit une API Python simple pour charger et manipuler des données de tableau volumineux. Dans l’exemple suivant, nous créons un objet TensorStore qui représente un voxel de 56 000 milliards Image 3D d’un cerveau de mouche et accéder à un petit patch 100×100 des données en tant que NumPy déployer:

>>> import tensorstore as ts
>>> import numpy as np

# Create a TensorStore object to work with fly brain data.
>>> dataset = ts.open({
...     'driver':
...         'neuroglancer_precomputed',
...     'kvstore':
...         'gs://neuroglancer-janelia-flyem-hemibrain/v1.1/segmentation/',
... }).result()

# Create a 3-d view (remove singleton 'channel' dimension):
>>> dataset_3d = dataset[ts.d['channel'][0]]
>>> dataset_3d.domain
{ "x": [0, 34432), "y": [0, 39552), "z": [0, 41408) }

# Convert a 100x100x1 slice of the data to a numpy ndarray
>>> slice = np.array(dataset_3d[15000:15100, 15000:15100, 20000])

Surtout, aucune donnée réelle n’est consultée ou stockée en mémoire tant que la tranche 100×100 spécifique n’est pas demandée ; par conséquent, des ensembles de données sous-jacents arbitrairement volumineux peuvent être chargés et manipulés sans avoir à stocker l’intégralité de l’ensemble de données en mémoire, en utilisant une syntaxe d’indexation et de manipulation largement identique aux opérations NumPy standard. TensorStore fournit également un support étendu pour fonctionnalités d’indexation avancéesy compris les transformations, alignement, diffusionet des vues virtuelles (conversion de type de données, sous-échantillonnage, tableaux générés paresseusement à la volée).

L’exemple suivant montre comment TensorStore peut être utilisé pour créer un tableau zarr et comment son API asynchrone permet un débit plus élevé :

>>> import tensorstore as ts
>>> import numpy as np

>>> # Create a zarr array on the local filesystem
>>> dataset = ts.open({
...     'driver': 'zarr',
...     'kvstore': 'file:///tmp/my_dataset/',
... },
... dtype=ts.uint32,
... chunk_layout=ts.ChunkLayout(chunk_shape=[256, 256, 1]),
... create=True,
... shape=[5000, 6000, 7000]).result()

>>> # Create two numpy arrays with example data to write.
>>> a = np.arange(100*200*300, dtype=np.uint32).reshape((100, 200, 300))
>>> b = np.arange(200*300*400, dtype=np.uint32).reshape((200, 300, 400))

>>> # Initiate two asynchronous writes, to be performed concurrently.
>>> future_a = dataset[1000:1100, 2000:2200, 3000:3300].write(a)
>>> future_b = dataset[3000:3200, 4000:4300, 5000:5400].write(b)

>>> # Wait for the asynchronous writes to complete
>>> future_a.result()
>>> future_b.result()

Mise à l’échelle sûre et performante
Le traitement et l’analyse de grands ensembles de données numériques nécessitent des ressources de calcul importantes. Ceci est généralement réalisé grâce à la parallélisation sur de nombreux cœurs de processeur ou d’accélérateur répartis sur de nombreuses machines. Par conséquent, un objectif fondamental de TensorStore a été de permettre un traitement parallèle d’ensembles de données individuels qui soit à la fois sûr (c’est-à-dire, évite la corruption ou les incohérences résultant de modèles d’accès parallèles) et hautes performances (c’est-à-dire que la lecture et l’écriture sur TensorStore ne sont pas un goulot d’étranglement pendant le calcul). . En fait, lors d’un test dans les centres de données de Google, nous avons constaté une mise à l’échelle presque linéaire des performances de lecture et d’écriture à mesure que le nombre de processeurs augmentait :

Performances en lecture et en écriture pour un ensemble de données TensorStore dans format zarr résidant sur Google Cloud Storage (GCS) accessible simultanément à l’aide d’un nombre variable de tâches de calcul monocœur dans les centres de données Google. Les performances de lecture et d’écriture évoluent presque linéairement avec le nombre de tâches de calcul.

Les performances sont obtenues en implémentant des opérations de base en C++, en utilisant largement le multithreading pour des opérations telles que l’encodage/décodage et les E/S réseau, et en partitionnant de grands ensembles de données en unités beaucoup plus petites grâce à la segmentation pour permettre une lecture et une écriture efficaces de sous-ensembles de l’ensemble de données. TensorStore fournit également une mise en cache en mémoire configurable (qui réduit les interactions plus lentes du système de stockage pour les données fréquemment consultées) et une API asynchrone qui permet à une opération de lecture ou d’écriture de se poursuivre en arrière-plan pendant qu’un programme effectue d’autres travaux.

La sécurité des opérations parallèles lorsque de nombreuses machines accèdent au même ensemble de données est obtenue grâce à l’utilisation de simultanéité optimistequi maintient la compatibilité avec diverses couches de stockage sous-jacentes (y compris les plates-formes de stockage Cloud, telles que CGV, ainsi que les systèmes de fichiers locaux) sans impact significatif sur les performances. TensorStore fournit également de solides garanties ACID pour toutes les opérations individuelles exécutées au cours d’une seule exécution.

Pour rendre l’informatique distribuée avec TensorStore compatible avec de nombreux workflows de traitement de données existants, nous avons également intégré TensorStore à des bibliothèques de calcul parallèle telles que Faisceau Apache (code d’exemple) et Tableau de bord (code d’exemple).

Cas d’utilisation : modèles de langage
Un développement récent passionnant dans ML est l’émergence de modèles de langage plus avancés tels que Palmier. Ces réseaux de neurones contiennent des centaines de milliards de paramètres et présentent des capacités surprenantes dans la compréhension et la génération du langage naturel. Ces modèles repoussent également les limites de l’infrastructure informatique ; en particulier, la formation d’un modèle de langage tel que PaLM nécessite des milliers de TPU fonctionnant en parallèle.

Un défi qui se pose au cours de ce processus de formation est de lire et d’écrire efficacement les paramètres du modèle. La formation est répartie sur de nombreuses machines distinctes, mais les paramètres doivent être régulièrement enregistrés dans un seul objet (« point de contrôle ») sur un système de stockage permanent sans ralentir le processus de formation global. Les tâches de formation individuelles doivent également être capables de lire uniquement l’ensemble spécifique de paramètres qui les concernent afin d’éviter la surcharge qui serait nécessaire pour charger l’ensemble complet des paramètres du modèle (qui pourrait représenter des centaines de gigaoctets).

TensorStore a déjà été utilisé pour relever ces défis. Il a été appliqué pour gérer les points de contrôle associés à grande échelle (« multipode”) modèles entraînés avec JAX (exemple de code) et a été intégré à des frameworks tels que T5X (exemple de code) et Parcours. Le parallélisme des modèles est utilisé pour partitionner l’ensemble complet des paramètres, qui peuvent occuper plus d’un téraoctet de mémoire, sur des centaines de TPU. Les points de contrôle sont stockés au format zarr à l’aide de TensorStore, avec une structure de blocs choisie pour permettre à la partition de chaque TPU d’être lue et écrite indépendamment en parallèle.

Lors de l’enregistrement d’un point de contrôle, chaque paramètre de modèle est écrit à l’aide de TensorStore au format zarr à l’aide d’une grille de blocs qui subdivise davantage la grille utilisée pour partitionner le paramètre sur les TPU. Les machines hôtes écrivent en parallèle les blocs zarr pour chacune des partitions attribuées aux TPU attachés à cet hôte. À l’aide de l’API asynchrone de TensorStore, la formation se poursuit même pendant que les données sont encore en cours d’écriture sur le stockage persistant. Lors de la reprise à partir d’un point de contrôle, chaque hôte lit uniquement les fragments qui composent les partitions attribuées à cet hôte.

Cas d’utilisation : Cartographie cérébrale 3D
Le domaine de la résolution des synapses connectomique vise à cartographier le câblage de animal et Humain cerveaux au niveau détaillé des connexions synaptiques individuelles. Cela nécessite une imagerie du cerveau à une résolution extrêmement élevée (nanomètres) sur des champs de vision allant jusqu’à des millimètres ou plus, ce qui produit des ensembles de données pouvant s’étendre sur des pétaoctets. À l’avenir, ces ensembles de données pourraient s’étendre à des exaoctets alors que les scientifiques envisagent de cartographier souris entière ou des cerveaux de primates. Cependant, même les ensembles de données actuels posent des défis importants liés au stockage, à la manipulation et au traitement ; en particulier, même un seul échantillon de cerveau peut nécessiter des millions de gigaoctets avec un système de coordonnées (espace de pixels) de centaines de milliers de pixels dans chaque dimension.

Nous avons utilisé TensorStore pour résoudre les problèmes de calcul associés aux ensembles de données connectomiques à grande échelle. Plus précisément, TensorStore a géré certains des ensembles de données connectomiques les plus vastes et les plus largement consultés, avec Google Cloud Storage comme système de stockage d’objets sous-jacent. Par exemple, il a été appliqué au jeu de données « h01 » du cortex humain, qui est une image 3D à résolution nanométrique du tissu cérébral humain. Les données d’imagerie brutes sont de 1,4 pétaoctets (environ 500 000 * 350 000 * 5 000 pixels de large, et sont en outre associées à du contenu supplémentaire tel que des segmentations 3D et des annotations qui résident dans le même système de coordonnées. Les données brutes sont subdivisées en blocs individuels de 128x128x16 pixels de large et stocké dans le « Neuroglancer précalculé” format, qui est optimisé pour visualisation interactive sur le Web et peut être facilement manipulé à partir de TensorStore.

Commencer
Pour commencer à utiliser l’API Python TensorStore, vous pouvez installer le paquet PyPI tensorstore utilisant:

pip install tensorstore

Se référer au tutoriels et Documentation API pour les détails d’utilisation. Pour d’autres options d’installation et pour utiliser l’API C++, reportez-vous à instructions d’installation.

Remerciements
Grâce à Tim Blakeley, Viren Jaïn, Yash Kataria, Jan-Matthias Luckmann, Michał Januszewski, Pierre Li, Adam Robert, Cerveau Williamset Hector Yee de Google Research, et Davis Bennett, Stuart Berg, Eric Perlman, Stephen Placeet Juan Nunez-Iglesias de la communauté scientifique au sens large pour des commentaires précieux sur la conception, les premiers tests et le débogage.