RD-WATCH/Smartflow Integration

Dockerfile

@@ -74,6 +74,8 @@ RUN mkdir /app/django/src \
  && touch /app/django/src/rdwatch/__init__.py \
  && mkdir /app/django/src/rdwatch_scoring \
  && touch /app/django/src/rdwatch_scoring/__init__.py \
+ && mkdir /app/django/src/rdwatch_smartflow \
+ && touch /app/django/src/rdwatch_smartflow/__init__.py \
  && touch /app/django/README.md \
  && poetry install --only main
 
@@ -88,6 +90,8 @@ RUN mkdir /app/django/src \
  && touch /app/django/src/rdwatch/__init__.py \
  && mkdir /app/django/src/rdwatch_scoring \
  && touch /app/django/src/rdwatch_scoring/__init__.py \
+ && mkdir /app/django/src/rdwatch_smartflow \
+ && touch /app/django/src/rdwatch_smartflow/__init__.py \
  && touch /app/django/README.md \
  && poetry install --with dev
 # Copy git metadata to enable display of version information

dev/.env.docker-compose

@@ -8,3 +8,6 @@ RDWATCH_MINIO_STORAGE_SECRET_KEY=secretkey
 RDWATCH_STORAGE_BUCKET_NAME=rdwatch
 AWS_REQUEST_PAYER=requester
 RDWATCH_POSTGRESQL_SCORING_URI="postgresql://scoring:secretkey@scoredb:5432/scoring"
+RDWATCH_SMARTFLOW_USERNAME="rdwatch"
+RDWATCH_SMARTFLOW_PASSWORD="rdwatch"
+RDWATCH_SMARTFLOW_URL="http://airflow:8080"

dev/.env.docker-compose-native

@@ -11,3 +11,6 @@ RDWATCH_POSTGRESQL_URI="postgresql://rdwatch:secretkey@localhost:5432/rdwatch"
 RDWATCH_REDIS_URI="redis://localhost:6379"
 RDWATCH_CELERY_BROKER_URL="redis://localhost:6379"
 RDWATCH_POSTGRESQL_SCORING_URI="postgresql://scoring:secretkey@localhost:5433/scoring"
+RDWATCH_SMARTFLOW_USERNAME="rdwatch"
+RDWATCH_SMARTFLOW_PASSWORD="rdwatch"
+RDWATCH_SMARTFLOW_URL="http://localhost:8002"

dev/airflow.Dockerfile

@@ -0,0 +1,11 @@
+FROM apache/airflow:slim-latest-python3.12
+
+RUN airflow db migrate
+
+RUN airflow users create \
+    --username rdwatch \
+    --password rdwatch \
+    --firstname rdwatch \
+    --lastname admin \
+    --role Admin \
+    --email admin@localhost

dev/airflow_sample_dag.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+
+from airflow import DAG
+from airflow.decorators import task
+from airflow.models import Param
+
+with DAG(
+    dag_id="RD-WATCH-AIRFLOW-DEMO-DAG",
+    description="Test DAG",
+    params={
+        "region_id": Param(
+            default="BR_R002", type="string", pattern=r"^[A-Z]{2}_[RCST]\d{3}$"
+        ),
+        "model_run_title": Param(default="test_run", type="string"),
+    },
+    start_date=datetime(2022, 3, 1),
+    catchup=False,
+    schedule_interval=None,
+    max_active_runs=1,
+    default_view="grid",
+    tags=["demo", "watch", "ta2", "rgd", "random"],
+) as dag1:
+
+    @task
+    def print_region(**context):
+        region_id = context["params"]["region_id"]
+        print(f"region_id: {region_id}")
+
+    @task
+    def print_model_run_title(**context):
+        model_run_title = context["params"]["model_run_title"]
+        print(f"model_run_title: {model_run_title}")
+
+    print_region() >> print_model_run_title()

django/pyproject.toml

@@ -22,6 +22,7 @@ classifiers = [
 packages = [
   { include = "rdwatch", from = "src" },
   { include = "rdwatch_scoring", from = "src"},
+  { include = "rdwatch_smartflow", from = "src"},
 ]
 
 [tool.poetry-dynamic-versioning]
@@ -30,7 +31,7 @@ vcs = "git"
 metadata = true
 
 [tool.poetry-dynamic-versioning.substitution]
-files = ["src/rdwatch/__init__.py", "src/rdwatch_scoring/__init__.py"]
+files = ["src/rdwatch/__init__.py", "src/rdwatch_scoring/__init__.py", "src/rdwatch_smartflow/__init__.py"]
 
 [tool.poetry.dependencies]
 python = ">=3.11.8,<4"
@@ -57,6 +58,8 @@ pystac-client = "^0.7.6"
 more-itertools = "^10.2.0"
 torch = {version = "^2.2.2+cpu", source = "pytorch"}
 torchvision = {version = "^0.17.2+cpu", source = "pytorch"}
+apache-airflow-client = "^2.9.0"
+beautifulsoup4 = "^4.12.3"
 
 [tool.poetry.group.dev.dependencies]
 django-stubs = "^4.2.7"
@@ -91,7 +94,7 @@ skip-string-normalization = true
 
 [tool.isort]
 profile = "black"
-known_first_party = ["rdwatch", "rdwatch_scoring"]
+known_first_party = ["rdwatch", "rdwatch_scoring", "rdwatch_smartflow"]
 known_django=["django"]
 sections=["FUTURE", "STDLIB", "THIRDPARTY", "DJANGO", "FIRSTPARTY", "LOCALFOLDER"]
 

django/src/rdwatch/server/settings.py

@@ -48,6 +48,15 @@ def INSTALLED_APPS(self):
         ]
         if 'RDWATCH_POSTGRESQL_SCORING_URI' in os.environ:
             base_applications.append('rdwatch_scoring')
+        if all(
+            os.environ.get(key)
+            for key in [
+                'RDWATCH_SMARTFLOW_URL',
+                'RDWATCH_SMARTFLOW_USERNAME',
+                'RDWATCH_SMARTFLOW_PASSWORD',
+            ]
+        ):
+            base_applications.append('rdwatch_smartflow')
         return base_applications
 
     MIDDLEWARE = [
@@ -132,6 +141,23 @@ def DATABASES(self):
         environ_required=True, environ_prefix=_ENVIRON_PREFIX
     )
 
+    SMARTFLOW_URL = values.Value(
+        None,
+        environ_required=False,
+        environ_prefix=_ENVIRON_PREFIX,
+    )
+
+    SMARTFLOW_USERNAME = values.Value(
+        None,
+        environ_required=False,
+        environ_prefix=_ENVIRON_PREFIX,
+    )
+    SMARTFLOW_PASSWORD = values.Value(
+        None,
+        environ_required=False,
+        environ_prefix=_ENVIRON_PREFIX,
+    )
+
     # django-celery-results configuration
     CELERY_RESULT_BACKEND = 'django-db'
     CELERY_CACHE_BACKEND = 'django-cache'

django/src/rdwatch/server/urls.py

@@ -11,6 +11,8 @@
 # Conditionally add the scoring URLs if the scoring app is installed
 if 'rdwatch_scoring' in apps.app_configs.keys():
     urlpatterns.append(path('api/scoring/', include('rdwatch_scoring.urls')))
+if 'rdwatch_smartflow' in apps.app_configs.keys():
+    urlpatterns.append(path('api/smartflow/', include('rdwatch_smartflow.urls')))
 
 if settings.DEBUG:
     try:

django/src/rdwatch/views/server_status.py

@@ -1,10 +1,13 @@
 import datetime
 import socket
+from typing import Any
 
 from ninja import Field, Router, Schema
 
 from django.http import HttpRequest
 
+from rdwatch_smartflow.utils import SmartFlowClient
+
 router = Router()
 
 SERVER_INSTANCE_EPOCH = datetime.datetime.now()
@@ -26,6 +29,7 @@ class ServerStatusSchema(Schema):
     hostname: str
     ip: str
     rdwatch_version: str
+    smartflow: dict[str, Any] | None
 
 
 @router.get('/', response=ServerStatusSchema)
@@ -35,6 +39,16 @@ def get_status(request: HttpRequest):
     uptime = datetime.datetime.now() - SERVER_INSTANCE_EPOCH
     hostname = socket.gethostname()
     ip = socket.gethostbyname(hostname)
+
+    try:
+        smartflow_status = SmartFlowClient().get_health().to_dict()
+    except Exception:
+        smartflow_status = None
+
     return ServerStatusSchema(
-        uptime=uptime, hostname=hostname, ip=ip, rdwatch_version=api.version
+        uptime=uptime,
+        hostname=hostname,
+        ip=ip,
+        rdwatch_version=api.version,
+        smartflow=smartflow_status,
     )

django/src/rdwatch_smartflow/api.py

@@ -0,0 +1,7 @@
+from ninja import NinjaAPI
+
+from . import views
+
+api = NinjaAPI(urls_namespace='smartflow')
+
+api.add_router('/', views.router)

django/src/rdwatch_smartflow/urls.py

@@ -0,0 +1,7 @@
+from django.urls import path
+
+from .api import api
+
+urlpatterns = [
+    path('', api.urls),
+]

django/src/rdwatch_smartflow/utils.py

@@ -0,0 +1,93 @@
+from functools import cached_property
+from typing import Any
+
+import airflow_client.client
+import requests
+from airflow_client.client.api import dag_api, dag_run_api, monitoring_api
+from airflow_client.client.models import (
+    DAGCollection,
+    DAGRun,
+    DAGRunCollection,
+    HealthInfo,
+)
+from bs4 import BeautifulSoup
+
+from django.conf import settings
+
+
+class SmartFlowClient:
+    @cached_property
+    def _session_token(self) -> str:
+        session = requests.Session()
+        smartflow_url: str = settings.SMARTFLOW_URL
+        username: str = settings.SMARTFLOW_USERNAME
+        password: str = settings.SMARTFLOW_PASSWORD
+
+        res = session.get(smartflow_url)
+        soup = BeautifulSoup(res.text, 'html.parser')
+        csrf_token = soup.find('input', {'name': 'csrf_token'})['value']
+
+        login_url = f'{smartflow_url}/login/'
+        login_form_data = {
+            'csrf_token': csrf_token,
+            'username': username,
+            'password': password,
+        }
+        login_headers = {'referer': login_url}
+        res = session.post(
+            url=login_url,
+            data=login_form_data,
+            headers=login_headers,
+        )
+        res.raise_for_status()
+        return session.cookies.get('session')
+
+    @property
+    def _client(self):
+        host: str = f'{settings.SMARTFLOW_URL}/api/v1'
+
+        configuration = airflow_client.client.Configuration(
+            host=host,
+            # api_key={'session': self._session_token}
+        )
+        with airflow_client.client.ApiClient(
+            configuration=configuration,
+            header_name='cookie',
+            header_value=f'session={self._session_token}',
+        ) as api_client:
+            return api_client
+
+    @property
+    def _monitoring_api(self):
+        return monitoring_api.MonitoringApi(self._client)
+
+    @property
+    def _dag_api(self):
+        return dag_api.DAGApi(self._client)
+
+    @property
+    def _dag_run_api(self):
+        return dag_run_api.DAGRunApi(self._client)
+
+    def get_health(self) -> HealthInfo:
+        return self._monitoring_api.get_health()
+
+    def list_dags(self, **kwargs) -> DAGCollection:
+        return self._dag_api.get_dags(**kwargs)
+
+    def list_dag_runs(self, dag_id: str, **kwargs) -> DAGRunCollection:
+        return self._dag_run_api.get_dag_runs(dag_id=dag_id, **kwargs)
+
+    def create_dag_run(
+        self, dag_id: str, dag_run_title: str, conf: dict[str, Any]
+    ) -> DAGRun:
+        """
+        Parameters:
+            dag_id: the id of the DAG to run
+            dag_run_title: the title to give this new DAG run
+        """
+        dag_run = DAGRun(conf=conf, dag_run_id=dag_run_title)
+        return self._dag_run_api.post_dag_run(dag_id=dag_id, dag_run=dag_run)
+
+    def get_dag_run(self, dag_id: str, dag_run_id: str) -> DAGRun:
+        return self._dag_run_api.get_dag_run(dag_id=dag_id, dag_run_id=dag_run_id)