def zeromq_aware_logging(f):
    """Ensure multiprocessing logging uses ZeroMQ queues.

    ZeroMQ and local stdout/stderr do not behave nicely when intertwined. This
    ensures the local logging uses existing ZeroMQ logging queues.
    """
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        config = None
        for arg in args:
            if config_utils.is_std_config_arg(arg):
                config = arg
                break
            elif config_utils.is_nested_config_arg(arg):
                config = arg["config"]
            elif isinstance(arg, (list, tuple)) and config_utils.is_nested_config_arg(arg[0]):
                config = arg[0]["config"]
                break
        assert config, "Could not find config dictionary in function arguments."
        if config.get("parallel", {}).get("log_queue") and not config.get("parallel", {}).get("wrapper"):
            handler = setup_local_logging(config, config["parallel"])
        else:
            handler = None
        try:
            out = f(*args, **kwargs)
        finally:
            if handler and hasattr(handler, "close"):
                handler.close()
        return out
    return wrapper
Ensure multiprocessing logging uses ZeroMQ queues.

    ZeroMQ and local stdout/stderr do not behave nicely when intertwined. This
    ensures the local logging uses existing ZeroMQ logging queues.