[flyteagent] FileSensor Timeout (Remote Execution Only) (#2745)

* Add timeout for task

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* new interface

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* upload pingsu discussion

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* Sensor TimeOut

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* update

Signed-off-by: Future-Outlier <eric901201@gmail.com>

* update

Signed-off-by: Future-Outlier <eric901201@gmail.com>

---------

Signed-off-by: Future-Outlier <eric901201@gmail.com>

Author: Future-Outlier

flytekit/extend/backend/base_agent.py

@@ -21,7 +21,7 @@
 from rich.logging import RichHandler
 from rich.progress import Progress
 
-from flytekit import FlyteContext, PythonFunctionTask
+from flytekit import FlyteContext, PythonFunctionTask, logger
 from flytekit.configuration import ImageConfig, SerializationSettings
 from flytekit.core import utils
 from flytekit.core.base_task import PythonTask
@@ -285,6 +285,8 @@ def execute(self: PythonTask, **kwargs) -> LiteralMap:
         ctx = FlyteContext.current_context()
         ss = ctx.serialization_settings or SerializationSettings(ImageConfig())
         task_template = get_serializable(OrderedDict(), ss, self).template
+        if task_template.metadata.timeout:
+            logger.info("Timeout is not supported for local execution.\n" "Ignoring the timeout.")
         output_prefix = ctx.file_access.get_random_remote_directory()
 
         agent = AgentRegistry.get_agent(task_template.type, task_template.task_type_version)
@@ -336,6 +338,8 @@ def execute(self: PythonTask, **kwargs) -> LiteralMap:
         from flytekit.tools.translator import get_serializable
 
         task_template = get_serializable(OrderedDict(), ss, self).template
+        if task_template.metadata.timeout:
+            logger.info("Timeout is not supported for local execution.\n" "Ignoring the timeout.")
         self._agent = AgentRegistry.get_agent(task_template.type, task_template.task_type_version)
 
         resource_meta = local_agent_loop.run_until_complete(

flytekit/sensor/base_sensor.py

@@ -1,14 +1,15 @@
 import collections
+import datetime
 import inspect
 import typing
 from abc import abstractmethod
 from dataclasses import asdict, dataclass
-from typing import Any, Dict, Optional, TypeVar
+from typing import Any, Dict, Optional, TypeVar, Union
 
 from typing_extensions import Protocol, get_type_hints, runtime_checkable
 
 from flytekit.configuration import SerializationSettings
-from flytekit.core.base_task import PythonTask
+from flytekit.core.base_task import PythonTask, TaskMetadata
 from flytekit.core.interface import Interface
 from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin, ResourceMeta
 
@@ -50,6 +51,7 @@ class BaseSensor(AsyncAgentExecutorMixin, PythonTask):
     def __init__(
         self,
         name: str,
+        timeout: Optional[Union[datetime.timedelta, int]] = None,
         sensor_config: Optional[T] = None,
         task_type: str = "sensor",
         **kwargs,
@@ -61,11 +63,17 @@ def __init__(
             annotation = type_hints.get(k, None)
             inputs[k] = annotation
 
+        if kwargs.get("metadata", None) and timeout:
+            raise ValueError("You cannot set timeout and metadata at the same time in the sensor")
+
+        metadata = TaskMetadata(timeout=timeout)
+
         super().__init__(
             task_type=task_type,
             name=name,
             task_config=None,
             interface=Interface(inputs=inputs),
+            metadata=metadata,
             **kwargs,
         )
         self._sensor_config = sensor_config

flytekit/sensor/file_sensor.py

@@ -1,10 +1,13 @@
+import datetime
+from typing import Optional, Union
+
 from flytekit import FlyteContextManager
 from flytekit.sensor.base_sensor import BaseSensor
 
 
 class FileSensor(BaseSensor):
-    def __init__(self, name: str, **kwargs):
-        super().__init__(name=name, **kwargs)
+    def __init__(self, name: str, timeout: Optional[Union[datetime.timedelta, int]] = None, **kwargs):
+        super().__init__(name=name, timeout=timeout, **kwargs)
 
     async def poke(self, path: str) -> bool:
         file_access = FlyteContextManager.current_context().file_access

tests/flytekit/unit/sensor/test_file_sensor.py

@@ -1,7 +1,12 @@
+import datetime
 import tempfile
 
+import pytest
+
 from flytekit import task, workflow
 from flytekit.configuration import ImageConfig, SerializationSettings
+from flytekit.core.task import TaskMetadata
+from flytekit.sensor.base_sensor import BaseSensor
 from flytekit.sensor.file_sensor import FileSensor
 from tests.flytekit.unit.test_translator import default_img
 
@@ -34,3 +39,48 @@ def wf():
 
     if __name__ == "__main__":
         wf()
+
+
+def test_base_sensor_timeout_initialization():
+    # Test with no timeout
+    sensor = BaseSensor(name="test_sensor")
+    assert sensor.metadata.timeout is None
+
+    # Test with integer timeout (converted to timedelta internally)
+    sensor = BaseSensor(name="test_sensor", timeout=60)
+    assert sensor.metadata.timeout == datetime.timedelta(seconds=60)
+
+    # Test with timedelta timeout
+    timeout = datetime.timedelta(minutes=5)
+    sensor = BaseSensor(name="test_sensor", timeout=timeout)
+    assert sensor.metadata.timeout == timeout
+
+
+def test_file_sensor_timeout_initialization():
+    # Test with no timeout
+    sensor = FileSensor(name="test_file_sensor")
+    assert sensor.metadata.timeout is None
+
+    # Test with integer timeout (converted to timedelta internally)
+    sensor = FileSensor(name="test_file_sensor", timeout=60)
+    assert sensor.metadata.timeout == datetime.timedelta(seconds=60)
+
+    # Test with timedelta timeout
+    timeout = datetime.timedelta(minutes=5)
+    sensor = FileSensor(name="test_file_sensor", timeout=timeout)
+    assert sensor.metadata.timeout == timeout
+
+
+def test_agent_executor_timeout_logging():
+    # Create a sensor with timeout
+    sensor = BaseSensor(name="test_sensor", timeout=60)
+
+    # Verify the timeout is set correctly (converted to timedelta internally)
+    assert sensor.metadata.timeout == datetime.timedelta(seconds=60)
+
+def test_file_sensor_set_timeout_and_metadata_at_the_same_time():
+    timeout = datetime.timedelta(seconds=60)
+    metadata = TaskMetadata(timeout=timeout)
+
+    with pytest.raises(ValueError, match="You cannot set timeout and metadata at the same time in the sensor"):
+        FileSensor(name="test_sensor", timeout=60, metadata=metadata)