From b112e41dea470b7bafabaa7b6c33155977b408ef Mon Sep 17 00:00:00 2001
From: aditya7302 <aditya7302@gmail.com>
Date: Wed, 19 Jun 2024 22:49:43 +0530
Subject: [PATCH] Boolean flag to show local container logs to the terminal

Signed-off-by: aditya7302 <aditya7302@gmail.com>
---
 flytekit/core/container_task.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/flytekit/core/container_task.py b/flytekit/core/container_task.py
index 66fe522c07..e35ff7ed7a 100644
--- a/flytekit/core/container_task.py
+++ b/flytekit/core/container_task.py
@@ -59,6 +59,7 @@ def __init__(
         secret_requests: Optional[List[Secret]] = None,
         pod_template: Optional["PodTemplate"] = None,
         pod_template_name: Optional[str] = None,
+        local_logs: bool = False,
         **kwargs,
     ):
         sec_ctx = None
@@ -93,6 +94,7 @@ def __init__(
             requests=requests if requests else Resources(), limits=limits if limits else Resources()
         )
         self.pod_template = pod_template
+        self.local_logs = local_logs
 
     @property
     def resources(self) -> ResourceSpec:
@@ -249,6 +251,11 @@ def execute(self, **kwargs) -> LiteralMap:
         )
         # Wait for the container to finish the task
         # TODO: Add a 'timeout' parameter to control the max wait time for the container to finish the task.
+
+        if self.local_logs:
+            for log in container.logs(stream=True):
+                print(log.strip())
+                
         container.wait()
 
         output_dict = self._get_output_dict(output_directory)