trying list comp

lsst-uk · Aug 15, 2024 · 11aec13 · 11aec13
1 parent 5872f8d
commit 11aec13
Showing 1 changed file with 22 additions and 22 deletions.
diff --git a/echo-side/dags/process_new_zips.py b/echo-side/dags/process_new_zips.py
@@ -26,16 +26,16 @@ def dl_bucket_names(**kwargs):
 }
 
 # Instantiate the DAG
-dag = DAG(
-    'process_zips',
-    default_args=default_args,
-    description='Runs process_collated_zips.py',
-    schedule_interval=timedelta(hours=1), # change to daily once in production
-    start_date=datetime(2024, 1, 1),
-    catchup=False,
-)
 
-with dag:
+
+with DAG(
+        'process_zips',
+        default_args=default_args,
+        description='Runs process_collated_zips.py',
+        schedule_interval=timedelta(hours=1), # change to daily once in production
+        start_date=datetime(2024, 1, 1),
+        catchup=False,
+    ) as dag:
 
     get_bucket_names = PythonOperator(
         task_id = 'get_bucket_names',
@@ -44,19 +44,19 @@ def dl_bucket_names(**kwargs):
     )
 
     if len(bucket_names) > 0:
-        print(f'Found {bucket_names} bucket names.')
-        for i, bucket_name in enumerate(bucket_names):
-            process_zips_task = KubernetesPodOperator(
-                task_id=f'process_zips_{i}',
-                image='ghcr.io/lsst-uk/csd3-echo-somerville:latest',
-                cmds=['./entrypoint.sh'],
-                arguments=['python', 'csd3-echo-somerville/scripts/process_collated_zips.py', '--bucket_name', bucket_name, '--extract', '--nprocs', '16'],
-                env_vars={
-                    'ECHO_S3_ACCESS_KEY': Variable.get("ECHO_S3_ACCESS_KEY"),
-                    'ECHO_S3_SECRET_KEY': Variable.get("ECHO_S3_SECRET_KEY"),
-                },
-                get_logs=True,
-            )
+        print(f'Bucket names found: {bucket_names}')
+        process_zips_task = [
+            KubernetesPodOperator(
+            task_id=f'process_zips_{i}',
+            image='ghcr.io/lsst-uk/csd3-echo-somerville:latest',
+            cmds=['./entrypoint.sh'],
+            arguments=['python', 'csd3-echo-somerville/scripts/process_collated_zips.py', '--bucket_name', bucket_name, '--extract', '--nprocs', '16'],
+            env_vars={
+                'ECHO_S3_ACCESS_KEY': Variable.get("ECHO_S3_ACCESS_KEY"),
+                'ECHO_S3_SECRET_KEY': Variable.get("ECHO_S3_SECRET_KEY"),
+            },
+            get_logs=True,
+        ) for bucket_name in bucket_names]
     else:
         print('No bucket names found.')