Repository: PrefectHQ/prefect

Issue:
Logs configured in tasks with `get_run_logger` using `DaskTaskRunner` don't make it to the Prefect 2.0 backend
## Opened from the [Prefect Public Slack Community](https://prefect.io/slack)

**marcin.grzybowski**: Hi, another strange question, when I use DaskTaskRunner i see my print logs but not logs from get_run_logger, when I do not specify runner then I see both ...
I run it locally using python3

**anna**: it's a general issue with `distributed` that is described <https://discourse.prefect.io/t/why-are-the-dask-logs-not-being-shown-in-the-prefect-ui/75 |here> - if you want to, you can open an issue so that we can track it but it looks like <https://github.com/dask/distributed/issues/2033 |it's still open> on distributed repo since 2018

**anna**: <@ULVA73B9P> open "Logs configured in tasks with `get_run_logger` using `DaskTaskRunner` don't make it to the Prefect 2.0 backend"

Original thread can be found [here](https://prefect-community.slack.com/archives/CL09KU1K7/p1654014315773189?thread_ts=1654014315.773189&cid=CL09KU1K7).




Discussion:
problem description:

On the Prefect side, the native Python logger gets serialized and then sent to Dask workers. When this Prefect logger gets deserialized, it loses the configuration attached to it, thus the logs are not directed to the Prefect 2.0 backend.


Related discussion
https://github.com/PrefectHQ/prefect-dask/issues/10
just mentioning there is a related PR here https://github.com/dask/distributed/pull/7276 
Hi,
I just looked through the different issues and PRs regarding login from dask distributed and found that this [PR](https://github.com/dask/distributed/pull/7297) got merged 2 days ago. Would this have a positive impact on the Prefect side? 

I am hoping that this will resolve the logging issue I encounter here:
```
from prefect import flow, get_run_logger, task
from prefect_dask import DaskTaskRunner

@task()
def a_task(i):
    logger = get_run_logger()
    logger.info(f"I am logging from a task. {i}")

@flow(
    name="Logger Test",
    task_runner=DaskTaskRunner(),
)
def flow():
    logger = get_run_logger()
    logger.info("Flow starts!")

    a_task.map(range(3))

    logger.info("Flow ends!")
```

Which produces only the following output:
```
16:05:00.779 | INFO    | prefect.engine - Created flow run 'hairy-owl' for flow 'Logger Test'
16:05:00.779 | INFO    | prefect.task_runner.dask - Creating a new Dask cluster with `distributed.deploy.local.LocalCluster`
16:05:02.965 | INFO    | prefect.task_runner.dask - The Dask dashboard is available at http://127.0.0.1:8787/status
16:05:04.322 | INFO    | Flow run 'hairy-owl' - Flow starts!
16:05:04.327 | INFO    | Flow run 'hairy-owl' - Flow ends!
16:05:04.546 | INFO    | Flow run 'hairy-owl' - Created task run 'a_task-800f1d79-2' for task 'a_task'
16:05:04.930 | INFO    | Flow run 'hairy-owl' - Submitted task run 'a_task-800f1d79-2' for execution.
16:05:04.937 | INFO    | Flow run 'hairy-owl' - Created task run 'a_task-800f1d79-1' for task 'a_task'
16:05:04.946 | INFO    | Flow run 'hairy-owl' - Submitted task run 'a_task-800f1d79-1' for execution.
16:05:04.953 | INFO    | Flow run 'hairy-owl' - Created task run 'a_task-800f1d79-0' for task 'a_task'
16:05:04.958 | INFO    | Flow run 'hairy-owl' - Submitted task run 'a_task-800f1d79-0' for execution.
16:05:06.653 | INFO    | Flow run 'hairy-owl' - Finished in state Completed('All states completed.')
```

The expected output would be (obtained with the default task runner):
```
16:08:10.958 | INFO    | prefect.engine - Created flow run 'eccentric-fennec' for flow 'Logger Test'
16:08:12.168 | INFO    | Flow run 'eccentric-fennec' - Flow starts!
16:08:12.171 | INFO    | Flow run 'eccentric-fennec' - Flow ends!
16:08:12.374 | INFO    | Flow run 'eccentric-fennec' - Created task run 'a_task-800f1d79-2' for task 'a_task'
16:08:12.375 | INFO    | Flow run 'eccentric-fennec' - Submitted task run 'a_task-800f1d79-2' for execution.
16:08:12.380 | INFO    | Flow run 'eccentric-fennec' - Created task run 'a_task-800f1d79-1' for task 'a_task'
16:08:12.380 | INFO    | Flow run 'eccentric-fennec' - Submitted task run 'a_task-800f1d79-1' for execution.
16:08:12.384 | INFO    | Flow run 'eccentric-fennec' - Created task run 'a_task-800f1d79-0' for task 'a_task'
16:08:12.385 | INFO    | Flow run 'eccentric-fennec' - Submitted task run 'a_task-800f1d79-0' for execution.
16:08:12.956 | INFO    | Task run 'a_task-800f1d79-2' - I am logging from a task. 2
16:08:12.962 | INFO    | Task run 'a_task-800f1d79-0' - I am logging from a task. 0
16:08:13.164 | INFO    | Task run 'a_task-800f1d79-0' - Finished in state Completed()
16:08:13.169 | INFO    | Task run 'a_task-800f1d79-1' - I am logging from a task. 1
16:08:13.246 | INFO    | Task run 'a_task-800f1d79-2' - Finished in state Completed()
16:08:13.352 | INFO    | Task run 'a_task-800f1d79-1' - Finished in state Completed()
16:08:13.677 | INFO    | Flow run 'eccentric-fennec' - Finished in state Completed('All states completed.')
```

All the best and thank you for developing Prefect!

Resolving patch:
diff --git a/src/prefect/__init__.py b/src/prefect/__init__.py
--- a/src/prefect/__init__.py
+++ b/src/prefect/__init__.py
@@ -20,6 +20,7 @@
 
 del _version, pathlib
 
+
 # Import user-facing API
 from prefect.states import State
 from prefect.logging import get_run_logger
@@ -63,6 +64,12 @@
 prefect.plugins.load_prefect_collections()
 
 
+# Configure logging
+import prefect.logging.configuration
+
+prefect.logging.configuration.setup_logging()
+
+
 # Ensure moved names are accessible at old locations
 import prefect.client
 
diff --git a/src/prefect/logging/configuration.py b/src/prefect/logging/configuration.py
--- a/src/prefect/logging/configuration.py
+++ b/src/prefect/logging/configuration.py
@@ -60,7 +60,12 @@ def load_logging_config(path: Path) -> dict:
     return flatdict_to_dict(flat_config)
 
 
-def setup_logging() -> None:
+def setup_logging() -> dict:
+    """
+    Sets up logging.
+
+    Returns the config used.
+    """
     global PROCESS_LOGGING_CONFIG
 
     # If the user has specified a logging path and it exists we will ignore the
@@ -73,15 +78,8 @@ def setup_logging() -> None:
         )
     )
 
-    if PROCESS_LOGGING_CONFIG:
-        # Do not allow repeated configuration calls, only warn if the config differs
-        if PROCESS_LOGGING_CONFIG != config:
-            warnings.warn(
-                "Logging can only be setup once per process, the new logging config "
-                f"will be ignored.",
-                stacklevel=2,
-            )
-        return
+    # Perform an incremental update if setup has already been run
+    config.setdefault("incremental", bool(PROCESS_LOGGING_CONFIG))
 
     logging.config.dictConfig(config)
 
@@ -91,9 +89,12 @@ def setup_logging() -> None:
     for logger_name in PREFECT_LOGGING_EXTRA_LOGGERS.value():
         logger = logging.getLogger(logger_name)
         for handler in extra_config.handlers:
-            logger.addHandler(handler)
+            if not config["incremental"]:
+                logger.addHandler(handler)
             if logger.level == logging.NOTSET:
                 logger.setLevel(extra_config.level)
             logger.propagate = extra_config.propagate
 
     PROCESS_LOGGING_CONFIG = config
+
+    return config
diff --git a/src/prefect/logging/handlers.py b/src/prefect/logging/handlers.py
--- a/src/prefect/logging/handlers.py
+++ b/src/prefect/logging/handlers.py
@@ -15,7 +15,7 @@
 from rich.theme import Theme
 
 import prefect.context
-from prefect.client import get_client
+from prefect.client.orion import get_client
 from prefect.exceptions import MissingContextError
 from prefect.logging.highlighters import PrefectConsoleHighlighter
 from prefect.orion.schemas.actions import LogCreate


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -35,6 +35,7 @@
 
 import prefect
 import prefect.settings
+from prefect.logging import get_logger
 from prefect.logging.configuration import setup_logging
 from prefect.settings import (
     PREFECT_API_URL,
@@ -451,3 +452,20 @@ def reset_registered_blocks():
 
     registry.clear()
     registry.update(before)
+
+
+@pytest.fixture
+def caplog(caplog):
+    """
+    Overrides caplog to apply to all of our loggers that do not propagate and
+    consequently would not be captured by caplog.
+    """
+
+    config = setup_logging()
+
+    for name, logger_config in config["loggers"].items():
+        if not logger_config.get("propagate", True):
+            logger = get_logger(name)
+            logger.handlers.append(caplog.handler)
+
+    yield caplog
diff --git a/tests/test_engine.py b/tests/test_engine.py
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -72,24 +72,6 @@ def flow_for_tests(dog: str, cat: int):
     return flow_for_tests
 
 
-@pytest.fixture
-def flow_run_caplog(caplog):
-    """
-    Capture logging from flow runs to ensure messages are correct.
-    """
-    import logging
-
-    logger = logging.getLogger("prefect.flow_runs")
-    logger2 = logging.getLogger("prefect")
-    logger.propagate = True
-    logger2.propagate = True
-
-    try:
-        yield caplog
-    finally:
-        logger.propagate = False
-
-
 @pytest.fixture
 async def get_flow_run_context(orion_client, result_factory, local_filesystem):
     partial_ctx = PartialModel(FlowRunContext)
diff --git a/tests/test_logging.py b/tests/test_logging.py
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -103,27 +103,19 @@ def test_setup_logging_uses_default_path(tmp_path, dictConfigMock):
         {PREFECT_LOGGING_SETTINGS_PATH: tmp_path.joinpath("does-not-exist.yaml")}
     ):
         expected_config = load_logging_config(DEFAULT_LOGGING_SETTINGS_PATH)
+        expected_config["incremental"] = False
         setup_logging()
 
     dictConfigMock.assert_called_once_with(expected_config)
 
 
-def test_setup_logging_allows_repeated_calls(dictConfigMock):
+def test_setup_logging_sets_incremental_on_repeated_calls(dictConfigMock):
     setup_logging()
-    dictConfigMock.assert_called_once()
+    assert dictConfigMock.call_count == 1
     setup_logging()
-    dictConfigMock.assert_called_once()
-
-
-def test_setup_logging_warns_on_repeated_calls_with_new_settings(dictConfigMock):
-    setup_logging()
-    dictConfigMock.assert_called_once()
-    with pytest.warns(
-        UserWarning, match="only be setup once per process.* will be ignored"
-    ):
-        with temporary_settings({PREFECT_LOGGING_LEVEL: "ERROR"}):
-            setup_logging()
-    dictConfigMock.assert_called_once()
+    assert dictConfigMock.call_count == 2
+    assert dictConfigMock.mock_calls[0][1][0]["incremental"] == False
+    assert dictConfigMock.mock_calls[1][1][0]["incremental"] == True
 
 
 def test_setup_logging_uses_settings_path_if_exists(tmp_path, dictConfigMock):
@@ -134,6 +126,7 @@ def test_setup_logging_uses_settings_path_if_exists(tmp_path, dictConfigMock):
 
         setup_logging()
         expected_config = load_logging_config(tmp_path.joinpath("exists.yaml"))
+        expected_config["incremental"] = False
 
     dictConfigMock.assert_called_once_with(expected_config)
 
@@ -146,6 +139,8 @@ def test_setup_logging_uses_env_var_overrides(tmp_path, dictConfigMock, monkeypa
         expected_config = load_logging_config(DEFAULT_LOGGING_SETTINGS_PATH)
     env = {}
 
+    expected_config["incremental"] = False
+
     # Test setting a value for a simple key
     env["PREFECT_LOGGING_HANDLERS_ORION_LEVEL"] = "ORION_LEVEL_VAL"
     expected_config["handlers"]["orion"]["level"] = "ORION_LEVEL_VAL"
@@ -239,6 +234,7 @@ def test_get_logger_does_not_duplicate_prefect_prefix():
 def test_default_level_is_applied_to_interpolated_yaml_values(dictConfigMock):
     with temporary_settings({PREFECT_LOGGING_LEVEL: "WARNING"}):
         expected_config = load_logging_config(DEFAULT_LOGGING_SETTINGS_PATH)
+        expected_config["incremental"] = False
 
         assert expected_config["loggers"]["prefect"]["level"] == "WARNING"
         assert expected_config["loggers"]["prefect.extra"]["level"] == "WARNING"
@@ -1095,22 +1091,8 @@ def test_init_override_kwargs(self):
         assert handler.level == logging.DEBUG
 
 
-@pytest.fixture
-def flow_run_caplog(caplog):
-    """
-    Capture logging from flow runs to ensure messages are correct.
-    """
-    logger = logging.getLogger("prefect.flow_runs")
-    logger.propagate = True
-
-    try:
-        yield caplog
-    finally:
-        logger.propagate = False
-
-
-def test_log_in_flow(flow_run_caplog):
-    msg = "10:21:34.114 | INFO    | Flow run 'polite-jackal' - Finished in state Completed()"
+def test_log_in_flow(caplog):
+    msg = "Hello world!"
 
     @flow
     def test_flow():
@@ -1118,30 +1100,17 @@ def test_flow():
         logger.warning(msg)
 
     test_flow()
-    for record in flow_run_caplog.records:
+
+    for record in caplog.records:
         if record.msg == msg:
             assert record.levelno == logging.WARNING
             break
     else:
-        raise AssertionError(f"{msg} was not found in records")
-
-
-@pytest.fixture
-def task_run_caplog(caplog):
-    """
-    Capture logging from task runs to ensure messages are correct.
-    """
-    logger = logging.getLogger("prefect.task_runs")
-    logger.propagate = True
-
-    try:
-        yield caplog
-    finally:
-        logger.propagate = False
+        raise AssertionError(f"{msg} was not found in records: {caplog.records}")
 
 
-def test_log_in_task(task_run_caplog):
-    msg = "10:21:34.114 | INFO    | Flow run 'polite-jackal' - Finished in state Completed()"
+def test_log_in_task(caplog):
+    msg = "Hello world!"
 
     @task
     def test_task():
@@ -1153,7 +1122,7 @@ def test_flow():
         test_task()
 
     test_flow()
-    for record in task_run_caplog.records:
+    for record in caplog.records:
         if record.msg == msg:
             assert record.levelno == logging.WARNING
             break