Traitement des données avec une logique personnalisée sur toutes les partitions¶

Définition de la fonction de traitement¶

Importez les classes nécessaires à l’exécution du traitement distribué :

from snowflake.ml.modeling.distributors.distributed_partition_function.dpf import DPF
from snowflake.ml.modeling.distributors.distributed_partition_function.dpf_run import DPFRun
from snowflake.ml.modeling.distributors.distributed_partition_function.entities import (
    RunStatus, ExecutionOptions
)

Définissez une fonction de traitement qui prend deux arguments :

• data_connector : Un DataConnector qui donne accès aux données de la partition. Appelez data_connector.to_pandas() pour le charger en tant que DataFrame Pandas, ou utilisez d’autres méthodes comme to_torch_dataset() ou to_ray_dataset().

• context : Un objet PartitionContext qui fournit l’ID de partition et les méthodes pour l’importation et le téléchargement d’artefacts.

import json

def process_sales_data(data_connector, context):
    df = data_connector.to_pandas()
    print(f"Processing {len(df)} records for region: {context.partition_id}")

    # Perform region-specific analytics
    summary = {
        'region': context.partition_id,
        'total_sales': df['amount'].sum(),
        'avg_order_value': df['amount'].mean(),
        'customer_count': df['customer_id'].nunique(),
        'record_count': len(df)
    }

    # Store results in stage for subsequent access

    context.upload_to_stage(summary, "sales_summary.json",
        write_function=lambda obj, path: json.dump(obj, open(path, 'w')))

Pour chaque région, cette fonction calcule des statistiques récapitulatives et enregistre les résultats sous forme de fichier JSON vers le répertoire de zone de préparation dédié à la partition.

Initialiser et exécuter une DPF¶

Créez une instance DPF avec votre fonction de traitement et un nom de zone de préparation de sortie, puis appelez run() pour démarrer le traitement distribué.

dpf = DPF(process_sales_data, "analytics_stage")
run = dpf.run(
    partition_by="region",
    snowpark_dataframe=sales_data,
    run_id="regional_analytics_2024"
)

La méthode run() accepte les paramètres suivants :

• partition_by (str): Nom de la colonne dans laquelle partitionner le DataFrame. Chaque valeur unique crée une partition distincte.

• snowpark_dataframe : Le DataFrame Snowpark à partitionner et à traiter.

• run_id (str): identificateur unique de l’exécution. Crée un répertoire dédié @{stage_name}/{run_id}/ pour tous les artefacts. Utilisez des noms descriptifs comme experiment_2024_01_15 ou model_v1_retrain.

• on_existing_artifacts (chaîne, facultatif) : Action lorsque des artefacts pour l’run_id existent déjà. "error" (par défaut) génère une erreur ; "overwrite" remplace les artefacts existants.

• execution_options (ExecutionOptions, facultatif) : Configuration pour l”=’allocation des ressources et le comportement d’exécution.

Contrôle de la progression et attente de la fin de l’exécution¶

Appelez wait() pour bloquer jusqu’à ce que l’exécution soit terminée. Par défaut, une barre de progression s’affiche.

final_status = run.wait()  # Shows progress bar by default
print(f"Job completed with status: {final_status}")

Voici un exemple de la sortie :

Progress: 100%|██████████| 4/4 [02:15<00:00, 33.75s/it]
Job completed with status: RunStatus.SUCCESS

Vous pouvez également vérifier le statut et la progression à tout moment sans blocage :

# Check overall status
current_status = run.status

# Get progress grouped by partition status
progress = run.get_progress()

Récupération des résultats de chaque partition¶

Une fois l’exécution terminée avec succès, récupérez les résultats de chaque partition en utilisant la propriété partition_details. Les détails de chaque partition incluent un stage_artifacts_manager pour accéder aux artefacts enregistrés.

if final_status == RunStatus.SUCCESS:
    import json
    all_results = []
    for partition_id, details in run.partition_details.items():
        # List available artifacts for this partition
        files = details.stage_artifacts_manager.list()
        print(f"Partition {partition_id} artifacts: {files}")

        # Load an artifact using a custom deserializer
        summary = details.stage_artifacts_manager.get("sales_summary.json",
            read_function=lambda path: json.load(open(path, 'r')))
        all_results.append(summary)

    # Combine results across all regions
    total_sales = sum(r['total_sales'] for r in all_results)
    total_customers = sum(r['customer_count'] for r in all_results)

Le stage_artifacts_manager propose trois méthodes :

• list() : Renvoie une liste de noms de fichiers enregistrés dans le répertoire de zone de préparation de la partition.

• get(filename, read_function=None) : Télécharge et désérialise un artefact. Utilise pickle par défaut, ou une read_function personnalisée si fournie.

• download(filename, local_dir) : Télécharge un fichier brut dans un répertoire local et renvoie le chemin du fichier local.

Traitement des erreurs¶

Si l’exécution échoue, vérifiez les détails des partitions individuelles pour diagnostiquer les échecs :

if final_status != RunStatus.SUCCESS:
    for partition_id, details in run.partition_details.items():
        if details.status != PartitionStatus.DONE:
            print(f"Partition {partition_id} status: {details.status}")
            try:
                error_logs = details.logs
                print(error_logs)
            except RuntimeError:
                print("Logs not available for this partition")

Le RunStatus global reflète le résultat agrégé :

• SUCCESS : Toutes les partitions se sont terminées avec succès.

• PARTIAL_FAILURE : Certaines partitions ont réussi, mais au moins une a échoué.

• FAILURE : Aucune partition a été complétée avec succès.

• CANCELLED : L’exécution a été annulée.

• IN_PROGRESS : L’exécution est toujours en cours.

Chaque partition dispose d’un PartitionStatus :

• PENDING : En attente de traitement.

• RUNNING : En cours de traitement.

• DONE : Terminé correctement.

• FAILED : La fonction utilisateur a soulevé une exception.

• CANCELLED : Annulé par l’utilisateur.

• INTERNAL_ERROR : Une erreur interne du système s’est produite (par exemple, manque de mémoire).

Pour importer ces énumérations :

from snowflake.ml.modeling.distributors.distributed_partition_function.entities import (
    RunStatus, PartitionStatus
)

Pour annuler une tâche en cours :

run.cancel()

Restauration des résultats d’une exécution terminée¶

Vous pouvez restaurer une exécution terminée à partir de son état persistant et accéder aux mêmes résultats sans relancer le processus :

from snowflake.ml.modeling.distributors.distributed_partition_function.dpf_run import DPFRun

restored_run = DPFRun.restore_from(run_id="regional_analytics_2024", stage_name="analytics_stage")

# Access results just like the original run
for partition_id, details in restored_run.partition_details.items():
    print(f"{partition_id}: {details.status}")

Définition de la fonction de traitement¶

La signature de la fonction de traitement est la même que le mode DataFrame. Le data_connector donne accès aux données du fichier, et context.partition_id est le chemin d’accès relatif au fichier dans la zone de préparation.

def process_file(data_connector, context):
    df = data_connector.to_pandas()
    print(f"Processing file {context.partition_id}: {len(df)} rows")

    # Process data and save results
    result = {"row_count": len(df), "columns": list(df.columns)}
    import json
    context.upload_to_stage(result, "result.json",
        write_function=lambda obj, path: json.dump(obj, open(path, 'w')))

Ajouter une DPF depuis une zone de préparation¶

Appelez run_from_stage() au lieu de run(). Spécifier la stage_location d’entrée contenant les fichiers sources et éventuellement un file_pattern pour filtrer les fichiers à traiter.

dpf = DPF(process_file, "output_stage")
run = dpf.run_from_stage(
    stage_location="@my_db.my_schema.input_stage/data/",
    run_id="file_processing_2024",
    file_pattern="*.parquet",
)
final_status = run.wait()

La méthode run_from_stage() accepte les paramètres suivants :

• stage_location (str): Chemin de la zone de préparation d’entrée contenant les fichiers de données source. Chaque fichier correspondant au file_pattern devient une partition. Prend en charge les noms de zone de préparation simples et complets :

  • Nom simple : "@my_stage/data/"

  • Nom complet : "@my_db.my_schema.my_stage/data/"

• run_id (chaîne) : identificateur unique de l’exécution.

• file_pattern (chaîne, facultatif) : Modèle global pour filtrer les fichiers. La valeur par défaut est "*.parquet".

• on_existing_artifacts (chaîne, obligatoire) : "error" (par défaut) ou "overwrite".

• execution_options (ExecutionOptions, facultatif) : Configuration pour l”=’allocation des ressources et le comportement d’exécution.

La surveillance, la récupération des résultats, la gestion des erreurs et la restauration des exécutions fonctionnent de la même manière qu’en mode DataFrame.

Pour le traitement des fichiers liés aux E/S, définissez num_cpus_per_worker=1 dans ExecutionOptions pour maximiser le parallélisme (un acteur par CPU). Pour les charges de travail liées au CPU, utilisez la valeur par défaut ou augmentez num_cpus_per_worker.

from snowflake.ml.modeling.distributors.distributed_partition_function.entities import ExecutionOptions

run = dpf.run_from_stage(
    stage_location="@my_stage/data/",
    run_id="io_bound_processing",
    execution_options=ExecutionOptions(num_cpus_per_worker=1),
)

Importer les artefacts¶

Utilisez upload_to_stage() pour enregistrer les résultats de dans votre fonction de traitement. Par défaut, les objets sont sérialisés à l’aide de pickle. Fournissez une write_function pour une sérialisation personnalisée.

def my_function(data_connector, context):
    df = data_connector.to_pandas()

    # Save a pickle object (default serialization)
    results = {'total': df['amount'].sum(), 'count': len(df)}
    context.upload_to_stage(results, "summary.pkl")

    # Save JSON data with a custom write function
    import json
    context.upload_to_stage(
        results, "summary.json",
        write_function=lambda obj, path: json.dump(obj, open(path, 'w'))
    )

    # Save a CSV file
    df_processed = df.groupby('category').sum()
    context.upload_to_stage(
        df_processed, "aggregated.csv",
        write_function=lambda df, path: df.to_csv(path, index=False)
    )

Télécharger les artefacts¶

Utilisez download_from_stage() pour charger des artefacts dans votre fonction de traitement. Vous pouvez utiliser cette fonction pour accéder aux artefacts d’une exécution précédente. Par exemple, vous pouvez l’utiliser pour charger un modèle pour l’inférence.

def my_inference_function(data_connector, context):
    # Load a pickle object from the current partition's stage path
    model = context.download_from_stage("model.pkl")

    # Load from a different stage path (e.g., from a prior training run)
    model = context.download_from_stage(
        "model.pkl",
        stage_path="@db.schema.stage/training_run/partition_0"
    )

    # Load JSON with a custom deserializer
    import json
    config = context.download_from_stage(
        "config.json",
        read_function=lambda path: json.load(open(path, 'r'))
    )

Utiliser les sessions Snowflake¶

Utilisez with_session() pour exécuter des opérations qui nécessitent une session Snowflake, telles que l’écriture de résultats dans une table. Cette méthode utilise un pool de sessions limité pour éviter d’atteindre les limites de session de Snowflake lors de l’exécution simultanée de plusieurs partitions.

def my_function(data_connector, context):
    df = data_connector.to_pandas()

    # Load a model from a prior training run
    model = context.download_from_stage("model.pkl")

    predictions = model.predict(df[['X1', 'X2']])

    results = df.copy()
    results['predictions'] = predictions
    results['partition_id'] = context.partition_id

    # Write results to a Snowflake table using the bounded session pool
    context.with_session(lambda session:
        session.create_dataframe(results)
            .write.mode("append")
            .save_as_table("predictions_table")
    )