Merge pull request #42 from airflow-laminar/tkp/foo

Make supervisor independent of DAG logic for better integration in airflow-config

Author: timkpaine

README.md

@@ -38,6 +38,86 @@ Here is a nice overview of the DAG, with annotations for code paths and the acti
 
 More docs and code examples coming soon!
 
+### Example Dag:
+
+```python
+from airflow import DAG
+from datetime import timedelta, datetime
+from airflow_supervisor import (
+    Supervisor,
+    SupervisorAirflowConfiguration,
+    ProgramConfiguration,
+    AirflowConfiguration,
+)
+
+
+# Create supervisor configuration
+cfg = SupervisorAirflowConfiguration(
+    airflow=AirflowConfiguration(port="*:9091"),
+    working_dir="/data/airflow/supervisor",
+    config_path="/data/airflow/supervisor/supervisor.conf",
+    program={
+        "test": ProgramConfiguration(
+            command="bash -c 'sleep 14400; exit 1'",
+        )
+    },
+)
+
+# Create DAG as normal
+with DAG(
+    dag_id="test-supervisor",
+    schedule=timedelta(days=1),
+    start_date=datetime(2024, 1, 1),
+    catchup=False,
+) as dag:
+
+    # Link supervisor config to dag
+    supervisor = Supervisor(dag=dag, cfg=cfg)
+```
+
+## Example DAG: [`airflow-config`](https://github.com/airflow-laminar/airflow-config)
+
+
+```yaml
+# @package _global_
+_target_: airflow_config.Configuration
+default_args:
+  _target_: airflow_config.DefaultArgs
+  retries: 0
+  depends_on_past: false
+all_dags:
+  _target_: airflow_config.DagArgs
+  start_date: "2024-01-01"
+  catchup: false
+extensions:
+  supervisor:
+    _target_: airflow_supervisor.SupervisorAirflowConfiguration
+    airflow:
+      _target_: airflow_supervisor.AirflowConfiguration
+      port: "*:9091"
+    working_dir: "/data/airflow/supervisor"
+    config_path: "/data/airflow/supervisor/supervisor.conf"
+    program:
+      test:
+        _target_: airflow_supervisor.ProgramConfiguration
+        command: "bash -c 'sleep 14400; exit 1'"
+```
+
+```python
+from datetime import timedelta
+from airflow_config import load_config, DAG
+from airflow_supervisor import Supervisor
+
+config = load_config(config_name="airflow")
+
+with DAG(
+    dag_id="test-supervisor",
+    schedule=timedelta(days=1),
+    config=config,
+) as dag:
+    supervisor = Supervisor(dag=dag, cfg=config.extensions["supervisor"])
+```
+
 ## How To: Use as a supervisord configuration frontend
 
 This library can be used outside airflow as a generic supervisord configuration framework, with the static typing benefits that entails. For an example, look at the [hydra configuration test](./airflow_supervisor/tests/hydra/test_hydra.py). This example generates a supervisor configuration file by composing independent hydra configs.

airflow_supervisor/airflow/local.py

@@ -13,30 +13,23 @@
 __all__ = ("Supervisor",)
 
 
-class Supervisor(DAG):
-    _supervisor_cfg: SupervisorAirflowConfiguration
-    _supervisor_kill: DAG
-    _supervisor_xmlrpc_client: SupervisorRemoteXMLRPCClient
+class Supervisor(object):
+    _dag: DAG
+    _cfg: SupervisorAirflowConfiguration
+    _kill_dag: DAG
+    _xmlrpc_client: SupervisorRemoteXMLRPCClient
 
-    def __init__(self, supervisor_cfg: SupervisorAirflowConfiguration, **kwargs):
+    def __init__(self, dag: DAG, cfg: SupervisorAirflowConfiguration, **kwargs):
         # store config
-        self._supervisor_cfg = supervisor_cfg
-        self._supervisor_xmlrpc_client = kwargs.pop(
-            "supervisor_xmlrpc_client", SupervisorRemoteXMLRPCClient(self._supervisor_cfg)
-        )
+        self._cfg = cfg
 
-        # setup role and tweak dag id
-        if "dag_id" not in kwargs:
-            kwargs["dag_id"] = list(self._supervisor_cfg.program.keys())[0]
+        # store or create client
+        self._xmlrpc_client = kwargs.pop("xmlrpc_client", SupervisorRemoteXMLRPCClient(self._cfg))
 
-        # override dag kwargs that dont make sense
-        kwargs["catchup"] = False
-        kwargs["concurrency"] = 1
-        kwargs["max_active_tasks"] = 1
-        kwargs["max_active_runs"] = 1
+        # store dag
+        self._dag = dag
 
-        # init with base DAG
-        super().__init__(**kwargs)
+        self.setup_dag()
 
         # initialize tasks
         self.initialize_tasks()
@@ -52,6 +45,13 @@ def __init__(self, supervisor_cfg: SupervisorAirflowConfiguration, **kwargs):
         # Default non running
         PythonOperator(task_id="skip", python_callable=skip_) >> self._force_kill
 
+    def setup_dag(self):
+        # override dag kwargs that dont make sense
+        self._dag.catchup = False
+        self._dag.concurrency = 1
+        self._dag.max_active_tasks = 1
+        self._dag.max_active_runs = 1
+
     def initialize_tasks(self):
         # tasks
         self._configure_supervisor = self.get_step_operator(step="configure-supervisor")
@@ -99,37 +99,35 @@ def unconfigure_supervisor(self) -> Operator:
 
     @property
     def supervisor_client(self) -> SupervisorRemoteXMLRPCClient:
-        return SupervisorRemoteXMLRPCClient(self._supervisor_cfg)
+        return SupervisorRemoteXMLRPCClient(self._cfg)
 
     def get_base_operator_kwargs(self) -> Dict:
-        return dict(dag=self)
+        return dict(dag=self._dag)
 
     def get_step_kwargs(self, step: SupervisorTaskStep) -> Dict:
         if step == "configure-supervisor":
             from .commands import write_supervisor_config
 
-            return dict(
-                python_callable=lambda: write_supervisor_config(self._supervisor_cfg, _exit=False), do_xcom_push=True
-            )
+            return dict(python_callable=lambda: write_supervisor_config(self._cfg, _exit=False), do_xcom_push=True)
         elif step == "start-supervisor":
             from .commands import start_supervisor
 
             return dict(
-                python_callable=lambda: start_supervisor(self._supervisor_cfg._pydantic_path, _exit=False),
+                python_callable=lambda: start_supervisor(self._cfg._pydantic_path, _exit=False),
                 do_xcom_push=True,
             )
         elif step == "start-programs":
             from .commands import start_programs
 
-            return dict(python_callable=lambda: start_programs(self._supervisor_cfg, _exit=False), do_xcom_push=True)
+            return dict(python_callable=lambda: start_programs(self._cfg, _exit=False), do_xcom_push=True)
         elif step == "stop-programs":
             from .commands import stop_programs
 
-            return dict(python_callable=lambda: stop_programs(self._supervisor_cfg, _exit=False), do_xcom_push=True)
+            return dict(python_callable=lambda: stop_programs(self._cfg, _exit=False), do_xcom_push=True)
         elif step == "check-programs":
             from .commands import check_programs
 
-            def _check_programs(supervisor_cfg=self._supervisor_cfg, **kwargs) -> CheckResult:
+            def _check_programs(supervisor_cfg=self._cfg, **kwargs) -> CheckResult:
                 # TODO formalize
                 if check_programs(supervisor_cfg, check_done=True, _exit=False):
                     # finish
@@ -144,35 +142,33 @@ def _check_programs(supervisor_cfg=self._supervisor_cfg, **kwargs) -> CheckResul
         elif step == "restart-programs":
             from .commands import restart_programs
 
-            return dict(python_callable=lambda: restart_programs(self._supervisor_cfg, _exit=False), do_xcom_push=True)
+            return dict(python_callable=lambda: restart_programs(self._cfg, _exit=False), do_xcom_push=True)
         elif step == "stop-supervisor":
             from .commands import stop_supervisor
 
-            return dict(python_callable=lambda: stop_supervisor(self._supervisor_cfg, _exit=False), do_xcom_push=True)
+            return dict(python_callable=lambda: stop_supervisor(self._cfg, _exit=False), do_xcom_push=True)
         elif step == "unconfigure-supervisor":
             from .commands import remove_supervisor_config
 
-            return dict(
-                python_callable=lambda: remove_supervisor_config(self._supervisor_cfg, _exit=False), do_xcom_push=True
-            )
+            return dict(python_callable=lambda: remove_supervisor_config(self._cfg, _exit=False), do_xcom_push=True)
         elif step == "force-kill":
             from .commands import kill_supervisor
 
-            return dict(python_callable=lambda: kill_supervisor(self._supervisor_cfg, _exit=False), do_xcom_push=True)
+            return dict(python_callable=lambda: kill_supervisor(self._cfg, _exit=False), do_xcom_push=True)
         raise NotImplementedError
 
     def get_step_operator(self, step: SupervisorTaskStep) -> Operator:
         if step == "check-programs":
             return HighAvailabilityOperator(
                 **{
-                    "task_id": f"{self.dag_id}-{step}",
-                    "poke_interval": self._supervisor_cfg.check_interval.total_seconds(),
-                    "timeout": self._supervisor_cfg.check_timeout.total_seconds(),
+                    "task_id": f"{self._dag.dag_id}-{step}",
+                    "poke_interval": self._cfg.check_interval.total_seconds(),
+                    "timeout": self._cfg.check_timeout.total_seconds(),
                     "mode": "poke",
                     **self.get_base_operator_kwargs(),
                     **self.get_step_kwargs(step),
                 }
             )
         return PythonOperator(
-            **{"task_id": f"{self.dag_id}-{step}", **self.get_base_operator_kwargs(), **self.get_step_kwargs(step)}
+            **{"task_id": f"{self._dag.dag_id}-{step}", **self.get_base_operator_kwargs(), **self.get_step_kwargs(step)}
         )

airflow_supervisor/airflow/ssh.py

@@ -1,6 +1,7 @@
 from shlex import quote
 from typing import Dict, List, Optional, Union
 
+from airflow.models.dag import DAG
 from airflow.models.operator import Operator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 from airflow.providers.ssh.operators.ssh import SSHOperator
@@ -17,7 +18,8 @@ class SupervisorSSH(Supervisor):
     # Mimic SSH Operator: https://airflow.apache.org/docs/apache-airflow-providers-ssh/stable/_api/airflow/providers/ssh/operators/ssh/index.html
     def __init__(
         self,
-        supervisor_cfg: SupervisorAirflowConfiguration,
+        dag: DAG,
+        cfg: SupervisorAirflowConfiguration,
         command_prefix: str = "",
         command_noescape: str = "",
         ssh_hook: Optional[SSHHook] = None,
@@ -52,7 +54,7 @@ def __init__(
             self._ssh_operator_kwargs["banner_timeout"] = banner_timeout
         if skip_on_exit_code:
             self._ssh_operator_kwargs["skip_on_exit_code"] = skip_on_exit_code
-        super().__init__(supervisor_cfg=supervisor_cfg, **kwargs)
+        super().__init__(dag=dag, cfg=cfg, **kwargs)
 
     def get_base_operator_kwargs(self) -> Dict:
         return dict(dag=self, **self._ssh_operator_kwargs)