debugs

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

flytekit/core/context_manager.py

@@ -984,7 +984,6 @@ def initialize():
         import sys
         import threading
         print(f"{threading.current_thread().name}: name: {__name__}")
-        print(f"{threading.current_thread().name}: {sys.modules}")
         print(f"{threading.current_thread().name}: {sys.path}")
         print(f"{threading.current_thread().name}: {sys.thread_info}")
 

flytekit/core/worker_queue.py

@@ -174,12 +174,13 @@ def __init__(self, remote: FlyteRemote, ss: SerializationSettings, tag: str, roo
         # Import this to ensure context is loaded... python is reloading this module because its in a different thread
         FlyteContextManager.current_context()
 
+        self.stopping_condition = threading.Event()
         # Things for actually kicking off and monitoring
         self.__runner_thread: threading.Thread = threading.Thread(
             target=self._execute, daemon=True, name="controller-thread"
         )
         self.__runner_thread.start()
-        atexit.register(self._close)
+        atexit.register(self._close, event=self.stopping_condition, runner=self.__runner_thread)
 
         # Executions should be tracked in the following way:
         #  a) you should be able to list by label, all executions generated by the current eager task,
@@ -200,8 +201,11 @@ def __init__(self, remote: FlyteRemote, ss: SerializationSettings, tag: str, roo
         self.tag = tag
         self.root_tag = root_tag
 
-    def _close(self) -> None:
-        self.__runner_thread.join()
+    @staticmethod
+    def _close(event: threading.Event, runner: threading.Thread) -> None:
+        event.set()
+        printf("----------------------------- in close ------------------------")
+        runner.join()
 
     def reconcile_one(self, update: Update):
         """
@@ -286,6 +290,9 @@ def _poll(self) -> None:
         FlyteContextManager.current_context()
         # This needs to be a while loop that runs forever,
         while True:
+            if self.stopping_condition.is_set():
+                print(f"STOPPING RUNNER THREAD!")
+                break
             # Gather all items that need processing
             update_items = self._get_update_items()
             print(f"DEBUG: poll loop {update_items}")
@@ -394,13 +401,14 @@ async def add(self, entity: RunnableEntity, input_kwargs: dict[str, typing.Any])
         # wait for it to finish one way or another
         while True:
             developer_logger.debug(f"Watching id {id(i)}")
+            print(f"In Add loop - Watching id {id(i)}")
             if i.status == ItemStatus.SUCCESS:
                 return i.result
             elif i.status == ItemStatus.FAILED:
                 assert i.error is not None
                 raise i.error
             else:
-                await asyncio.sleep(1)  # Small delay to avoid busy-waiting
+                await asyncio.sleep(2)  # Small delay to avoid busy-waiting
 
     def render_html(self) -> str:
         """Render the callstack as a deck presentation to be shown after eager workflow execution."""