refactor: wait for variables sleep 10s and get-variable function

dags/ml-training.py

@@ -11,30 +11,32 @@ import time
 
 import time
 
-def get_variable(var_name, default=None):
+def get_variable(var_name, default=None, retries=5, delay=5):
     """
-    Retrieves an Airflow variable safely. If the variable does not exist, it logs a warning and returns a default value.
-
-    Args:
-        var_name (str): The name of the variable.
-        default (any): Default value if the variable is not found.
-
-    Returns:
-        str: The variable value or default if not found.
+    Retrieves an Airflow variable safely with retries.
+    If the variable does not exist or is empty, it waits and retries.
     """
-    try:
-        value = Variable.get(var_name)
-        print(f"✅ DEBUG: Retrieved {var_name} = {value}")  # Debugging log
-        return value
-    except KeyError:
-        print(f"⚠️ WARNING: {var_name} not found. Using default: {default}")  # Debug warning
-        return default  # Instead of retrying infinitely, just use the default
+    for attempt in range(retries):
+        try:
+            value = Variable.get(var_name)
+            if value.strip():  # Check if not empty
+                print(f"✅ DEBUG: Retrieved {var_name} = {value}")
+                return value
+            else:
+                print(f"⚠️ WARNING: {var_name} is set but EMPTY in Airflow! Retrying in {delay}s...")
+        except KeyError:
+            print(f"⚠️ WARNING: {var_name} not found. Retrying in {delay}s...")
+
+        time.sleep(delay)  # Wait before retrying
+
+    print(f"❌ ERROR: {var_name} could not be retrieved after {retries} retries!")
+    return default
 
 # Fetch Airflow Variables
 GCP_PROJECT_ID = get_variable("GCP_PROJECT_ID")
 GCP_ZONE = get_variable("GCP_ZONE")
-GCE_INSTANCE_NAME = f"training-instance-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
 SERVICE_ACCOUNT_EMAIL = get_variable("SERVICE_ACCOUNT_EMAIL")
+GCE_INSTANCE_NAME = f"training-instance-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
 
 # Debug Log
 print("\n===== ✅ ENVIRONMENT CONFIGURATION CHECK (Inside DAG) =====")

start.sh

@@ -78,7 +78,6 @@ fi
 export SERVICE_ACCOUNT_EMAIL
 log "✅ Retrieved SERVICE_ACCOUNT_EMAIL: $SERVICE_ACCOUNT_EMAIL"
 
-
 ### 🔹 Enable Required GCP APIs
 log "🔹 Enabling required GCP APIs..."
 GCP_APIS=(
@@ -112,8 +111,10 @@ log "✅ Airflow variables set."
 
 ### 🚀 Start Airflow Services
 log "🚀 Starting Airflow services..."
+sleep 10  # Ensure variables are set before DAG execution
 airflow scheduler > /dev/null 2>&1 &
 airflow webserver > /dev/null 2>&1 &
+log "✅ Airflow services started."
 
 log "✅ Airflow services started."
 wait