<@ULVA73B9P>: how can I rewrite the task decorator...
# ask-marvin
i
@Marvin: how can I rewrite the task decorator so that an option log_prints=True is enforced?
m
thought for 101.0 seconds
Great question! Assuming Prefect 3.x, you can wrap the
@task
decorator to force
log_prints=True
. Here are two options depending on how strict you want to be. 1) Strict enforcement (always True, even if someone tries to override)
Copy code
from prefect import task as _task

def logged_task(*args, **kwargs):
    # Enforce log_prints=True no matter what
    kwargs["log_prints"] = True
    return _task(*args, **kwargs)

# Usage
@logged_task
def my_task():
    print("This will be logged")
Works with and without parentheses:
Copy code
@logged_task(retries=2)
def another():
    print("Also logged")
2) Default to True but allow override
Copy code
from prefect import task as _task

def logged_task(*args, **kwargs):
    # Default to True only if not provided
    kwargs.setdefault("log_prints", True)
    return _task(*args, **kwargs)

# Usage
@logged_task
def my_task():
    print("This will be logged")

@logged_task(log_prints=False)
def quiet_task():
    print("This will NOT be logged")
Bonus: If you’d rather enforce this globally without wrapping the decorator, you can set the global config so tasks/flows inherit it by default: - Environment variable:
PREFECT_LOGGING_LOG_PRINTS=true
- Or via CLI:
prefect config set PREFECT_LOGGING_LOG_PRINTS=true
i
@Marvin If I also want to print "Starting {task_name}" at the beginning of the task and "Finished {task_name}" when finished, how can it be done?
m
thought for 56.7 seconds
You can wrap the Prefect task decorator and the underlying function to emit those messages. Here are two options. Option A: Use prints (pair with log_prints=True)
Copy code
import inspect
from functools import wraps
from prefect import task as _task
from prefect.runtime import task_run

def _safe_task_name(fn):
    try:
        return task_run.task_name or fn.__name__
    except Exception:
        return fn.__name__

def logged_task(_fn=None, /, **task_kwargs):
    # Default to logging prints unless explicitly overridden
    task_kwargs.setdefault("log_prints", True)

    def decorator(fn):
        if inspect.iscoroutinefunction(fn):
            @wraps(fn)
            async def wrapped(*args, **kwargs):
                name = _safe_task_name(fn)
                print(f"Starting {name}")
                try:
                    result = await fn(*args, **kwargs)
                except Exception as e:
                    print(f"Failed {name}: {e}")
                    raise
                else:
                    print(f"Finished {name}")
                    return result
        else:
            @wraps(fn)
            def wrapped(*args, **kwargs):
                name = _safe_task_name(fn)
                print(f"Starting {name}")
                try:
                    result = fn(*args, **kwargs)
                except Exception as e:
                    print(f"Failed {name}: {e}")
                    raise
                else:
                    print(f"Finished {name}")
                    return result

        return _task(**task_kwargs)(wrapped)

    return decorator if _fn is None else decorator(_fn)

# Usage
@logged_task
def my_task(x: int) -> int:
    print("Doing work...")
    return x + 1

@logged_task(retries=2, retry_delay_seconds=1)  # still supports task kwargs
def other_task():
    return "ok"
Option B: Use Prefect’s logger (does not depend on log_prints) ``` import inspect from functools import wraps from prefect import task as _task, get_run_logger from prefect.runtime import task_run def _safe_task_name(fn): try: return task_run.task_name or fn.name except Exception: return fn.name def logged_task_with_logger(_fn=None, /, **task_kwargs): # No need to force log_prints when using the logger def decorator(fn): if inspect.iscoroutinefunction(fn): @wraps(fn) async def wrapped(*args, **kwargs): name = _safe_task_name(fn) try: logger = get_run_logger() logger.info(f"Starting {name}") except Exception: print(f"Starting {name}") try: result = await fn(*args, **kwargs) except Exception as e: try: get_run_logger().exception(f"Failed {name}: {e}") except Exception: print(f"Failed {name}: {e}") raise else: try: get_run_logger().info(f"Finished {name}") except Exception: print(f"Finished {name}") return result else: @wraps(fn) def wrapped(*args, **kwargs): name = _safe_task_name(fn) try: get_run_logger().info(f"Starting {name}") except Exception: print(f"Starting {name}") try: result = fn(*args, **kwargs) except Exception as e: try: get_run_logger().exception(f"Failed {name}: {e}") except Exception: print(f"Failed {name}: {e}") raise else: try: get_run_logger().info(f"Finished {name}") except Exception:
print(f"Finished {name}") return result return _task(**task_kwargs)(wrapped) return decorator if _fn is None else decorator(_fn) # Usage @logged_task_with_logger def my_logged_task(): pass ``` Notes: -
prefect.runtime.task_run.task_name
gives you the Prefect task’s name at runtime; we fall back to the function name when not in a task run context. - If you want to strictly enforce
log_prints=True
, replace
task_kwargs.setdefault("log_prints", True)
with
task_kwargs["log_prints"] = True
.
i
@Marvin I also need the same for a flow
m
thought for 245.0 seconds
Absolutely — here are two solid ways to do the same for flows. Option A (recommended): Use flow state hooks + Prefect logger (works regardless of log_prints)
Copy code
from prefect import flow as _flow, get_run_logger

def flow_with_logging(_fn=None, /, **flow_kwargs):
    # Enable print capture for any print() inside your flow body
    flow_kwargs.setdefault("log_prints", True)

    def start_hook(flow, flow_run, state):
        try:
            get_run_logger().info(f"Starting {flow.name}")
        except Exception:
            print(f"Starting {flow.name}")

    def finish_hook(flow, flow_run, state):
        try:
            get_run_logger().info(f"Finished {flow.name}")
        except Exception:
            print(f"Finished {flow.name}")

    def failure_hook(flow, flow_run, state):
        try:
            get_run_logger().error(f"Failed {flow.name}: {state.message}")
        except Exception:
            print(f"Failed {flow.name}: {state.message}")

    # Merge with any user-provided hooks
    flow_kwargs["on_running"] = [start_hook] + list(flow_kwargs.get("on_running") or [])
    flow_kwargs["on_completion"] = list(flow_kwargs.get("on_completion") or []) + [finish_hook]
    flow_kwargs["on_failure"] = list(flow_kwargs.get("on_failure") or []) + [failure_hook]

    def decorator(fn):
        return _flow(**flow_kwargs)(fn)

    return decorator if _fn is None else decorator(_fn)

# Usage
@flow_with_logging
def my_flow(x: int, y: int):
    print("Flow body running")  # captured because log_prints=True
    return x + y
Option B: Wrap the flow body with prints (pair with log_prints=True)
Copy code
import inspect
from functools import wraps
from prefect import flow as _flow

def flow_with_prints(_fn=None, /, **flow_kwargs):
    # Ensure printed messages are captured by Prefect logs
    flow_kwargs.setdefault("log_prints", True)

    def decorator(fn):
        flow_name = flow_kwargs.get("name") or fn.__name__

        if inspect.iscoroutinefunction(fn):
            @wraps(fn)
            async def wrapped(*args, **kwargs):
                print(f"Starting {flow_name}")
                try:
                    result = await fn(*args, **kwargs)
                except Exception as e:
                    print(f"Failed {flow_name}: {e}")
                    raise
                else:
                    print(f"Finished {flow_name}")
                    return result
        else:
            @wraps(fn)
            def wrapped(*args, **kwargs):
                print(f"Starting {flow_name}")
                try:
                    result = fn(*args, **kwargs)
                except Exception as e:
                    print(f"Failed {flow_name}: {e}")
                    raise
                else:
                    print(f"Finished {flow_name}")
                    return result

        return _flow(**flow_kwargs)(wrapped)

    return decorator if _fn is None else decorator(_fn)

# Usage
@flow_with_prints
def another_flow():
    print("Work…")
Notes: - Flow-level
log_prints=True
will be inherited by tasks that don’t override
log_prints
, so both your flow and its tasks’ prints can be captured. - Hooks (Option A) are more robust: they log at state transitions (start, finish, fail) and don’t rely on
print
. You can keep both (hooks + log_prints) if you want captured prints from your flow body too.