Charger les tables d’Unity Catalog dans PyTorch avec `UCDataset`

info

Bêta

Cette fonctionnalité est en bêta.

UCDataset est un IterableDataset PyTorch qui transmet en streaming des données de tables Unity Catalog vers une boucle d'entraînement sur AI Runtime. Vous décrivez ce qui doit être lu comme un DataFrame : choisissez une table, puis sélectionnez les colonnes et filtrez les lignes. Le dataset lit les lignes correspondantes directement à partir des fichiers de données de la table dans le stockage cloud, vous n'avez donc pas besoin d'une session Spark ou d'un SQL warehouse. Il applique les vecteurs de suppression de la table lors de la lecture, divise les lignes correspondantes uniformément entre les GPU et les Worker du chargeur de données, et génère un dictionnaire par batch. Les colonnes numériques et booléennes deviennent des tenseurs. Les autres colonnes deviennent des listes Python.

Choisissez un dataset en fonction de l'emplacement de vos données d'entraînement :

Données d'entraînement	Jeu de données
Lignes dans une table Unity Catalog, telles que les colonnes de caractéristiques et d’étiquettes dans une table Delta Lake	`UCDataset` (Bêta)
Fichiers dans un volume Unity Catalog, tels que des images, du contenu audio ou des documents	`UCVolumeDataset`

Données d'entraînement	Jeu de données
Lignes dans une table Unity Catalog, telles que les colonnes de caractéristiques et d’étiquettes dans une table Delta Lake	`UCDataset` (Bêta)
Fichiers dans un volume Unity Catalog, tels que des images, du contenu audio ou des documents	`UCVolumeDataset`

Les deux datasets fonctionnent côte à côte. Par exemple, un job d’entraînement peut lire des étiquettes à partir d’une table avec UCDataset et les images correspondantes à partir d’un volume avec UCVolumeDataset.

Pour d'autres données, utilisez plutôt ces options :

Pour les petites tables qui tiennent en mémoire, ou pour les queries qui nécessitent des jointures ou des agrégations, utilisez Spark Connect.
Pour les charges de travail Ray, utilisez ray.data.read_databricks_tables().

Exigences

Une charge de travail AI Runtime. UCDataset est disponible partout où AI Runtime est disponible. Voir disponibilité AI Runtime par Geo.
Linux sur x86_64 avec Python 3.10 ou supérieur. Le package est conçu pour x86_64 uniquement. Le compute GPU AI Runtime répond à cette exigence.
Une table Delta Lake inscrite dans le Unity Catalog.
Le privilège USE CATALOG sur le catalogue de la table, le privilège USE SCHEMA sur son schéma et le privilège SELECT sur la table.
Pour effectuer un entraînement sur plus d’un GPU, un lanceur qui définit les variables d’environnement RANK et WORLD_SIZE pour chaque processus, comme torchrun. Si aucune variable n’est définie, comme dans un notebook AI Runtime, UCDataset lit la table comme une seule partition. Si une seule d’entre elles est définie, UCDataset génère une erreur.

important

Lorsque vous soumettez une charge de travail avec databricks air run, start le script d'entraînement avec torchrun, même sur un seul GPU. Runtime AI définit WORLD_SIZE pour chaque charge de travail mais ne définit pas RANK, et UCDataset génère une erreur lorsque seul l'un des deux est défini. torchrun définit les deux pour chaque processus. Voir Exécuter un job multi-nœuds sur le Runtime AI.

Charger une table

Installez le package databricks-uc-dataset. Pour une charge de travail que vous soumettez avec databricks air run, ajoutez le package à environment.dependencies. Voir Exécuter un job multi-nœuds sur AI Runtime. Dans un notebook, exécutez ce qui suit :
Python
```
%pip install databricks-uc-dataset==0.0.1b1
```

Sélectionnez la table, les colonnes et les lignes à lire. table() accepte un nom de table à trois niveaux, et l’argument optionnel version_as_of lit une version de table spécifique :

Python
from databricks.uc_dataset.beta.dataframe import table, col

df = (
    table("main.ml.clicks")
    .select("features", "label")
    .filter(col("label").isNotNull() & col("features").isNotNull())
)

Créez le dataset et enveloppez-le dans un DataLoader PyTorch. Le dataset retourne déjà des batchs, donc définissez batch_size=None sur le chargeur. Si vous entraînez avec PyTorch DistributedDataParallel (DDP), appelez torch.distributed.init_process_group() avant de créer le dataset. Le rang 0 planifie ensuite la lecture une fois et partage le plan avec les autres rangs :

Python
import torch
from torch.utils.data import DataLoader
from databricks.uc_dataset.beta import UCDataset

dataset = UCDataset(df=df, batch_size=256)
loader = DataLoader(dataset, batch_size=None, num_workers=4)

for epoch in range(num_epochs):
    for batch in loader:
        # Array columns arrive as Python lists, and a NULL array arrives as None.
        # Convert the lists before you train.
        features = torch.as_tensor(batch["features"], dtype=torch.float32)
        label = batch["label"]
        train_step(features, label)

Pour entraîner sur plus d'un GPU, enregistrez le code sous le nom train.py et lancez-le avec torchrun. Dans un notebook, ignorez cette étape. Par exemple, pour effectuer l’entraînement sur 8 GPU sur un nœud, exécutez la commande suivante :
Bash
```
torchrun --nproc_per_node=8 train.py
```

Chaque rang et chaque worker lit son propre shard des lignes correspondantes. Chaque batch est un dictionnaire qui mappe chaque nom de colonne sélectionné à un tenseur ou à une liste.

Exécuter un job multinœud sur AI Runtime

Pour entraîner sur plus d'un nœud, décrivez la charge de travail dans workload.yaml et soumettez-la avec databricks air run:

Réglez environment.version sur databricks_ai_v6, qui inclut PyTorch, et ajoutez databricks-uc-dataset==0.0.1b1 à environment.dependencies.
start le script d’entraînement avec torchrun. Le Runtime IA exécute la commande sur chaque nœud et définit NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR et MASTER_PORT. torchrun les utilise pour définir RANK pour chaque processus :
YAML
```
command: |
  torchrun --nnodes=$NUM_NODES --node_rank=$NODE_RANK --nproc_per_node=$LOCAL_WORLD_SIZE \
    --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT "$CODE_SOURCE_PATH/train.py"
```

Pour une charge de travail complète avec train.py, workload.yaml et une table d’exemple, voir Modèle de taux de clic sur une table Unity Catalog dans le livre de recettes AI Runtime. Pour entraîner sur plus de GPU, modifiez num_accelerators dans workload.yaml. train.py ne change pas.

Filtrer les lignes

UCDataset applique des filtres pendant la lecture des fichiers, de sorte que les lignes qui ne correspondent pas n'atteignent pas votre boucle d'entraînement. Créez un filtre avec col() et les expressions suivantes :

Expression	Exemple
Comparaison	`col("amount") > 50`. Ainsi que `==`, `!=`, `<`, `<=` et `>=`.
Vérification nulle	`col("label").isNull()`, `col("label").isNotNull()`
Égalité sécurisée contre les valeurs nulles	`col("region").eqNullSafe("us")`
Logique booléenne	`(col("a") > 1) & (col("b") == "x")`. Ainsi que `\|` (ou) et `~` (non).

Expression	Exemple
Comparaison	`col("amount") > 50`. Ainsi que `==`, `!=`, `<`, `<=` et `>=`.
Vérification nulle	`col("label").isNull()`, `col("label").isNotNull()`
Égalité sécurisée contre les valeurs nulles	`col("region").eqNullSafe("us")`
Logique booléenne	`(col("a") > 1) & (col("b") == "x")`. Ainsi que `\|` (ou) et `~` (non).

Les filtres ne prennent pas en charge les chaînes SQL, isin ou IN, BETWEEN, la correspondance de chaînes telle que like ou contains, l’arithmétique sur les colonnes, cast ou les fonctions définies par l’utilisateur (UDF). Transmettez chaque condition sous la forme d’une expression Column et combinez les conditions avec &, | et ~. Une comparaison prend une colonne et une valeur littérale, dans n’importe quel ordre. Vous ne pouvez pas comparer deux colonnes. Pour correspondre à une liste de valeurs, combinez des conditions == avec |. Les filtres très longs peuvent échouer. Voir Limitations.

L’exemple suivant lit les événements d’achat à partir d’une version de table fixe. La version doit exister :

Python
df = (
    table("main.ml.events", version_as_of=120)
    .filter((col("event_type") == "purchase") & (col("amount") > 50))
    .select("user_id", "amount", "label")
)

Format de sortie

Chaque batch est un dictionnaire Python qui mappe chaque nom de colonne sélectionné à ses valeurs :

Type de colonne	Valeur de batch
Entier, sans valeurs nulles dans le batch	Tenseur 1D avec le type entier de la colonne, tel que `int32`
Entier, avec des valeurs nulles dans le batch (sauf 64 bits)	Tenseur `float64` 1D, avec des valeurs nulles comme `NaN`
Virgule flottante	Tenseur 1D, avec des valeurs nulles en tant que `NaN`
Booléen, sans valeurs nulles	Tenseur booléen
Binaire	Liste des objets `bytes`
Autres types, tels que les chaînes, les tableaux, les structs et les Timestamp	Liste Python

Type de colonne	Valeur de batch
Entier, sans valeurs nulles dans le batch	Tenseur 1D avec le type entier de la colonne, tel que `int32`
Entier, avec des valeurs nulles dans le batch (sauf 64 bits)	Tenseur `float64` 1D, avec des valeurs nulles comme `NaN`
Virgule flottante	Tenseur 1D, avec des valeurs nulles en tant que `NaN`
Booléen, sans valeurs nulles	Tenseur booléen
Binaire	Liste des objets `bytes`
Autres types, tels que les chaînes, les tableaux, les structs et les Timestamp	Liste Python

Comme le type d’une colonne d’entiers dépend du fait que le batch contienne ou non un nul, le type de colonne peut changer d’un batch à l’autre. Si votre modèle nécessite un type entier fixe, par exemple pour les recherches d'embeddings ou les libellés de classe, filtrez les nuls ou remplacez-les dans une fonction transform.

Si une colonne ne peut pas être convertie sans perte de données, UCDataset génère une erreur. Cela se produit pour une colonne d'entiers 64 bits qui contient des valeurs nulles et pour une colonne booléenne qui contient des valeurs nulles. L'erreur apparaît lorsqu'un batch avec une valeur nulle est lu, elle peut donc se produire au milieu d'une époque. Pour l'éviter, filtrez les valeurs nulles avec isNotNull(), ou remplacez-les dans une fonction transform.

Pour modifier le layout de batch ou gérer les valeurs nulles, transmettez une fonction transform. UCDataset l'appelle sur chaque batch en tant que PyArrow RecordBatch, avant toute conversion en tenseurs, et renvoie ce que la fonction retourne, alors convertissez chaque colonne que vous utilisez. L'exemple suivant renvoie un tuple de tenseurs :

Python
import torch

def to_features_and_label(batch):
    features = torch.as_tensor(batch.column("features").to_pylist(), dtype=torch.float32)
    label = torch.from_numpy(batch.column("label").to_numpy(zero_copy_only=False))
    return features, label

dataset = UCDataset(df=df, transform=to_features_and_label)

L’exemple suivant remplace les valeurs nulles d’une colonne d’entiers 64 bits par 0, afin que la colonne reste int64 dans chaque batch :

Python
import pyarrow.compute as pc
import torch

def fill_nulls(batch):
    user_ids = pc.fill_null(batch.column("user_id"), 0)
    return {"user_id": torch.tensor(user_ids.to_numpy())}

dataset = UCDataset(df=df, transform=fill_nulls)

Formation distribuée et mélange

UCDataset divise les lignes qui correspondent à votre filtre sur chaque rang et chaque worker de chargeur de données. Vous n’avez pas besoin d’un échantillonneur distribué. Par default, les tailles de shard diffèrent d’au plus une ligne.

Des tailles de shard égales ne garantissent pas des nombres de batch égaux. Un rang peut obtenir un batch de plus qu’un autre, par exemple lorsque certains shards contiennent un multiple exact de batch_size lignes et que d’autres en contiennent une de plus, ou lorsqu’un filtre conserve moins de lignes qu’il n’y a de shards. Avec DDP, le rang avec le batch supplémentaire attend ensuite un gradient all-reduce auquel les autres rangs ne se joignent jamais, et le job échoue. Pour permettre aux rangs de terminer une époque à des moments différents, enveloppez chaque époque dans le gestionnaire de contexte PyTorch Join, où model est le modèle enveloppé DDP :

Python
from torch.distributed.algorithms.join import Join

for epoch in range(num_epochs):
    with Join([model]):
        for batch in loader:
            train_step(batch)

dataset.num_rows() renvoie le nombre de lignes qui correspondent au filtre, à travers tous les rangs et les Worker. len(dataset) et len(loader) lèvent une TypeError. Si une bibliothèque a besoin du nombre de pas dans une époque, comme un planificateur de taux d'apprentissage, estimez-le à partir de num_rows(), de la taille du monde et de batch_size. Chaque rang peut lire un batch supplémentaire par rapport à l'estimation, et chaque worker de chargeur de données supplémentaire peut ajouter un autre batch partiel.

UCDataset lit les fichiers de la table dans un ordre fixe et aléatoire. Chaque époque utilise le même ordre, tout comme un nouveau dataset sur une table inchangée. Vous ne pouvez pas définir de seed, et les fichiers ne sont pas mélangés à nouveau entre les époques.

Authentification

UCDataset lit la table en tant qu'identité qui exécute la charge de travail, et les autorisations Unity Catalog sur la table s'appliquent.

Le Unity Catalog ne fournit pas d'identifiants de stockage pour les tables avec des filtres de ligne ou des masques de colonne. Par conséquent, UCDataset ne peut pas les lire. La création du dataset échoue avec une erreur qui inclut le message du Unity Catalog.

Limitations

UCDataset lit uniquement les tables Delta Lake dans Unity Catalog. Les vues, les vues matérialisées et les tables de streaming ne sont pas prises en charge.
Les tables avec filtres de ligne ou masques de colonne ne peuvent pas être lues, car Unity Catalog ne délivre pas d'identifiants de stockage pour elles.
Les clones superficiels ne peuvent pas être lus, car leurs fichiers de données sont stockés avec la table source. Pour pin les lignes d’une exécution d’entraînement, lisez plutôt la table source avec version_as_of. Un clone start son propre historique de version, il n’a donc pas les versions antérieures de la table source.
Les tables avec une colonne VARIANT ne peuvent pas être lues, même si vous ne sélectionnez pas cette colonne.
version_as_of prend un numéro de version de table. Les Timestamp ne sont pas pris en charge : version_as_of="2026-01-01" génère une erreur versionAsOf must be a non-negative integer.
Sélectionnez au moins une colonne qui n’est pas une colonne de partition. La sélection exclusive de colonnes de partition provoque une erreur Delta reads require at least one non-partition column.
Un filtre avec une longue chaîne de conditions, comme de nombreuses comparaisons == combinées avec |, peut échouer avec une erreur failed to decode Spark Connect plan. Si cela se produit, utilisez moins de conditions.
Les filtres ne prennent en charge que les expressions décrites dans Filtrer des lignes.
Les colonnes de tableau, telles que les encastrements, sont renvoyées sous forme de listes Python, ce qui est beaucoup plus lent que les colonnes numériques.
UCDataset ne définit pas de point de contrôle pour sa position de lecture. Si l'entraînement redémarre, le dataset repart du début de la table et vous ne pouvez pas reprendre au milieu d'une époque.

Pour signaler un problème ou nous faire part de vos commentaires pendant la bêta, contactez l’équipe de votre compte Databricks.

Ressources supplémentaires

Charger les données sur le Runtime IA
Formation distribuée dans les notebooks
Modèle de taux de clic sur une table Unity Catalog dans le livre de recettes Runtime AI
Améliorez les performances d'entraînement et la résilience sur AI Runtime

Exigences​

Charger une table​

Exécuter un job multinœud sur AI Runtime​

Filtrer les lignes​

Format de sortie​

Formation distribuée et mélange​

Authentification​

Limitations​

Ressources supplémentaires​