Add ability to set timeout health check in workflows (#58)

* Add ability to set timeout health check in workflows --------- Co-authored-by: Brent Johnson <[email protected]>
Nike-Inc · Nov 10, 2023 · ae84716 · ae84716
1 parent b2af63d
commit ae84716
Show file tree

Hide file tree

Showing 12 changed files with 65 additions and 4 deletions.
diff --git a/brickflow/codegen/databricks_bundle.py b/brickflow/codegen/databricks_bundle.py
@@ -581,6 +581,7 @@ def proj_to_bundle(self) -> DatabricksAssetBundles:
                 tasks=tasks,
                 git_source=git_conf,
                 tags=workflow.tags,
+                health=workflow.health,
                 job_clusters=[JobsJobClusters(**c) for c in workflow_clusters],
                 schedule=self.workflow_obj_to_schedule(workflow),
                 max_concurrent_runs=workflow.max_concurrent_runs,

diff --git a/brickflow/engine/task.py b/brickflow/engine/task.py
@@ -32,7 +32,11 @@
     BrickflowProjectDeploymentSettings,
     get_brickflow_version,
 )
-from brickflow.bundles.model import JobsTasksNotebookTask, JobsTasksNotificationSettings
+from brickflow.bundles.model import (
+    JobsTasksNotebookTask,
+    JobsTasksNotificationSettings,
+    JobsTasksHealthRules,
+)
 from brickflow.cli.projects import DEFAULT_BRICKFLOW_VERSION_MODE
 from brickflow.context import (
     BrickflowBuiltInTaskVariables,
@@ -479,6 +483,7 @@ class Task:
     task_settings: Optional[TaskSettings] = None
     custom_execute_callback: Optional[Callable] = None
     ensure_brickflow_plugins: bool = False
+    health: Optional[List[JobsTasksHealthRules]] = None
 
     def __post_init__(self) -> None:
         self.is_valid_task_signature()

diff --git a/brickflow/engine/workflow.py b/brickflow/engine/workflow.py
@@ -11,6 +11,7 @@
     JobsWebhookNotifications,
     JobsNotificationSettings,
     JobsTrigger,
+    JobsHealthRules,
 )
 from brickflow.context import BrickflowInternalVariables
 from brickflow.engine import ROOT_NODE
@@ -116,6 +117,8 @@ class Workflow:
     schedule_pause_status: str = "UNPAUSED"
     default_cluster: Optional[Cluster] = None
     clusters: List[Cluster] = field(default_factory=lambda: [])
+
+    health: Optional[List[JobsHealthRules]] = None
     default_task_settings: TaskSettings = TaskSettings()
     email_notifications: Optional[WorkflowEmailNotifications] = None
     webhook_notifications: Optional[WorkflowWebhookNotifications] = None

diff --git a/docs/workflows.md b/docs/workflows.md
@@ -20,7 +20,8 @@ wf = Workflow(  # (1)!
         email_notifications=EmailNotifications(
             on_start=["[email protected]"],
             on_success=["[email protected]"],
-            on_failure=["[email protected]"]
+            on_failure=["[email protected]"],
+            on_duration_warning_threshold_exceeded=["[email protected]"]
         ),
         timeout_seconds=timedelta(hours=2).seconds
     ),
@@ -44,6 +45,11 @@ wf = Workflow(  # (1)!
         "catalog": "development",
         "database": "your_database"
     },
+    health = { # (16)!
+        "metric": "RUN_DURATION_SECONDS",
+        "op": "GREATER_THAN",
+        "value": 7200
+    }
 )
 
 
@@ -67,6 +73,7 @@ def task_function(*, test="var"):
 13. Define the common task parameters that can be used in all the tasks
 14. Define a workflow task and associate it to the workflow
 15. Define the schedule pause status. It is defaulted to "UNPAUSED"
+16. Define health check condition that triggers duration warning threshold exceeded notifications
 
 ### Clusters
 
@@ -207,7 +214,8 @@ default_task_settings=TaskSettings(
    email_notifications=EmailNotifications(
       on_start=["[email protected]"],
       on_success=["[email protected]"],
-      on_failure=["[email protected]"]
+      on_failure=["[email protected]"],
+      on_duration_warning_threshold_exceeded=["[email protected]"]      
    ),
    timeout_seconds=timedelta(hours=2).seconds,
    max_retries=2,
@@ -249,4 +257,5 @@ common_task_parameters={
         "catalog": "development",
         "database": "your_database"
     }
-```
+```
+
diff --git a/tests/codegen/expected_bundles/dev_bundle_monorepo.yml b/tests/codegen/expected_bundles/dev_bundle_monorepo.yml
@@ -8,6 +8,11 @@ environments:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          health: 
+            rules:
+              - metric: "RUN_DURATION_SECONDS"
+                op: "GREATER_THAN" 
+                value: 7200.0
           git_source:
             git_commit: a
             git_provider: github

diff --git a/tests/codegen/expected_bundles/dev_bundle_polyrepo.yml b/tests/codegen/expected_bundles/dev_bundle_polyrepo.yml
@@ -8,6 +8,11 @@ environments:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          health: 
+            rules:
+              - metric: "RUN_DURATION_SECONDS"
+                op: "GREATER_THAN" 
+                value: 7200.0
           git_source:
             git_commit: a
             git_provider: github

diff --git a/tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml b/tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml
@@ -6,6 +6,7 @@ environments:
       jobs:
         some_wf:
           email_notifications: null
+          health: null
           git_source:
             git_commit: a
             git_provider: github
@@ -64,6 +65,11 @@ environments:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          health: 
+            rules:
+              - metric: "RUN_DURATION_SECONDS"
+                op: "GREATER_THAN" 
+                value: 7200.0          
           git_source:
             git_commit: a
             git_provider: github

diff --git a/tests/codegen/expected_bundles/local_bundle.yml b/tests/codegen/expected_bundles/local_bundle.yml
@@ -8,6 +8,11 @@ environments:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          health: 
+            rules:
+              - metric: "RUN_DURATION_SECONDS"
+                op: "GREATER_THAN" 
+                value: 7200.0
           git_source: null
           job_clusters: []
           max_concurrent_runs: 1.0

diff --git a/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml b/tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml
@@ -8,6 +8,11 @@ environments:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          health: 
+            rules:
+              - metric: "RUN_DURATION_SECONDS"
+                op: "GREATER_THAN" 
+                value: 7200.0
           git_source: null
           job_clusters: []
           max_concurrent_runs: 1.0

diff --git a/tests/codegen/sample_workflow.py b/tests/codegen/sample_workflow.py
@@ -23,6 +23,11 @@
     run_as_user="[email protected]",
     tags={"test": "test2"},
     common_task_parameters={"all_tasks1": "test", "all_tasks3": "123"},  # type: ignore
+    health={
+        "rules": [
+            {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200.0}
+        ]
+    },
 )
 
 

diff --git a/tests/engine/sample_workflow.py b/tests/engine/sample_workflow.py
@@ -19,6 +19,11 @@
     ),
     tags={"test": "test2"},
     common_task_parameters={"all_tasks1": "test", "all_tasks3": "123"},  # type: ignore
+    health={
+        "rules": [
+            {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200}
+        ]
+    },
 )
 
 

diff --git a/tests/engine/test_workflow.py b/tests/engine/test_workflow.py
@@ -182,6 +182,13 @@ def test_tags(self):
     def test_default_task_settings(self):
         assert wf.default_task_settings is not None
 
+    def test_health_settings(self):
+        assert wf.health == {
+            "rules": [
+                {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200}
+            ]
+        }
+
     def test_user(self):
         principal = "[email protected]"
         u = User(principal)