Add timeout parameter to workflow class (#83)

* Add timeout parameter to workflow class

* Add warning if task timeout exceeds workflow timeout

* Change timeout_seconds from 'float' to 'int'

---------

Co-authored-by: Brent Johnson <[email protected]>

brickflow/bundles/model.py

@@ -1525,7 +1525,7 @@ class Config:
         ...,
         description='A unique name for the task. This field is used to refer to this task from other tasks.\nThis field is required and must be unique within its parent job.\nOn Update or Reset, this field is used to reference the tasks to be updated or reset.',
     )
-    timeout_seconds: Optional[float] = Field(
+    timeout_seconds: Optional[int] = Field(
         None,
         description='An optional timeout applied to each run of this job task. A value of `0` means no timeout.',
     )
@@ -1661,7 +1661,7 @@ class Config:
     tasks: Optional[List[JobsTasks]] = Field(
         None, description='A list of task specifications to be executed by this job.'
     )
-    timeout_seconds: Optional[float] = Field(
+    timeout_seconds: Optional[int] = Field(
         None,
         description='An optional timeout applied to each run of this job. A value of `0` means no timeout.',
     )

brickflow/codegen/databricks_bundle.py

@@ -594,6 +594,7 @@ def proj_to_bundle(self) -> DatabricksAssetBundles:
                 permissions=self.workflow_obj_to_permissions(
                     workflow
                 ),  # will be none if not set
+                timeout_seconds=workflow.timeout_seconds,
                 email_notifications=workflow.email_notifications,
                 notification_settings=workflow.notification_settings,
                 webhook_notifications=workflow.webhook_notifications,

brickflow/engine/workflow.py

@@ -1,4 +1,5 @@
 import abc
+import logging
 import functools
 from dataclasses import dataclass, field
 from typing import Callable, List, Optional, Dict, Union, Iterator, Any
@@ -117,8 +118,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
+    timeout_seconds: Optional[int] = None
     default_task_settings: TaskSettings = TaskSettings()
     email_notifications: Optional[WorkflowEmailNotifications] = None
     webhook_notifications: Optional[WorkflowWebhookNotifications] = None
@@ -255,6 +256,13 @@ def pop_task(self, task_id: str) -> None:
     def task_exists(self, task_id: str) -> bool:
         return task_id in self.tasks
 
+    def log_timeout_warning(self, task_settings: TaskSettings) -> bool:
+        if task_settings is not None and self.timeout_seconds is not None:
+            if task_settings.timeout_seconds is not None:
+                if task_settings.timeout_seconds > self.timeout_seconds:
+                    return True
+        return False
+
     def _set_active_task(self, task_id: str) -> None:
         self.active_task = task_id
 
@@ -301,6 +309,11 @@ def _add_task(
                 "Some how default cluster wasnt set please raise a github issue."
             )
 
+        if self.log_timeout_warning(task_settings):  # type: ignore
+            logging.warning(
+                "Task timeout_seconds should not exceed workflow timeout_seconds",
+            )
+
         _libraries = libraries or [] + self.libraries
         _depends_on = (
             [depends_on]

docs/workflows.md

@@ -74,6 +74,7 @@ def task_function(*, test="var"):
 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
+17. Define timeout_seconds check condition that triggers workflow failure if duration exceeds threshold
 
 ### Clusters
 

tests/codegen/expected_bundles/dev_bundle_monorepo.yml

@@ -8,6 +8,7 @@ targets:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          timeout_seconds: 42
           health: 
             rules:
               - metric: "RUN_DURATION_SECONDS"

tests/codegen/expected_bundles/dev_bundle_polyrepo.yml

@@ -8,6 +8,7 @@ targets:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          timeout_seconds: 42
           health: 
             rules:
               - metric: "RUN_DURATION_SECONDS"

tests/codegen/expected_bundles/dev_bundle_polyrepo_with_auto_libs.yml

@@ -65,6 +65,7 @@ targets:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          timeout_seconds: 42
           health: 
             rules:
               - metric: "RUN_DURATION_SECONDS"

tests/codegen/expected_bundles/local_bundle.yml

@@ -8,6 +8,7 @@ targets:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          timeout_seconds: 42
           health: 
             rules:
               - metric: "RUN_DURATION_SECONDS"

tests/codegen/expected_bundles/local_bundle_prefix_suffix.yml

@@ -8,6 +8,7 @@ targets:
           email_notifications: null
           notification_settings: null
           webhook_notifications: null
+          timeout_seconds: 42
           health: 
             rules:
               - metric: "RUN_DURATION_SECONDS"

tests/codegen/sample_workflow.py

@@ -28,6 +28,7 @@
             {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200.0}
         ]
     },
+    timeout_seconds=42,
 )
 
 

tests/engine/sample_workflow.py

@@ -24,6 +24,7 @@
             {"metric": "RUN_DURATION_SECONDS", "op": "GREATER_THAN", "value": 7200}
         ]
     },
+    timeout_seconds=42,
 )
 
 

tests/engine/test_task.py

@@ -1,3 +1,4 @@
+import datetime
 from collections import namedtuple
 from unittest.mock import Mock, patch
 
@@ -97,7 +98,6 @@ def test_custom_task_params(self):
             "all_tasks3": "123",
             "test": "var",
         }
-
         wf.common_task_parameters = {}
         assert wf.get_task(task_function_nokwargs.__name__).custom_task_parameters == {}
 
@@ -323,6 +323,25 @@ def test_task_settings_merge(self):
             "retry_on_timeout": default_bool,
         }
 
+    def test_task_settings_small_timeout(self):
+        small_to = TaskSettings(
+            timeout_seconds=30,
+        )
+        print(small_to.timeout_seconds)
+        assert not wf.log_timeout_warning(small_to)
+
+    def test_task_settings_big_timeout_warning(self):
+        big_to = TaskSettings(
+            timeout_seconds=datetime.timedelta(hours=0.5).seconds,
+        )
+        print(big_to.timeout_seconds)
+        assert wf.log_timeout_warning(big_to)
+
+    def test_task_settings_no_timeout_warning(self):
+        no_to = TaskSettings()
+        print(no_to.timeout_seconds)
+        assert not wf.log_timeout_warning(no_to)
+
     def test_task_libraries(self):
         s3_path = "s3://somepath-in-s3"
         repo = "somerepo"

tests/engine/test_workflow.py

@@ -189,6 +189,9 @@ def test_health_settings(self):
             ]
         }
 
+    def test_timeout_seconds(self):
+        assert wf.timeout_seconds == 42
+
     def test_user(self):
         principal = "[email protected]"
         u = User(principal)