breeders - centralize dag task dependencies again

godon-dev · Nov 19, 2023 · d8dcdfd · d8dcdfd
1 parent 398e9c1
commit d8dcdfd
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 17 deletions.
diff --git a/breeder/linux_network_stack/effectuation.py b/breeder/linux_network_stack/effectuation.py
@@ -13,7 +13,6 @@ def create_target_interaction_dag(dag_id, config, target, identifier):
 
         @dag.task(task_id="pull_optimization_step")
         def run_pull_optimization():
-            import asyncio
 
             task_logger.debug("Entering")
 
@@ -31,7 +30,6 @@ def run_pull_optimization():
         def run_aquire_lock():
             task_logger.debug("Entering")
 
-            import pals
 
             locker = pals.Locker('network_breeder_effectuation', DLM_DB_CONNECTION)
 
@@ -61,10 +59,6 @@ def run_release_lock():
         @dag.task(task_id="push_optimization_step")
         def run_push_optimization(ti=None):
 
-            import asyncio
-            from sqlalchemy import create_engine
-            from sqlalchemy import text
-
             archive_db_engine = create_engine(f'postgresql://{ARCHIVE_DB_USER}:{ARCHIVE_DB_PASSWORD}@{ARCHIVE_DB_HOST}:{ARCHIVE_DB_PORT}/{ARCHIVE_DB_DATABASE}')
             task_logger.debug("Entering")
 
@@ -98,10 +92,6 @@ def run_push_optimization(ti=None):
         @dag.task(task_id="recon_step")
         def run_reconnaissance():
 
-            from prometheus_api_client import PrometheusConnect, MetricsList, Metric
-            from prometheus_api_client.utils import parse_datetime
-            import urllib3
-
             task_logger.debug("Entering")
             prom_conn = PrometheusConnect(url=PROMETHEUS_URL,
                                           retry=urllib3.util.retry.Retry(total=3, raise_on_status=True, backoff_factor=0.5),

diff --git a/breeder/linux_network_stack/optimization.py b/breeder/linux_network_stack/optimization.py
@@ -8,9 +8,6 @@ def objective(trial, identifier):
     {{ local_coroutines_include()|indent }} # default is indent of 4 spaces!
 ###--- end coroutines ---###
 
-    import logging
-    from sqlalchemy import create_engine
-    from sqlalchemy import text
 
     logger = logging.getLogger('objective')
     logger.setLevel(logging.DEBUG)
@@ -84,10 +81,6 @@ def create_optimization_dag(dag_id, config, identifier):
         ## perform optimiziation run
         @dag.task(task_id="optimization_step")
         def run_optimization():
-            import optuna
-            from optuna.storages import InMemoryStorage
-            from optuna.integration import DaskStorage
-            from distributed import Client, wait
 
             __directions = list()
 

diff --git a/breeder/linux_network_stack/root_dag.py b/breeder/linux_network_stack/root_dag.py
@@ -36,6 +36,25 @@
 import hashlib
 import os
 
+from sqlalchemy import create_engine
+from sqlalchemy import text
+
+import optuna
+from optuna.storages import InMemoryStorage
+from optuna.integration import DaskStorage
+from distributed import Client, wait
+
+import asyncio
+import pals
+
+import asyncio
+from sqlalchemy import create_engine
+from sqlalchemy import text
+
+from prometheus_api_client import PrometheusConnect, MetricsList, Metric
+from prometheus_api_client.utils import parse_datetime
+import urllib3
+
 task_logger = logging.getLogger("airflow.task")
 task_logger.setLevel(logging.DEBUG)