refactor: simplify timeout handling in ContainerTask to use only timedelta

Signed-off-by: catfish <johnnyfish0221@gmail.com>

Author: Narwhal-fish

flytekit/core/container_task.py

@@ -62,7 +62,7 @@ def __init__(
         pod_template_name: Optional[str] = None,
         local_logs: bool = False,
         resources: Optional[Resources] = None,
-        timeout: Optional[typing.Union[float, int, "datetime.timedelta"]] = None,
+        timeout: Optional["datetime.timedelta"] = None,
         **kwargs,
     ):
         sec_ctx = None
@@ -286,10 +286,7 @@ def execute(self, **kwargs) -> LiteralMap:
         # Wait for the container to finish the task, with timeout if specified
         timeout_seconds = None
         if self._timeout is not None:
-            if isinstance(self._timeout, datetime.timedelta):
-                timeout_seconds = self._timeout.total_seconds()
-            else:
-                timeout_seconds = float(self._timeout)
+            timeout_seconds = self._timeout.total_seconds()
 
         if self.local_logs:
             for log in container.logs(stream=True):
@@ -341,12 +338,7 @@ def get_k8s_pod(self, settings: SerializationSettings) -> _task_model.K8sPod:
             return None
         pod_spec = _serialize_pod_spec(self.pod_template, self._get_container(settings), settings)
         if self._timeout is not None:
-            import datetime
-
-            if isinstance(self._timeout, datetime.timedelta):
-                timeout_seconds = int(self._timeout.total_seconds())
-            else:
-                timeout_seconds = int(float(self._timeout))
+            timeout_seconds = int(self._timeout.total_seconds())
             pod_spec["activeDeadlineSeconds"] = timeout_seconds
         return _task_model.K8sPod(
             pod_spec=pod_spec,

tests/flytekit/unit/core/test_container_task.py

@@ -244,29 +244,12 @@ def test_container_task_image_spec(mock_image_spec_builder):
     assert pod.pod_spec["containers"][1]["image"] == image_spec_2.image_name()
 
 def test_container_task_timeout():
-    ct_with_timeout = ContainerTask(
-        name="timeout-test",
-        input_data_dir="/var/inputs",
-        output_data_dir="/var/outputs",
-        image="busybox",
-        command=["sleep", "100"],
-        timeout=1,
-    )
-
-
-
-    with pytest.raises((docker.errors.APIError, Exception)):
-        ct_with_timeout.execute()
-
     ct_with_timedelta = ContainerTask(
         name="timedelta-timeout-test",
-        input_data_dir="/var/inputs",
-        output_data_dir="/var/outputs",
         image="busybox",
         command=["sleep", "100"],
         timeout=timedelta(seconds=1),
     )
-    
 
     with pytest.raises((docker.errors.APIError, Exception)):
         ct_with_timedelta.execute()
@@ -278,24 +261,13 @@ def test_container_task_timeout_k8s_serialization():
         containers=[], tolerations=[V1Toleration(effect="NoSchedule", key="nvidia.com/gpu", operator="Exists")]
     )
     pt = PodTemplate(pod_spec=ps, labels={"test": "timeout"})
-
-    ct_numeric = ContainerTask(
-        name="timeout-k8s-test",
-        image="busybox",
-        command=["echo", "hello"],
-        pod_template=pt,
-        timeout=60,
-    )
-
+    
     default_image = Image(name="default", fqn="docker.io/xyz", tag="some-git-hash")
     default_image_config = ImageConfig(default_image=default_image)
     default_serialization_settings = SerializationSettings(
         project="p", domain="d", version="v", image_config=default_image_config
     )
 
-    k8s_pod = ct_numeric.get_k8s_pod(default_serialization_settings)
-    assert k8s_pod.pod_spec["activeDeadlineSeconds"] == 60
-
     ct_timedelta = ContainerTask(
         name="timeout-k8s-timedelta-test",
         image="busybox",
@@ -308,16 +280,13 @@ def test_container_task_timeout_k8s_serialization():
     assert k8s_pod_timedelta.pod_spec["activeDeadlineSeconds"] == 120
 
 def test_container_task_no_timeout():
-
-    ct = ContainerTask(
+    ct_timedelta = ContainerTask(
         name="no-timeout-task",
         input_data_dir="/var/inputs",
         output_data_dir="/var/outputs",
         image="busybox",
         command=["sleep", "1"],  
-        timeout=500,
+        timeout=timedelta(seconds=500),
     )
 
-
-    
-    ct.execute()
+    ct_timedelta.execute()