From c4ae62e44b207f1d8fd795915736030c3f3dadfa Mon Sep 17 00:00:00 2001
From: connor-mccarthy <mccarthy.connor.james@gmail.com>
Date: Mon, 18 Dec 2023 21:05:50 -0500
Subject: [PATCH] log to stdout, instead of stderr

---
 sdk/python/kfp/local/logging_utils.py        |  7 +++-
 sdk/python/kfp/local/task_dispatcher_test.py | 43 +++++++++++---------
 2 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/sdk/python/kfp/local/logging_utils.py b/sdk/python/kfp/local/logging_utils.py
index dd4d2d908435..a2b7a0eaa25f 100644
--- a/sdk/python/kfp/local/logging_utils.py
+++ b/sdk/python/kfp/local/logging_utils.py
@@ -16,6 +16,7 @@
 import contextlib
 import datetime
 import logging
+import sys
 from typing import Any, Dict, Generator, List
 
 from kfp import dsl
@@ -52,7 +53,11 @@ def local_logger_context() -> Generator[None, None, None]:
         fmt='%(asctime)s - %(levelname)s - %(message)s',
         datefmt='%H:%M:%S.%f',
     )
-    handler = logging.StreamHandler()
+    # use sys.stdout so that both inner process and outer process logs
+    # go to stdout
+    # this is needed for logs to present sequentially in a colab notebook,
+    # since stderr will print above stdout
+    handler = logging.StreamHandler(sys.stdout)
     handler.setFormatter(formatter)
     logger.handlers.clear()
     logger.addHandler(handler)
diff --git a/sdk/python/kfp/local/task_dispatcher_test.py b/sdk/python/kfp/local/task_dispatcher_test.py
index f0bacbe89551..982f27eb1f77 100644
--- a/sdk/python/kfp/local/task_dispatcher_test.py
+++ b/sdk/python/kfp/local/task_dispatcher_test.py
@@ -20,6 +20,7 @@
 should seek to minimize it.
 """
 import io
+import re
 import unittest
 from unittest import mock
 
@@ -216,11 +217,9 @@ def fail_comp():
         )
 
     @mock.patch('sys.stdout', new_callable=io.StringIO)
-    @mock.patch('sys.stderr', new_callable=io.StringIO)
     def test_user_code_no_exception_if_not_raise_on_error(
         self,
         runner,
-        mock_stderr,
         mock_stdout,
     ):
         local.init(runner=runner, raise_on_error=False)
@@ -233,7 +232,7 @@ def fail_comp():
         self.assertDictEqual(task.outputs, {})
 
         self.assertRegex(
-            mock_stderr.getvalue(),
+            mock_stdout.getvalue(),
             r"\d+:\d+:\d+\.\d+ - ERROR - Task \x1b\[96m'fail-comp'\x1b\[0m finished with status \x1b\[91mFAILURE\x1b\[0m",
         )
         self.assertIn(
@@ -242,11 +241,9 @@ def fail_comp():
         )
 
     @mock.patch('sys.stdout', new_callable=io.StringIO)
-    @mock.patch('sys.stderr', new_callable=io.StringIO)
     def test_all_logs(
         self,
         runner,
-        mock_stderr,
         mock_stdout,
     ):
         local.init(runner=runner)
@@ -262,24 +259,30 @@ def many_type_component(
 
         many_type_component(num=2)
 
-        # outer process logs in stderr
-        outer_log_regex = (
-            r"\d+:\d+:\d+\.\d+ - INFO - Executing task \x1b\[96m'many-type-component'\x1b\[0m\n"
-            + r'\d+:\d+:\d+\.\d+ - INFO - Streamed logs:\n\n' +
-            r"\d+:\d+:\d+\.\d+ - INFO - Task \x1b\[96m'many-type-component'\x1b\[0m finished with status \x1b\[92mSUCCESS\x1b\[0m\n"
-            +
-            r"\d+:\d+:\d+\.\d+ - INFO - Task \x1b\[96m'many-type-component'\x1b\[0m outputs:\n    Output: hellohello\n    model: Model\( name=model,\n                  uri=\./local_outputs/many-type-component-\d+-\d+-\d+-\d+-\d+-\d+-\d+/many-type-component/model,\n                  metadata={'foo': 'bar'} \)\n\n"
-        )
+        # inner process logs correctly nested inside outer process logs
+        outer_log_regex_sections = [
+            r"\d+:\d+:\d+\.\d+ - INFO - Executing task \x1b\[96m'many-type-component'\x1b\[0m\n",
+            r'\d+:\d+:\d+\.\d+ - INFO - Streamed logs:\n\n',
+            r'.*',
+            r'Looking for component ',
+            r'.*',
+            r'Loading KFP component ',
+            r'.*',
+            r'Got executor_input:',
+            r'.*',
+            r'Inside of my component!',
+            r'.*',
+            r'Wrote executor output file to',
+            r'.*',
+            r"\d+:\d+:\d+\.\d+ - INFO - Task \x1b\[96m'many-type-component'\x1b\[0m finished with status \x1b\[92mSUCCESS\x1b\[0m\n",
+            r"\d+:\d+:\d+\.\d+ - INFO - Task \x1b\[96m'many-type-component'\x1b\[0m outputs:\n    Output: hellohello\n    model: Model\( name=model,\n                  uri=\./local_outputs/many-type-component-\d+-\d+-\d+-\d+-\d+-\d+-\d+/many-type-component/model,\n                  metadata={'foo': 'bar'} \)\n\n",
+        ]
 
         self.assertRegex(
-            mock_stderr.getvalue(),
-            outer_log_regex,
+            mock_stdout.getvalue(),
+            # use dotall os that .* include newline characters
+            re.compile(''.join(outer_log_regex_sections), re.DOTALL),
         )
-        # inner process logs in stdout
-        self.assertIn('[KFP Executor', mock_stdout.getvalue())
-        self.assertIn('Got executor_input:', mock_stdout.getvalue())
-        self.assertIn('Inside of my component!', mock_stdout.getvalue())
-        self.assertIn('Wrote executor output file to', mock_stdout.getvalue())
 
 
 if __name__ == '__main__':