adding checkpoint to airflow pipelines

Author: Dpbm

dags/dataset.py

@@ -4,7 +4,7 @@
 
 from airflow import DAG
 from airflow.providers.standard.operators.bash import BashOperator
-from airflow.providers.standard.operators.python import PythonOperator
+from airflow.providers.standard.operators.python import PythonOperator, BranchPythonOperator
 from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
 
 from generate.dataset import (
@@ -13,6 +13,8 @@
     remove_duplicated_files,
     transform_images,
     start_df,
+    Checkpoint,
+    Stages
 )
 from utils.constants import (
     DEFAULT_NUM_QUBITS,
@@ -22,6 +24,7 @@
     DEFAULT_SHOTS,
     DEFAULT_DATASET_SIZE,
     DEFAULT_THREADS,
+    images_gen_checkpoint_file
 )
 from generate.ghz import gen_circuit
 from export.kaggle import upload_dataset as upload_dataset_kaggle
@@ -31,6 +34,34 @@
     "depends_on_past": True,
 }
 
+GEN_IMAGES_TASK_ID = 'gen_images'
+REMOVE_DUPLICATES_TASK_ID = 'remove_duplicates'
+TRANSFORM_TASK_ID = 'transform_images'
+
+def next_step(checkpoint:Checkpoint) -> str:
+    """
+    checks teh current checkpoint and returns the next
+    task_id.
+    """
+
+    if (checkpoint.stage == Stages.GEN_IMAGES):
+        return GEN_IMAGES_TASK_ID
+    
+    if(checkpoint.stage == Stages.DUPLICATES):
+        return REMOVE_DUPLICATES_TASK_ID
+
+    return TRANSFORM_TASK_ID
+
+def update_checkpoint(checkpoint:Checkpoint, stage:Stages):
+    """
+    Updates the checkpoint to start next task.
+    """
+    checkpoint.index = 0
+    checkpoint.files = []
+    checkpoint.stage = Stages.TRANSFORM
+    checkpoint.save()
+
+
 with DAG(
     dag_id="build_dataset",
     default_args=default_args,
@@ -47,6 +78,17 @@
     create_folder.doc_md = """
     Create a folder (if it doesn't exist) to store images.
     """
+    
+    checkpoint = Checkpoint(images_gen_checkpoint_file(folder))
+
+    branch_checkpoint = BranchPythonOperator(
+        task_id="check_checkpoint",
+        python_callable=next_step,
+        op_args=[checkpoint]
+    )
+    branch_checkpoint.doc_md = """
+    Choose the next task based on the current checkpoint.
+    """
 
     gen_df = PythonOperator(
         task_id="gen_df", python_callable=start_df, op_args=[folder]
@@ -56,7 +98,7 @@
     """
 
     gen_images = PythonOperator(
-        task_id="gen_images",
+        task_id=GEN_IMAGES_TASK_ID,
         python_callable=generate_images,
         op_args=[
             folder,
@@ -65,6 +107,7 @@
             DEFAULT_SHOTS,
             DEFAULT_DATASET_SIZE,
             DEFAULT_THREADS,
+            checkpoint
         ],
     )
 
@@ -73,20 +116,40 @@
     Qiskit framework.
     """
 
+    transtion_gen_to_remove = PythonOperator(
+        task_id="gen_to_remove",
+        python_callable=update_checkpoint,
+        op_args=[checkpoint, Stages.DUPLICATES]
+    )
+
+    transtion_gen_to_remove.doc_md = """
+    Update checkpoint to start removing duplicated files.
+    """
+
     remove_duplicates = PythonOperator(
-        task_id="remove_duplicates",
+        task_id=REMOVE_DUPLICATES_TASK_ID,
         python_callable=remove_duplicated_files,
-        op_args=[folder],
+        op_args=[folder, checkpoint],
     )
 
     remove_duplicates.doc_md = """
     Remove files that have the same hashes.
     """
 
+    transition_remove_to_transform = PythonOperator(
+        task_id="remove_to_transform",
+        python_callable=update_checkpoint,
+        op_args=[checkpoint, Stages.TRANSFORM]
+    )
+    
+    transition_remove_to_transform.doc_md = """
+    Update checkpoint to start transforming images.
+    """
+
     transform_img = PythonOperator(
-        task_id="transform_images",
+        task_id=TRANSFORM_TASK_ID,
         python_callable=transform_images,
-        op_args=[folder, DEFAULT_NEW_DIM],
+        op_args=[folder, DEFAULT_NEW_DIM, checkpoint],
     )
 
     transform_img.doc_md = """
@@ -137,9 +200,16 @@
     """
 
     create_folder >> [gen_ghz, gen_df]
-    gen_df >> gen_images
-    gen_images >> remove_duplicates
-    remove_duplicates >> transform_img
+    gen_df >> branch_checkpoint
+
+    branch_checkpoint >> gen_images
+    branch_checkpoint >> remove_duplicates
+    branch_checkpoint >> transform_img
+
+    gen_images >> transtion_gen_to_remove
+    transtion_gen_to_remove >> remove_duplicates
+    remove_duplicates >> transition_remove_to_transform
+    transition_remove_to_transform >> transform_img
     transform_img >> pack_img
 
     [gen_ghz, pack_img] >> trigger_dag_train

dags/train.py

@@ -8,6 +8,7 @@
 from train import setup_and_run_training
 from args.parser import Arguments
 from utils.constants import DEFAULT_TARGET_FOLDER
+from utils.helpers import get_latest_model_checkpoint
 from export.kaggle import upload_model as upload_model_kaggle
 from export.huggingface import upload_model as upload_model_hf
 
@@ -17,6 +18,10 @@
     args = Arguments()
     args.target_folder = folder
 
+    checkpoint = get_latest_model_checkpoint(folder)
+    if(checkpoint):
+        args.checkpoint = checkpoint
+
     train = PythonOperator(
         task_id="train_model", 
         python_callable=setup_and_run_training, 

utils/helpers.py

@@ -2,10 +2,12 @@
 
 from typing import List, Optional
 import random
+import os
 
 import torch
 
-from utils.constants import DEBUG
+from utils.constants import DEBUG, CHECKPOINT_FILE_PREFIX
+from utils.datatypes import FilePath
 
 
 class PlotImages:
@@ -53,3 +55,31 @@ def get_measurements(n_qubits: int) -> List[int]:
     total_measurements = random.randint(1, n_qubits)
     qubits = list(range(n_qubits))
     return random.sample(qubits, total_measurements)
+
+def get_latest_model_checkpoint(target_folder:FilePath) -> Optional[FilePath]:
+    """Returns the path of the lastest checkpoint"""
+
+    if not os.path.exists(target_folder):
+        return None
+
+    files = [
+        os.path.join(target_folder, file)
+        for file in os.listdir(target_folder)
+        if file.startswith(CHECKPOINT_FILE_PREFIX)
+    ]
+
+    if not files:
+        return None
+
+    modification_time = [os.path.getmtime(file) for file in files]
+
+    latest_time = max(modification_time)
+
+    file_index = modification_time.index(latest_time)
+
+    return files[file_index]
+
+
+
+
+