Repository: PrefectHQ/prefect

Issue:
extra_loggers doesn't work for "processes" scheduler in LocalDaskExecutor
## Description
When the `extra_loggers` config is set and the processes scheduler is used in `LocalDaskExecutor`, log messages don't ultimately make it to the flow.

## Expected Behavior
Log outputs in tasks are expected to make it to the main process and to stdout (when running the flow directly from the flow script).

## Reproduction
```
import sys
import logging

from prefect import task, Flow
from prefect.engine.executors import LocalDaskExecutor
from prefect.environments import LocalEnvironment


LOGGER = logging.getLogger("my-named-logger")
LOGGER.addHandler(logging.StreamHandler(sys.stdout))
LOGGER.setLevel(logging.DEBUG)

@task()
def log_message():
    LOGGER.info("Hello world!")
    LOGGER.error("Foo!")


with Flow('flow-with-extra-logger') as flow:
    log_message()
flow.environment = LocalEnvironment(executor=LocalDaskExecutor(scheduler="processes"))

flow.run()
```

## Environment
The local agent can be used with `extra_loggers` set with an env variable:
```
prefect agent local start --env PREFECT__LOGGING__EXTRA_LOGGERS="['my-named-logger']"
```


Discussion:
This issue is the result of a discussion on slack with @madkinsz and @jcrist which can be found here:
https://prefect-community.slack.com/archives/CL09KU1K7/p1607631235485700

Resolving patch:
diff --git a/src/prefect/executors/dask.py b/src/prefect/executors/dask.py
--- a/src/prefect/executors/dask.py
+++ b/src/prefect/executors/dask.py
@@ -397,6 +397,15 @@ def wait(self, futures: Any) -> Any:
         return self.client.gather(futures)
 
 
+def _multiprocessing_pool_initializer() -> None:
+    """Initialize a process used in a `multiprocssing.Pool`.
+
+    Ensures the standard atexit handlers are run."""
+    import signal
+
+    signal.signal(signal.SIGTERM, lambda signum, frame: sys.exit())
+
+
 class LocalDaskExecutor(Executor):
     """
     An executor that runs all functions locally using `dask` and a configurable
@@ -436,7 +445,13 @@ def __setstate__(self, state: dict) -> None:
 
     def _interrupt_pool(self) -> None:
         """Interrupt all tasks in the backing `pool`, if any."""
-        if self.scheduler == "threads" and self._pool is not None:
+        if self._pool is None:
+            return
+
+        # Terminate the pool
+        self._pool.terminate()
+
+        if self.scheduler == "threads":
             # `ThreadPool.terminate()` doesn't stop running tasks, only
             # prevents new tasks from running. In CPython we can attempt to
             # raise an exception in all threads. This exception will be raised
@@ -502,7 +517,9 @@ def _posttask(self, key, value, dsk, state, id):  # type: ignore
                     from dask.multiprocessing import get_context
 
                     context = get_context()
-                    self._pool = context.Pool(num_workers)
+                    self._pool = context.Pool(
+                        num_workers, initializer=_multiprocessing_pool_initializer
+                    )
             try:
                 exiting_early = False
                 yield
@@ -511,9 +528,10 @@ def _posttask(self, key, value, dsk, state, id):  # type: ignore
                 raise
             finally:
                 if self._pool is not None:
-                    self._pool.terminate()
                     if exiting_early:
                         self._interrupt_pool()
+                    else:
+                        self._pool.close()
                     self._pool.join()
                     self._pool = None
 
diff --git a/src/prefect/utilities/logging.py b/src/prefect/utilities/logging.py
--- a/src/prefect/utilities/logging.py
+++ b/src/prefect/utilities/logging.py
@@ -9,13 +9,13 @@
 When running locally, log levels and message formatting are set via your Prefect configuration file.
 """
 import atexit
-import json
 import logging
 import sys
 import threading
 import time
+import warnings
 from queue import Empty, Queue
-from typing import Any
+from typing import Any, List, Optional
 
 import pendulum
 
@@ -32,147 +32,150 @@
     "task_run_id",
 )
 
+MAX_LOG_LENGTH = 1_000_000  # 1 MB - max length of a single log message
+MAX_BATCH_LOG_LENGTH = 20_000_000  # 20 MB - max total batch size for log messages
 
-class CloudHandler(logging.StreamHandler):
-    def __init__(self) -> None:
-        super().__init__(sys.stdout)
-        self.client = None
-        self.logger = logging.getLogger("CloudHandler")
-        handler = logging.StreamHandler(sys.stdout)
-        formatter = logging.Formatter(
-            context.config.logging.format, context.config.logging.datefmt
-        )
-        formatter.converter = time.gmtime  # type: ignore
-        handler.setFormatter(formatter)
-        self.logger.addHandler(handler)
-        self.logger.setLevel(context.config.logging.level)
-
-    @property
-    def queue(self) -> Queue:
-        if not hasattr(self, "_queue"):
-            self._queue = Queue()  # type: Queue
-            self._flush = False
-            self.start()
-        return self._queue
 
-    def flush(self) -> None:
-        self._flush = True
-        if self.client is not None:
-            self.batch_upload()
-            self._thread.join()
+class LogManager:
+    """A global log manager for managing all logs to be sent to Prefect"""
 
-    def batch_upload(self) -> None:
-        logs = []
+    def __init__(self) -> None:
+        self.queue = Queue()  # type: Queue[dict]
+        self.pending_logs = []  # type: List[dict]
+        self.thread = None  # type: Optional[threading.Thread]
+        self.client = None  # type: Optional[prefect.Client]
+        self.pending_length = 0
+        self._stopped = False
+
+    def ensure_started(self) -> None:
+        """Ensure the log manager is started"""
+        if self.thread is None:
+            self.client = prefect.Client()
+            self.logging_period = context.config.cloud.logging_heartbeat
+            self.thread = threading.Thread(
+                target=self._write_logs_loop,
+                name="prefect-log-manager",
+                daemon=True,
+            )
+            self.thread.start()
+            atexit.register(self._on_shutdown)
+
+    def _on_shutdown(self) -> None:
+        """Called via atexit, flushes all logs and stops the background thread"""
+        # Sometimes a signal can hit the process at shutdown multiple times,
+        # interrupting an active shutdown hook. To give a better chance of the
+        # shutdown hook succeeding, we retry a few times, ignoring extra
+        # `SystemExit` exceptions raised here. Note that this won't prevent
+        # shutdown (the interpreter is already shutting down regardless).
+        for _ in range(3):
+            try:
+                self.stop()
+                return
+            except SystemExit:
+                pass
+
+    def stop(self) -> None:
+        """Flush all logs and stop the background thread"""
+        if self.thread is not None:
+            self._stopped = True
+            self.thread.join()
+            self._write_logs()
+            self.thread = None
+            self.client = None
+
+    def _write_logs_loop(self) -> None:
+        """Runs in a background thread, uploads logs periodically in a loop"""
+        while not self._stopped:
+            cont = True
+            while cont:
+                cont = self._write_logs()
+            time.sleep(self.logging_period)
+
+    def _write_logs(self) -> bool:
+        """Upload a single batch of logs.
+
+        Returns:
+            - bool: Whether `_write_logs` should be called again this round.
+        """
+        assert self.client is not None  # mypy
+
+        # Read all logs from the queue into the `pending_logs` list. This
+        # is stored on the manager to ensure that logs aren't dropped in
+        # the case of an upload error, and will be retried later. This
+        # could be due to an api error, or due to a shutdown signal
+        # interrupting a log upload.
+        #
+        # We batch uploads with a max total length to prevent uploading too
+        # large a payload at once. This call will continue to loop until
+        # the queue is empty or an error occurs on upload (usually only one
+        # iteration is sufficient)
+        cont = True
         try:
-            while True:
-                log = self.queue.get(False)
-                logs.append(log)
+            while self.pending_length < MAX_BATCH_LOG_LENGTH:
+                log = self.queue.get_nowait()
+                self.pending_length += len(log.get("message", ""))
+                self.pending_logs.append(log)
         except Empty:
-            pass
+            cont = False
 
-        if logs:
+        if self.pending_logs:
             try:
-                assert self.client is not None
-                self.client.write_run_logs(logs)
+                self.client.write_run_logs(self.pending_logs)
+                self.pending_logs = []
+                self.pending_length = 0
             except Exception as exc:
-                message = "Failed to write log with error: {}".format(str(exc))
-                self.logger.critical(message)
-
-                # Attempt to write batch error log otherwise log invalid cloud communication
-                try:
-                    assert self.client is not None
-                    self.client.write_run_logs([self._make_error_log(message)])
-                except Exception:
-                    self.logger.critical("Unable to write logs to Prefect Cloud")
-
-    def _monitor(self) -> None:
-        while not self._flush:
-            self.batch_upload()
-            time.sleep(self.heartbeat)
-
-    def __del__(self) -> None:
-        if hasattr(self, "_thread"):
-            self.flush()
-            atexit.unregister(self.flush)
-
-    def start(self) -> None:
-        if not hasattr(self, "_thread"):
-            self.heartbeat = context.config.cloud.logging_heartbeat
-            self._thread = t = threading.Thread(
-                target=self._monitor, name="PrefectCloudLoggingThread"
-            )
-            t.daemon = True
-            t.start()
-            atexit.register(self.flush)
+                # An error occurred on upload, warn and exit the loop (will
+                # retry later)
+                warnings.warn(f"Failed to write logs with error: {exc!r}")
+                cont = False
+        return cont
 
-    def put(self, log: dict) -> None:
-        try:
-            json.dumps(log)  # make sure the payload is serializable
-            self.queue.put(log)
-        except TypeError as exc:
-            message = "Failed to write log with error: {}".format(str(exc))
-            self.logger.critical(message)
+    def enqueue(self, message: dict) -> None:
+        """Enqueue a new log message to be uploaded.
 
-            self.queue.put(self._make_error_log(message))
+        Args:
+            - message (dict): a log message to upload.
+        """
+        self.ensure_started()
+        self.queue.put(message)
 
-    def emit(self, record) -> None:  # type: ignore
-        # if we shouldn't log to cloud, don't emit
-        if not prefect.context.config.logging.log_to_cloud:
-            return
 
-        try:
-            from prefect.client import Client
+class CloudHandler(logging.Handler):
+    """A handler for sending logs to the prefect API"""
 
-            if self.client is None:
-                self.client = Client()  # type: ignore
+    def emit(self, record: logging.LogRecord) -> None:  # type: ignore
+        """Emit a new log"""
+        # if we shouldn't log to cloud, don't emit
+        if not context.config.logging.log_to_cloud:
+            return
 
-            assert isinstance(self.client, Client)  # mypy assert
+        # ensures emitted logs respect configured logging level
+        config_level = getattr(logging, context.config.logging.level, logging.INFO)
 
-            record_dict = record.__dict__.copy()
+        if record.levelno < config_level:
+            return
 
-            # ensures emitted logs respect configured logging level
-            config_level = getattr(
-                logging, prefect.context.config.logging.level, logging.INFO
+        msg = self.format(record)
+        if len(msg) > MAX_LOG_LENGTH:
+            get_logger("prefect.logging").warning(
+                "Received a log message of %d bytes, exceeding the limit of %d. "
+                "The output will be truncated",
+                len(msg),
+                MAX_LOG_LENGTH,
             )
+            msg = msg[:MAX_LOG_LENGTH]
 
-            if record_dict["levelno"] < config_level:
-                return
-
-            # remove potentially non-json serializable formatting args
-            record_dict.pop("args", None)
-
-            log = dict()
-            log["flow_run_id"] = prefect.context.get("flow_run_id", None)
-            log["task_run_id"] = prefect.context.get("task_run_id", None)
-            log["timestamp"] = pendulum.from_timestamp(
-                record_dict.pop("created", time.time())
-            ).isoformat()
-            log["name"] = record_dict.pop("name", None)
-            log["message"] = record_dict.pop("message", None)
-            log["level"] = record_dict.pop("levelname", None)
-
-            if record_dict.get("exc_text") is not None:
-                log["message"] += "\n" + record_dict.pop("exc_text", "")
-                record_dict.pop("exc_info", None)
-
-            log["info"] = record_dict
-            self.put(log)
-        except Exception as exc:
-            message = "Failed to write log with error: {}".format(str(exc))
-            self.logger.critical(message)
-
-            self.put(self._make_error_log(message))
-
-    def _make_error_log(self, message: str) -> dict:
-        log = dict()
-        log["flow_run_id"] = prefect.context.get("flow_run_id", None)
-        log["timestamp"] = pendulum.from_timestamp(time.time()).isoformat()
-        log["name"] = self.logger.name
-        log["message"] = message
-        log["level"] = "CRITICAL"
-        log["info"] = {}
-
-        return log
+        log = {
+            "flow_run_id": context.get("flow_run_id"),
+            "task_run_id": context.get("task_run_id"),
+            "timestamp": pendulum.from_timestamp(
+                getattr(record, "created", None) or time.time()
+            ).isoformat(),
+            "name": getattr(record, "name", None),
+            "level": getattr(record, "levelname", None),
+            "message": msg,
+        }
+        LOG_MANAGER.enqueue(log)
 
 
 def _log_record_context_injector(*args: Any, **kwargs: Any) -> logging.LogRecord:
@@ -192,7 +195,7 @@ def _log_record_context_injector(*args: Any, **kwargs: Any) -> logging.LogRecord
     additional_attrs = context.config.logging.get("log_attributes", [])
 
     for attr in PREFECT_LOG_RECORD_ATTRIBUTES + tuple(additional_attrs):
-        value = prefect.context.get(attr, None)
+        value = context.get(attr, None)
         if value or attr in additional_attrs:
             setattr(record, attr, value)
 
@@ -220,13 +223,7 @@ def _create_logger(name: str) -> logging.Logger:
     handler.setFormatter(formatter)
     logger.addHandler(handler)
     logger.setLevel(context.config.logging.level)
-
-    # we set the cloud handler to DEBUG level
-    # but the handler itself will dynamically respond
-    # to the configured level in the emit() method call
-    cloud_handler = CloudHandler()
-    cloud_handler.setLevel("DEBUG")
-    logger.addHandler(cloud_handler)
+    logger.addHandler(CloudHandler())
     return logger
 
 
@@ -292,6 +289,9 @@ def get_logger(name: str = None) -> logging.Logger:
         return prefect_logger.getChild(name)
 
 
+LOG_MANAGER = LogManager()
+
+
 class RedirectToLog:
     """
     Custom redirect of stdout messages to logs


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -20,9 +20,11 @@
 
 
 @pytest.fixture(autouse=True)
-def logging_heartbeat():
-    with configuration.set_temporary_config({"cloud.logging_heartbeat": 0.15}):
-        yield
+def no_cloud_logs(monkeypatch):
+    """Prevent cloud logging from doing anything actually sending requests to
+    Prefect, regardless of status of `logging.log_to_cloud`. Test checking
+    cloud logging works explicitly may need to override this mock."""
+    monkeypatch.setattr("prefect.utilities.logging.LOG_MANAGER.enqueue", MagicMock())
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/engine/cloud/test_cloud_flow_runner.py b/tests/engine/cloud/test_cloud_flow_runner.py
--- a/tests/engine/cloud/test_cloud_flow_runner.py
+++ b/tests/engine/cloud/test_cloud_flow_runner.py
@@ -504,12 +504,7 @@ def my_run_task(*args, **kwargs):
         ):
             return run_task(*args, **kwargs)
 
-    calls = []
-
     class Client(MagicMock):
-        def write_run_logs(self, *args, **kwargs):
-            calls.append(args)
-
         def set_task_run_state(self, *args, **kwargs):
             return kwargs.get("state")
 
@@ -551,19 +546,21 @@ def log_stuff():
     assert flow_state.is_successful()
     assert flow_state.result[log_stuff].result == (42, "original")
 
-    time.sleep(0.75)
-    logs = [log for call in calls for log in call[0]]
+    # LOG_MANAGER.enqueue is mocked out in `no_cloud_logs` in conftest.py
+    logs = [
+        c[0][0] for c in prefect.utilities.logging.LOG_MANAGER.enqueue.call_args_list
+    ]
     assert len(logs) >= 5  # actual number of logs
 
-    loggers = [c["name"] for c in logs]
-    assert set(loggers) == {
+    loggers = {l["name"] for l in logs}
+    assert loggers == {
         "prefect.CloudTaskRunner",
         "prefect.CloudFlowRunner",
         "prefect.log_stuff",
     }
 
-    task_run_ids = [c["task_run_id"] for c in logs if c["task_run_id"]]
-    assert set(task_run_ids) == {"TESTME"}
+    task_run_ids = {c["task_run_id"] for c in logs if c["task_run_id"]}
+    assert task_run_ids == {"TESTME"}
 
 
 class TestCloudFlowRunnerQueuedState:
diff --git a/tests/engine/test_flow_runner.py b/tests/engine/test_flow_runner.py
--- a/tests/engine/test_flow_runner.py
+++ b/tests/engine/test_flow_runner.py
@@ -1352,66 +1352,6 @@ def log_stuff():
     assert len([r for r in caplog.records if r.levelname == "CRITICAL"]) == 1
 
 
-def test_task_runners_submitted_to_remote_machines_respect_original_config(monkeypatch):
-    """
-    This test is meant to simulate the behavior of running a Cloud Flow against an external
-    cluster which has _not_ been configured for Prefect.  The idea is that the configuration
-    settings which were present on the original machine are respected in the remote job, reflected
-    here by having the CloudHandler called during logging and the special values present in context.
-    """
-
-    from prefect.engine.flow_runner import run_task
-
-    def my_run_task(*args, **kwargs):
-        with prefect.utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": False, "cloud.auth_token": ""}
-        ):
-            return run_task(*args, **kwargs)
-
-    calls = []
-
-    class Client:
-        def write_run_logs(self, *args, **kwargs):
-            calls.append(args)
-
-    monkeypatch.setattr("prefect.engine.flow_runner.run_task", my_run_task)
-    monkeypatch.setattr("prefect.client.Client", Client)
-
-    @prefect.task
-    def log_stuff():
-        logger = prefect.context.get("logger")
-        logger.critical("important log right here")
-        return (
-            prefect.context.config.special_key,
-            prefect.context.config.cloud.auth_token,
-        )
-
-    with prefect.utilities.configuration.set_temporary_config(
-        {
-            "logging.log_to_cloud": True,
-            "special_key": 42,
-            "cloud.auth_token": "original",
-        }
-    ):
-        # captures config at init
-        flow = Flow("test", tasks=[log_stuff])
-        flow_state = flow.run(task_contexts={log_stuff: dict(special_key=99)})
-
-    assert flow_state.is_successful()
-    assert flow_state.result[log_stuff].result == (42, "original")
-
-    time.sleep(0.75)
-    assert len(calls) >= 1
-    assert len([log for call in calls for log in call[0]]) == 5  # actual number of logs
-
-    loggers = [log["name"] for call in calls for log in call[0]]
-    assert set(loggers) == {
-        "prefect.TaskRunner",
-        "prefect.FlowRunner",
-        "prefect.log_stuff",
-    }
-
-
 def test_constant_tasks_arent_submitted(caplog):
     calls = []
 
diff --git a/tests/utilities/test_logging.py b/tests/utilities/test_logging.py
--- a/tests/utilities/test_logging.py
+++ b/tests/utilities/test_logging.py
@@ -1,6 +1,5 @@
 import contextlib
 import datetime
-import json
 import logging
 import time
 from unittest.mock import MagicMock
@@ -8,322 +7,237 @@
 import click
 import pytest
 
+import prefect
 from prefect import context, utilities
+from prefect.utilities.logging import CloudHandler, LogManager
 
 
-def test_root_logger_level_responds_to_config():
-    try:
-        with utilities.configuration.set_temporary_config({"logging.level": "DEBUG"}):
-            assert (
-                utilities.logging.configure_logging(testing=True).level == logging.DEBUG
-            )
-
-        with utilities.configuration.set_temporary_config({"logging.level": "WARNING"}):
-            assert (
-                utilities.logging.configure_logging(testing=True).level
-                == logging.WARNING
-            )
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
-
-
-@pytest.mark.parametrize("datefmt", ["%Y", "%Y -- %D"])
-def test_root_logger_datefmt_responds_to_config(caplog, datefmt):
+@pytest.fixture
+def log_manager(monkeypatch):
+    log_manager = LogManager()
+    Client = MagicMock()
+    monkeypatch.setattr(log_manager, "enqueue", MagicMock(wraps=log_manager.enqueue))
+    monkeypatch.setattr(prefect, "Client", Client)
+    monkeypatch.setattr(utilities.logging, "LOG_MANAGER", log_manager)
     try:
-        with utilities.configuration.set_temporary_config({"logging.datefmt": datefmt}):
-            logger = utilities.logging.configure_logging(testing=True)
-            logger.error("badness")
-            logs = [r for r in caplog.records if r.levelname == "ERROR"]
-            assert logs[0].asctime == datetime.datetime.now().strftime(datefmt)
+        yield log_manager
     finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
+        log_manager.stop()
 
 
-def test_remote_handler_is_configured_for_cloud():
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger = utilities.logging.configure_logging(testing=True)
-            assert hasattr(logger.handlers[-1], "client")
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
-
-
-def test_diagnostic_logger_has_no_cloud_handler():
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger = utilities.logging.create_diagnostic_logger(name="diagnostic-test")
-            assert logger.handlers
-            assert all(
-                [
-                    not isinstance(h, utilities.logging.CloudHandler)
-                    for h in logger.handlers
-                ]
-            )
-    finally:
-        # reset logger
-        logger = utilities.logging.create_diagnostic_logger(name="diagnostic-test")
-        logger.handlers = []
-
-
-def test_remote_handler_captures_errors_and_logs_them(caplog, monkeypatch):
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True, "cloud.auth_token": None}
-        ):
-            logger = utilities.logging.configure_logging(testing=True)
-            assert hasattr(logger.handlers[-1], "client")
-            logger.handlers[-1].client = MagicMock()
-            child_logger = logger.getChild("sub-test")
-            child_logger.critical("this should raise an error in the handler")
-
-            time.sleep(1.5)  # wait for batch upload to occur
-            critical_logs = [r for r in caplog.records if r.levelname == "CRITICAL"]
-            assert len(critical_logs) == 2
-
-            assert (
-                "Failed to write log with error"
-                in [
-                    log.message for log in critical_logs if log.name == "CloudHandler"
-                ].pop()
-            )
-    finally:
-        # reset root_logger
+@pytest.fixture
+def logger(log_manager):
+    # Clean logger for every test run
+    with utilities.configuration.set_temporary_config(
+        {
+            "logging.level": "INFO",
+            "cloud.logging_heartbeat": 0.25,
+            "logging.log_to_cloud": True,
+        }
+    ):
         logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
+        yield logger
+        logger.handlers.clear()
 
 
-def test_remote_handler_captures_tracebacks(caplog, monkeypatch):
-    monkeypatch.setattr("prefect.client.Client", MagicMock)
-    client = MagicMock()
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger = utilities.logging.configure_logging(testing=True)
-            assert hasattr(logger.handlers[-1], "client")
-            logger.handlers[-1].client = client
-
-            child_logger = logger.getChild("sub-test")
-            try:
-                ## informative comment
-                1 + "2"
-            except:
-                child_logger.exception("unexpected error")
-
-            time.sleep(0.75)
-            error_logs = [r for r in caplog.records if r.levelname == "ERROR"]
-            assert len(error_logs) >= 1
-
-            cloud_logs = client.write_run_logs.call_args[0][0]
-            assert len(cloud_logs) == 1
-            logged_msg = cloud_logs[0]["message"]
-            assert "TypeError" in logged_msg
-            assert '1 + "2"' in logged_msg
-            assert "unexpected error" in logged_msg
+def test_root_logger_level_responds_to_config():
+    with utilities.configuration.set_temporary_config({"logging.level": "DEBUG"}):
+        assert utilities.logging.configure_logging(testing=True).level == logging.DEBUG
 
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
+    with utilities.configuration.set_temporary_config({"logging.level": "WARNING"}):
+        assert (
+            utilities.logging.configure_logging(testing=True).level == logging.WARNING
+        )
 
 
-def test_cloud_handler_formats_messages_and_removes_args(caplog, monkeypatch):
-    monkeypatch.setattr("prefect.client.Client", MagicMock)
-    client = MagicMock()
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger = utilities.logging.configure_logging(testing=True)
-            assert hasattr(logger.handlers[-1], "client")
-            logger.handlers[-1].client = client
-
-            child_logger = logger.getChild("sub-test")
-            child_logger.info("Here's a number: %d", 42)
-
-            time.sleep(0.75)
-
-            cloud_logs = client.write_run_logs.call_args[0][0]
-            assert len(cloud_logs) == 1
-            assert cloud_logs[0]["message"] == "Here's a number: 42"
-            assert "args" not in cloud_logs[0]["info"]
-    finally:
-        # reset root_logger
+@pytest.mark.parametrize("datefmt", ["%Y", "%Y -- %D"])
+def test_root_logger_datefmt_responds_to_config(caplog, datefmt):
+    with utilities.configuration.set_temporary_config({"logging.datefmt": datefmt}):
         logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
+        logger.error("badness")
+        logs = [r for r in caplog.records if r.levelname == "ERROR"]
+        assert logs[0].asctime == datetime.datetime.now().strftime(datefmt)
 
 
-def test_cloud_handler_emit_responds_to_config(caplog, monkeypatch):
-    monkeypatch.setattr("prefect.client.Client", MagicMock)
-    client = MagicMock()
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True, "logging.level": "DEBUG"}
-        ):
-            logger = utilities.logging.configure_logging(testing=True)
-            assert hasattr(logger.handlers[-1], "client")
-            logger.handlers[-1].client = client
-
-            child_logger = logger.getChild("sub-test")
+def test_root_logger_has_cloud_handler(logger):
+    assert logger.handlers
+    assert any(isinstance(h, CloudHandler) for h in logger.handlers)
 
-            with utilities.configuration.set_temporary_config(
-                {"logging.log_to_cloud": True, "logging.level": "INFO"}
-            ):
-                child_logger.debug("heres a log")
 
-            time.sleep(0.75)
-
-            assert not client.write_run_logs.called
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
+def test_diagnostic_logger_has_no_cloud_handler():
+    logger = utilities.logging.create_diagnostic_logger(name="diagnostic-test")
+    assert logger.handlers
+    assert not any(isinstance(h, CloudHandler) for h in logger.handlers)
+
+
+def test_cloud_handler_emit_noop_if_cloud_logging_disabled(logger, log_manager):
+    with utilities.configuration.set_temporary_config({"logging.log_to_cloud": False}):
+        logger.info("testing")
+    assert not log_manager.enqueue.called
+    assert log_manager.client is None
+    assert log_manager.thread is None
+
+
+def test_cloud_handler_emit_noop_if_below_log_level(logger, log_manager):
+    logger.debug("testing")
+    assert not log_manager.enqueue.called
+    assert log_manager.client is None
+    assert log_manager.thread is None
+
+
+def test_cloud_handler_emit_noop_if_below_log_level_in_context(logger, log_manager):
+    # Log level in context is higher than log level of logger
+    assert logger.level == logging.INFO
+    with utilities.configuration.set_temporary_config({"logging.level": "WARNING"}):
+        logger.info("testing")
+    assert not log_manager.enqueue.called
+    assert log_manager.client is None
+    assert log_manager.thread is None
+
+
+def test_cloud_handler_emit_warns_and_truncates_long_messages(
+    monkeypatch, logger, log_manager, caplog
+):
+    # Smaller value for testing
+    monkeypatch.setattr(prefect.utilities.logging, "MAX_LOG_LENGTH", 100)
+
+    logger.info("h" * 120)
+    # warning about truncating long messages
+    assert any(
+        r.getMessage().startswith("Received a log message of 120 bytes")
+        for r in caplog.records
+    )
+    assert log_manager.enqueue.call_count == 2
+    assert log_manager.enqueue.call_args_list[0][0][0]["message"].startswith(
+        "Received a log message of 120 bytes"
+    )
+    # Truncated log message
+    assert log_manager.enqueue.call_args_list[1][0][0]["message"] == "h" * 100
+
+
+@pytest.mark.parametrize(
+    "flow_run_id, task_run_id",
+    [("flow-run-id", "task-run-id"), ("flow-run-id", None), (None, None)],
+)
+def test_cloud_handler_emit_json_spec(logger, log_manager, flow_run_id, task_run_id):
+    with prefect.context(flow_run_id=flow_run_id, task_run_id=task_run_id):
+        logger.info("testing %d %s", 1, "hello")
+
+    log = log_manager.enqueue.call_args[0][0]
+
+    # Timestamp set on log
+    timestamp = log.pop("timestamp")
+    assert isinstance(timestamp, str)
+
+    # Remaining fields are deterministic
+    assert log == {
+        "message": "testing 1 hello",
+        "flow_run_id": flow_run_id,
+        "task_run_id": task_run_id,
+        "name": logger.name,
+        "level": "INFO",
+    }
 
 
-def test_remote_handler_ships_json_payloads(caplog, monkeypatch):
-    monkeypatch.setattr("prefect.client.Client", MagicMock)
-    client = MagicMock()
-    try:
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger = utilities.logging.configure_logging(testing=True)
-            assert hasattr(logger.handlers[-1], "client")
-            logger.handlers[-1].client = client
-
-            child_logger = logger.getChild("sub-test")
-            try:
-                ## informative comment
-                f = lambda x: x + "2"
-                f(1)
-            except:
-                child_logger.exception("unexpected error")
-
-            time.sleep(0.75)
-            error_logs = [r for r in caplog.records if r.levelname == "ERROR"]
-            assert len(error_logs) >= 1
-
-            cloud_logs = client.write_run_logs.call_args[0][0]
-            assert len(cloud_logs) == 1
-            info = cloud_logs[0]["info"]
-            assert json.loads(json.dumps(info))
+@pytest.mark.parametrize(
+    "flow_run_id, task_run_id",
+    [("flow-run-id", "task-run-id"), ("flow-run-id", None), (None, None)],
+)
+def test_cloud_handler_emit_json_spec_exception(
+    logger, log_manager, flow_run_id, task_run_id
+):
+    with prefect.context(flow_run_id=flow_run_id, task_run_id=task_run_id):
+        try:
+            1 / 0
+        except Exception:
+            logger.exception("An error occurred:")
+
+    log = log_manager.enqueue.call_args[0][0]
+
+    # Timestamp set on log
+    timestamp = log.pop("timestamp")
+    assert isinstance(timestamp, str)
+
+    message = log.pop("message")
+    assert "An error occurred:" in message
+    assert "1 / 0" in message
+    assert "ZeroDivisionError" in message
+
+    # Remaining fields are deterministic
+    assert log == {
+        "flow_run_id": flow_run_id,
+        "task_run_id": task_run_id,
+        "name": logger.name,
+        "level": "ERROR",
+    }
 
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
 
+def test_log_manager_startup_and_shutdown(logger, log_manager):
+    # On creation, neither thread or client are initialized
+    assert log_manager.client is None
+    assert log_manager.thread is None
 
-def test_cloud_handler_responds_to_config(caplog, monkeypatch):
-    calls = []
+    # After enqueue, thread and client are started
+    logger.info("testing")
+    assert log_manager.client is not None
+    assert log_manager.thread is not None
 
-    class Client:
-        def write_run_logs(self, *args, **kwargs):
-            calls.append(1)
+    client = log_manager.client
 
-    monkeypatch.setattr("prefect.client.Client", Client)
-    try:
-        logger = utilities.logging.configure_logging(testing=True)
-        cloud_handler = logger.handlers[-1]
-        assert isinstance(cloud_handler, utilities.logging.CloudHandler)
+    # Calling `_on_shutdown` (which calls stop) will flush the queue and
+    # cleanup resources
+    log_manager._on_shutdown()
+    assert log_manager.queue.empty()
+    assert client.write_run_logs.called
+    assert log_manager.client is None
+    assert log_manager.thread is None
 
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": False}
-        ):
-            logger.critical("testing")
+    # Calling `stop` is idempotent
+    log_manager.stop()
 
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger.critical("testing")
 
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": False}
-        ):
-            logger.critical("testing")
+def test_log_manager_batches_logs(logger, log_manager, monkeypatch):
+    monkeypatch.setattr(prefect.utilities.logging, "MAX_BATCH_LOG_LENGTH", 100)
+    # Fill up log queue with multiple logs exceeding the total batch length
+    for i in range(10):
+        logger.info(str(i) * 50)
+    time.sleep(0.5)
 
-        time.sleep(0.75)
-        assert len(calls) == 1
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
+    assert log_manager.queue.empty()
+    messages = [
+        l["message"]
+        for call in log_manager.client.write_run_logs.call_args_list
+        for l in call[0][0]
+    ]
+    assert messages == [f"{i}" * 50 for i in range(10)]
+    assert log_manager.client.write_run_logs.call_count == 5
 
 
-def test_cloud_handler_removes_bad_logs_from_queue_and_logs_error(caplog, monkeypatch):
-    calls = []
+def test_log_manager_warns_and_retries_on_client_error(
+    logger, log_manager, monkeypatch
+):
+    first_call = True
 
-    class Client:
-        def write_run_logs(self, *args, **kwargs):
-            calls.append(dict(args=args, kwargs=kwargs))
+    def write_run_logs(*args, **kwargs):
+        nonlocal first_call
+        if first_call:
+            first_call = False
+            raise ValueError("Oh no!")
 
-    monkeypatch.setattr("prefect.client.Client", Client)
-    try:
-        logger = utilities.logging.configure_logging(testing=True)
-        cloud_handler = logger.handlers[-1]
-        assert isinstance(cloud_handler, utilities.logging.CloudHandler)
+    log_manager.ensure_started()
+    log_manager.client.write_run_logs = MagicMock(wraps=write_run_logs)
 
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger.critical("one")
-            logger.critical(b"two")
-            logger.critical("three")
+    with pytest.warns(UserWarning, match="Failed to write logs with error:"):
+        logger.info("testing")
 
         time.sleep(0.75)
-        assert len(calls) == 1
-        msgs = [c["message"] for c in calls[0]["args"][0]]
-
-        assert msgs[0] == "one"
-        assert "is not JSON serializable" in msgs[1]
-        assert msgs[2] == "three"
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
-
-
-def test_cloud_handler_client_error(caplog, monkeypatch):
-    class Client:
-        def write_run_logs(self, *args, **kwargs):
-            raise utilities.exceptions.ClientError("Error")
-
-    monkeypatch.setattr("prefect.client.Client", Client)
-    try:
-        logger = utilities.logging.configure_logging(testing=True)
-        cloud_handler = logger.handlers[-1]
-        assert isinstance(cloud_handler, utilities.logging.CloudHandler)
-
-        with utilities.configuration.set_temporary_config(
-            {"logging.log_to_cloud": True}
-        ):
-            logger.critical("one")
-    finally:
-        # reset root_logger
-        logger = utilities.logging.configure_logging(testing=True)
-        logger.handlers = []
-
-
-def test_make_error_log(caplog):
 
-    with context({"flow_run_id": "f_id", "task_run_id": "t_id"}):
-        log = utilities.logging.CloudHandler()._make_error_log("test_message")
-        assert log["flow_run_id"] == "f_id"
-        assert log["timestamp"]
-        assert log["name"] == "CloudHandler"
-        assert log["message"] == "test_message"
-        assert log["level"] == "CRITICAL"
-        assert log["info"] == {}
+    assert log_manager.client.write_run_logs.call_count == 2
+    for call in log_manager.client.write_run_logs.call_args_list:
+        logs = call[0][0]
+        assert isinstance(logs, list)
+        assert len(logs) == 1
+        assert logs[0]["message"] == "testing"
 
 
 def test_get_logger_returns_root_logger():