Add ability to set timeout health check in workflows

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,

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()

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

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"
  }
-```
+```
+

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

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

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

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

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

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}
+ ]
+ },
 )
 
 

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}
+ ]
+ },
 )
 
 

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)