Great question. Short version: with Dask each mapped task runs in its own worker process, so a transaction you open in the flow won’t automatically apply inside those workers. To get per-task rollback, open the transaction inside the task itself and register the compensating action there. A concrete pattern you can use (Prefect 3.x):

from pathlib import Path
import shutil
from dask_jobqueue import PBSCluster
from prefect import flow, task
from prefect_dask.task_runners import DaskTaskRunner
from prefect.transactions import transaction

@task
def download_from_api(item: str) -> Path:
    # ... produce a local zip Path for this item ...
    ...

@task
def save_zip(zip_path: Path, dest_dir: Path) -> Path:
    # Open a transaction inside the task so it applies on the Dask worker
    with transaction(key=f"save-zip:{zip_path.name}") as txn:
        dest = dest_dir / zip_path.name

        # Register compensation for this unit of work
        # Use a closure to capture the concrete dest path
        txn.add_rollback_hook(lambda t, p=dest: p.exists() and p.unlink())

        # Do the side-effect
        shutil.copy2(zip_path, dest)

        # Optionally also register a commit action (e.g., tag/mark)
        # txn.add_commit_hook(lambda t, p=dest: ...)

        return dest

@flow(
    task_runner=DaskTaskRunner(
        cluster_class=PBSCluster,
        cluster_kwargs=...,         # your PBS config
        adapt_kwargs={"maximum": 20}
    ),
    persist_result=True,
)
def my_flow(items: list[str], dest_dir: Path) -> list[Path]:
    download_futures = download_from_api.map(items)
    saved_zipfile_futures = save_zip.map(zip_path=download_futures, dest_dir=dest_dir)
    return saved_zipfile_futures

Notes and options - Why inside the task? Dask executes mapped tasks in separate processes; opening

with transaction():

in the flow won’t be visible to those workers. Opening it in the task ensures the rollback hooks run if that unit fails or if its transaction is rolled back. - Task-level hooks: You can also attach hooks at task definition time using

@task(on_rollback=[...], on_commit=[...])

. These run in the context of the active transaction for that task run. They’re great for fixed behavior, but registering hooks inside the task body (as above) is more flexible because you can capture per-item details like a specific file path. - State hooks vs transactions:

on_failure

state hooks (docs) run when a task/flow enters a failed state. Use these for “cleanup only if this task fails.” Transactions (docs) add “commit or roll back” semantics and can run compensation even when a later error occurs elsewhere in the transaction scope. If you want “all-or-nothing” rollback across the entire map - Per-item transactions (above) only clean up the failing item. If you want to roll back all saved files when any item fails, you can: - Create a “parent” transaction in the flow, pass a transaction key into each task, and have the task open/join a child transaction using that key; then if any failure occurs and the parent rolls back, all child hooks run. - This pattern requires a shared ResultStore and propagating a transaction key to workers. It’s a bit more advanced; see Transactions docs for details and test in your environment. Quick sanity check - Dask runner import: in Prefect 3.x, use

from prefect_dask.task_runners import DaskTaskRunner

. - PBSCluster: ensure you’ve configured

dask-jobqueue

and PBS resources correctly. Could you confirm which behavior you need? - A) Roll back only the side-effects for the item that failed - B) Roll back all previously saved items if any mapped task fails I can tailor the code to your exact rollback semantics once I know which you’re aiming for.

Perfect — per-item rollback with task-level hooks is a good fit. Key points for Prefect 3.x: - Task-level transaction hooks receive exactly one argument: a Transaction object. - Hooks run only if the task executes inside a transaction; make sure to open one in the task body. - To get item-specific cleanup, stash any values the hook needs into the transaction via txn.set(...), then read them in the hook with txn.get(...). Example using task-level hooks for per-item cleanup on failure (works with Dask):

from pathlib import Path
import shutil

from dask_jobqueue import PBSCluster
from prefect import flow, task, get_run_logger
from prefect.transactions import transaction, Transaction
from prefect_dask.task_runners import DaskTaskRunner


# Hook signatures: each receives a single Transaction
def cleanup_saved_file(txn: Transaction) -> None:
    # Read values captured during the task run
    dest_str = txn.get("dest_path", None)
    if not dest_str:
        return
    dest = Path(dest_str)
    try:
        if dest.exists():
            dest.unlink()
    except Exception as exc:
        # Avoid raising from hooks; just log if needed
        logger = get_run_logger()
        logger.warning(f"Rollback hook could not remove {dest}: {exc}")


def log_commit(txn: Transaction) -> None:
    logger = get_run_logger()
    <http://logger.info|logger.info>(f"Committed file: {txn.get('dest_path', '<unknown>')}")


@task(
    on_rollback=[cleanup_saved_file],  # runs if the transaction rolls back
    on_commit=[log_commit],            # runs if the transaction commits
)
def save_zip(zip_path: Path, dest_dir: Path) -> Path:
    # Open a transaction INSIDE the task so it applies on the Dask worker
    with transaction(key=f"save-zip:{zip_path.name}") as txn:
        dest = dest_dir / zip_path.name

        # Make data available to hooks
        txn.set("dest_path", str(dest))

        # Side-effect
        shutil.copy2(zip_path, dest)

        return dest


@task
def download_from_api(item: str) -> Path:
    # ... produce a local zip Path for this item ...
    ...


@flow(
    task_runner=DaskTaskRunner(
        cluster_class=PBSCluster,
        cluster_kwargs=...,            # your PBS config
        adapt_kwargs={"maximum": 20},
    ),
    persist_result=True,
)
def my_flow(items: list[str], dest_dir: Path) -> list[Path]:
    download_futures = download_from_api.map(items)
    saved_zipfile_futures = save_zip.map(zip_path=download_futures, dest_dir=dest_dir)
    return saved_zipfile_futures

Alternative hook registration style - You can also attach hooks after defining a task:

@save_zip.on_rollback
def cleanup(txn: Transaction):
    ...

@save_zip.on_commit
def committed(txn: Transaction):
    ...

Notes and gotchas - Transactions are not opened automatically just because hooks are present; you must call

with transaction() as txn:

inside the task for the hooks to run. - To make hooks item-aware, put per-item values into the transaction with

txn.set("key", value)

in the task, then read them in the hook via

txn.get("key")

. - Hooks run in the Dask worker where the task executes, so

get_run_logger()

works there. - If you ever see inconsistencies accessing data from on_commit hooks in nested/parent transactions, see this GitHub report for context: issue #17864. Docs - Transactions: Advanced transactions - State hooks overview: State-change hooks

Short answer: yes—that pattern is correct in Prefect 3.x, and it’s a good way to guard a critical section across Dask workers when they share an NFS volume. A few notes and a corrected example with 3.x imports: - Use a shared lock directory that every Dask worker can access. On NFS, make sure file locking is actually enabled and reliable (e.g., no nolock mount option; NFSv4 generally handles this better). - Prefer getting your Prefect home path via settings in 3.x rather than

PREFECT_HOME.value()

. - Make sure the transaction key is stable and uniquely identifies the resource you’re protecting (e.g., the target file path or dataset/date partition). -

persist_result=True

on tasks is independent of the transaction; the transaction’s commit record is written via the

ResultStore

you pass to

transaction(...)

. Example (Prefect 3.x):

from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager
from prefect.settings import get_current_settings

lock_dir = get_current_settings().home / "locks"

with transaction(
    key=transaction_key,  # e.g., f"write:{target_path}"
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)
    ),
) as txn:
    if txn.is_committed():
        # Someone already committed this work
        return

    # Do the work here (ideally write to a tmp path)
    # ...

    # Let the context commit at exit; keep the critical section minimal

Caveats and tips: - NFS locking reliability and performance vary by implementation and mount options. If you see contention or flaky behavior, verify NFS lock daemon/opts and consider a smaller critical section or a single-writer pattern. Prefect’s built-in lock managers are memory (single-process) and filesystem; there isn’t a built-in Redis/ZK lock manager in 3.x. - SERIALIZABLE gives the strongest guarantees but can increase contention. Keep the code inside the

with transaction(...)

block as short as possible—just the check and the finalization step. - If your Dask workers run on nodes that do not share the same filesystem, file-based locks won’t coordinate across hosts. Docs and references: - Results and result storage (Prefect 3) - FileSystemLockManager API - Blog: Transactional pipelines with Prefect 3.0 If you can share how your Dask workers are deployed (same host vs multiple hosts, NFS mount options), I can sanity-check the locking setup further.

Yes—you can open the transaction context anywhere inside your task. The

transaction(...)

just scopes the critical section you want to protect; it does not need to be at the top of the task. How to choose placement: - If your goal is to skip work when it’s already done, put a quick “preflight” check near the top, then do heavy work outside the lock, and re-enter a transaction only for the final write. This minimizes lock time while still avoiding duplicates. - If your goal is to protect only the final write/rename, open the transaction right around that write. This is usually best for Dask to reduce contention. - If you wrap the whole task in a transaction, you’ll hold the lock longer and increase contention—only do this if truly necessary. Two recommended patterns 1) Preflight check + lock only for commit (minimal lock time)

from pathlib import Path
import os
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager
from prefect.settings import get_current_settings

lock_dir = get_current_settings().home / "locks"
txn_key = f"save_zip:{dataset}:{partition}"  # make this stable/unique
final_path = Path("/nfs/path/to/output.zip")

# Preflight read (no lock held long-term)
with transaction(
    key=txn_key,
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
) as txn:
    if txn.is_committed():
        return

# Do heavy work outside the transaction
tmp_path = final_path.with_suffix(".zip.tmp")
# ... compute and write to tmp_path ...

# Commit phase (short critical section + double-check)
with transaction(
    key=txn_key,
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
) as txn:
    if txn.is_committed():
        # Another worker already committed while we were computing
        # Optionally clean up tmp
        if tmp_path.exists():
            tmp_path.unlink(missing_ok=True)
        return

    os.replace(tmp_path, final_path)  # atomic on POSIX
    # exiting the context records the commit

2) Protect only the final write

with transaction(
    key=txn_key,
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
) as txn:
    if txn.is_committed():
        return
    os.replace(tmp_path, final_path)

Notes: - Keep the code inside the

with transaction(...)

block as short as possible—ideally just the finalization (e.g., atomic rename). This reduces contention between Dask workers. - Your

persist_result=True

task setting is independent of transactions. The transaction’s commit record lives in the

ResultStore

you pass to `transaction(...)`; make sure that store and the lock directory are shared across workers (e.g., same NFS mount). - On NFS, ensure locking is reliable (NFSv4 recommended; avoid nolock). If workers don’t share the same filesystem, file locks won’t coordinate. References: - Results and result storage - FileSystemLockManager API - Transactional pipelines with Prefect 3.0 If you share how your workers are mounted to NFS (v3/v4, mount options), I can help validate the locking approach.

Great catch — by default a transaction will write a commit record on clean exit. So if you run a “preflight” inside a transaction with the same key and then exit without error, that first context would commit and the second context would always see it as committed. Two safe patterns: A) One transaction only (preferred) - Do the heavy work first, outside any transaction - Open a single transaction only for the finalization step, check once, then finalize and let the context commit

from pathlib import Path
import os
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager
from prefect.settings import get_current_settings

lock_dir = get_current_settings().home / "locks"
txn_key = f"save_zip:{dataset}:{partition}"
final_path = Path("/nfs/path/output.zip")
tmp_path = final_path.with_suffix(".zip.tmp")

# Heavy work outside transaction
# ... compute/write to tmp_path ...

with transaction(
    key=txn_key,
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
) as txn:
    if txn.is_committed():
        # Another worker finalized already
        if tmp_path.exists():
            tmp_path.unlink(missing_ok=True)
        return

    os.replace(tmp_path, final_path)  # atomic on POSIX
    # clean exit writes the commit record

B) If you really want a preflight, make it read‑only - Use

write_on_commit=False

for the preflight so it won’t create a commit record - Re-check inside a second transaction before finalizing

# Preflight: read-only, do NOT write a commit record
with transaction(
    key=txn_key,
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
    write_on_commit=False,
) as txn:
    if txn.is_committed():
        return

# Heavy work outside transaction
# ... compute/write to tmp_path ...

# Finalization with commit record
with transaction(
    key=txn_key,
    isolation_level=IsolationLevel.SERIALIZABLE,
    store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
) as txn:
    if txn.is_committed():
        if tmp_path.exists():
            tmp_path.unlink(missing_ok=True)
        return
    os.replace(tmp_path, final_path)

Notes: - Keep the locked section as short as possible (ideally just the final atomic rename). - Using a preflight transaction still acquires the lock; if you want less contention, skip the preflight and use pattern A. - Ensure your Dask workers all see the same lock directory and that NFS locking is reliable. Refs: - prefect.transactions API - FileSystemLockManager API - Results and result storage

That Dask message usually means “the worker could not import or cloudpickle your module,” not that there’s anything inherently wrong with your transaction code. Common causes and fixes: 1) The module isn’t importable on workers - Symptom: distributed workers try to import your functions by module name, but that file isn’t on their sys.path or filesystem. - Fix options: - Package your project and install it in the worker environment. - Or, put your code on shared storage and ensure PYTHONPATH points to it on all workers. - Or, use Prefect deployments with source code distribution: - In Prefect 3.x, use

flow.from_source("https://...", entrypoint="pkg_or_path:flow_fn")

and

deploy(...)

so all workers pull the same code. - Docs: Store flow code 2) Running as a script (main) without guarding imports - If your flow is launched as a script, Dask may struggle to re-import functions defined in main. - Fix: move task/flow definitions into an importable module and run with: -

python -m your_package.flow_script

or - guard your entrypoint with

if __name__ == "__main__":

- Keep all task and hook function definitions at top-level of the module (no nesting/closures). 3) Non-picklable objects captured at import time - Don’t create loggers/contexts/globals that aren’t picklable at import. In Prefect 3.x: - Call

get_run_logger()

and

get_run_context()

inside the task body, not at module import time. - Ensure your

@save_zip.on_commit

/

@save_zip.on_rollback

callbacks are top-level functions that don’t close over non-picklable state. 4) Dependency/version mismatches on workers - Make sure workers have the same Python, Prefect, prefect-dask, and cloudpickle versions as the scheduler/flow runner. A minimal pattern that’s friendly to Dask serialization

from pathlib import Path
import os

from prefect import flow, task, get_run_logger
from prefect_dask import DaskTaskRunner

from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager
from prefect.settings import get_current_settings

@task(tags=["some-tag"], log_prints=True, timeout_seconds=4500, persist_result=True)
def save_zip(zip_path: Path, txn_lock_dir: str):
    logger = get_run_logger()

    # Heavy work outside any transaction if possible...
    # ...

    # Finalize with a single transaction
    lock_dir = Path(txn_lock_dir)
    txn_key = f"esa-id:{zip_path.stem}"
    final_path = Path("/nfs/path/to/output.zip")
    tmp_path = final_path.with_suffix(".zip.tmp")

    with transaction(
        key=txn_key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=lock_dir)),
    ) as txn:
        if txn.is_committed():
            logger.warning("Another worker already committed for %s", txn_key)
            if tmp_path.exists():
                tmp_path.unlink(missing_ok=True)
            return
        os.replace(tmp_path, final_path)

@flow(task_runner=DaskTaskRunner())
def my_flow(zip_paths: list[Path], txn_lock_dir: str):
    for p in zip_paths:
        save_zip.submit(p, txn_lock_dir)

If you need preflight without committing, use write_on_commit=False only for the preflight, then open a second transaction for the finalization (so the preflight doesn’t write a commit record). Questions to help pinpoint the issue - How are you running the flow (local script vs deployment)? If a script, what’s your working directory and how do you start it? - Are workers on separate machines? Do they all have the repo installed or mounted on the same NFS path and in PYTHONPATH? - Can you share a slightly longer traceback around the “Failed to deserialize <flow_script.py> module” line? References - Prefect + Dask integration

- Store and distribute flow code - Transactions API

You’re right to call that out—here’s how to register and use the commit/rollback hooks with your task, plus a correction to your transaction pattern. Key points: - The decorator style you used is supported in Prefect 3.x:

@task_fn.on_commit

and

@task_fn.on_rollback

register hook functions that receive a

Transaction

. - Keep hook functions at module top-level (not nested) so Dask can import/pickle them cleanly. - Your current single transaction uses

write_on_commit=False

, which means it will never record a commit. That will release the lock without recording completion, so another worker can redo the work. If you want a “preflight” check, use a read-only transaction for the check and then a second transaction for the final commit; or more simply, use one transaction only for the finalization with a commit record. Example A: single transaction (minimal contention, hooks fire on commit)

from pathlib import Path
import os

from prefect import flow, task, get_run_logger
from prefect.transactions import transaction, IsolationLevel, Transaction
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager
from prefect_dask import DaskTaskRunner

LOCK_DIR = Path("/nfs/locks")

@task(tags=["some-tag"], log_prints=True, timeout_seconds=4500)
def save_zip(zip_path: Path):
    logger = get_run_logger()

    txn_key = f"esa-id:{zip_path.stem}"
    final_path = Path("/nfs/out") / f"{zip_path.stem}.zip"
    tmp_path = final_path.with_suffix(".zip.tmp")

    # Heavy work outside the transaction; produce tmp_path
    # ... compute & write to tmp_path ...

    with transaction(
        key=txn_key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=LOCK_DIR)),
    ) as txn:
        if txn.is_committed():
            logger.warning("Already committed for %s", txn_key)
            if tmp_path.exists():
                tmp_path.unlink(missing_ok=True)
            return

        os.replace(tmp_path, final_path)  # atomic rename on POSIX

@save_zip.on_commit
def committed(txn: Transaction):
    # e.g. emit a log/event or clean up staging artifacts
    print(f"Committed transaction: {txn.key}")

@save_zip.on_rollback
def cleanup(txn: Transaction):
    # e.g. clean up temp files you can derive from txn.key
    print(f"Rolled back transaction: {txn.key}")

@flow(task_runner=DaskTaskRunner())
def my_flow(paths: list[Path]):
    for p in paths:
        save_zip.submit(p)

Example B: two-phase (read-only preflight + final commit) ``` @task(tags=["some-tag"], log_prints=True, timeout_seconds=4500) def save_zip(zip_path: Path): logger = get_run_logger() txn_key = f"esa-id:{zip_path.stem}" final_path = Path("/nfs/out") / f"{zip_path.stem}.zip" tmp_path = final_path.with_suffix(".zip.tmp") # Preflight (no commit record) with transaction( key=txn_key, isolation_level=IsolationLevel.SERIALIZABLE, write_on_commit=False, # read-only check store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=LOCK_DIR)), ) as txn: if txn.is_committed(): logger.warning("Already committed for %s", txn_key) return # Heavy work outside any lock # ... compute & write to tmp_path ... # Finalize with commit record (hooks will run here) with transaction( key=txn_key, isolation_level=IsolationLevel.SERIALIZABLE, store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=LOCK_DIR)), ) as txn: if txn.is_committed(): logger.warning("Already committed for %s", txn_key) if tmp_path.exists(): tmp_path.unlink(missing_ok=True) return os.replace(tmp_path, final_path) @save_zip.on_commit def committed(txn: Transaction): print(f"Committed transaction: {txn.key}") @save_zip.on_rollback def cleanup(txn: Transaction):

print(f"Rolled back transaction: {txn.key}") ``` About the Dask “Failed to deserialize <flow_script.py> module” error: - Ensure the module is importable on workers: - Install your project as a package on each worker; or - Use a shared code location on the workers’ PYTHONPATH; or - Use a deployment with source code distribution so workers pull the code: - In Prefect 3.x, use

flow.from_source(...).deploy(...)

and run via a work pool. - Keep all defs at top-level (flow, tasks, hooks). No closures around non-picklable state. - Use

get_run_logger()

and

get_run_context()

inside functions, not at import time. - Match Python/Prefect/prefect-dask versions across scheduler and workers. - If you run as a script, guard with

if __name__ == "__main__":

or run as a module. Docs: - Prefect + Dask - Transactions API - Store/distribute flow code If you can share how workers see your code (installed package vs shared mount vs deployment), I can tailor the import/distribution fix.

✅ thought for 41.3 seconds

Got it — thanks for sharing the workaround. Why that likely “fixed” deserialization - Dask needs to import/pickle your hook function on workers. Decorating it with

@task

can change how it’s serialized and, importantly, ensures it’s a top-level, importable symbol with a stable module path. - The root issue is usually importability/pickling of the hook function (e.g., nested/closure, or module not on workers’ PYTHONPATH). Making it a task incidentally made it picklable for your environment. A couple of important implications and safer patterns 1) Prefer plain, top-level functions for hooks if you need them to run synchronously - Transaction hooks are called synchronously during commit/rollback. If the hook is a plain function, you’re guaranteed it completes before the transaction finishes. - Example:

from prefect.transactions import Transaction

def _cleanup_impl(txn: Transaction):
    # Do cleanup deterministically here
    ...

@save_zip.on_rollback
def cleanup(txn: Transaction):
    _cleanup_impl(txn)

@save_zip.on_commit
def committed(txn: Transaction):
    # e.g., emit a metric/log
    ...

2) If you keep

@task

on the hook, consider explicitly waiting - A

@task

hook will schedule a sub-task and typically return a future. Unless you explicitly wait, the transaction may finish before the sub-task completes. - If you want deterministic cleanup, wrap and wait:

from prefect import task

@task
def cleanup_task(txn: Transaction):
    _cleanup_impl(txn)

@save_zip.on_rollback
def cleanup(txn: Transaction):
    fut = cleanup_task.submit(txn)
    fut.result()  # ensure it completes before rollback finishes

- If you’re okay with cleanup running asynchronously after rollback, you can skip

fut.result()

. General Dask-friendly tips - Keep hooks and tasks at module top-level (no closures), and ensure workers can import your module (package install, shared mount on PYTHONPATH, or deploy with

flow.from_source(...).deploy(...)

). - Avoid non-picklable globals at import time; get loggers/contexts inside functions. If you’d like, share a small snippet of your hook and how workers get your code (package vs shared mount vs from_source). I can confirm the most robust approach for your setup. References: - Prefect + Dask - Transactions API - Store/distribute flow code

You’ve uncovered the key gotcha: decorating the hook with @task turns it into a Prefect Task, but transaction hooks expect a plain callable that runs synchronously during rollback/commit. When you register a Task object as a hook, nothing runs because: - The transaction tries to call a regular function, not schedule a task - There’s no flow-level scheduling happening inside the transaction rollback path - Result: your hook appears “disabled” What to do instead - Keep hooks as plain, top-level functions (Dask-importable), and register them with the task using the decorator or method style. - Fix Dask deserialization by making your module importable on workers (package install, shared PYTHONPATH/NFS, or a deployment with source code) rather than wrapping hooks in @task. Correct pattern

from prefect.transactions import Transaction

@save_zip.on_rollback
def cleanup(txn: Transaction) -> None:
    # rollback routines here
    ...

@save_zip.on_commit
def committed(txn: Transaction) -> None:
    # commit routines here
    ...

If you need to run “heavier” cleanup in the hook - Keep the hook itself synchronous and minimal. If you truly need to fan out work, call a plain helper function or trigger an external side-effect (e.g., delete tmp files). Avoid trying to submit new Prefect tasks from a hook. - If you must use a Prefect task, be explicit that you’re okay with it running outside the transactional timing guarantees. For example, have the hook enqueue a message to a durable queue your system consumes, rather than trying to schedule a Prefect subtask inline. Also double-check these points - Ensure the failure occurs inside the with transaction(...) block; rollback hooks only fire when the transaction rolls back. If the task fails after the context exits, the hook won’t run. - If you set write_on_commit=False, commit hooks will never run by design. That does not affect rollback hooks, but worth noting. How to fix Dask deserialization without @task on the hook - Make your code importable on workers: - Package and pip-install your project in the worker environment, or - Put the repo on a shared mount and add it to PYTHONPATH on all workers, or - Use a Prefect 3 deployment with source distribution so workers pull the code: - flow.from_source(...).deploy(...) - Store/distribute flow code - Keep all defs (flow, tasks, hooks) at module top-level; no closures. If you share how your workers access the code (installed package vs shared mount vs deployment), I can suggest the cleanest import/distribution setup for your cluster.

Great question — there are two reliable ways to make your cloned repo importable on the worker so Dask can deserialize your module and your plain function hooks will run. Recommended approach: install the repo as a package after cloning - This makes imports work without messing with PYTHONPATH and is the most robust across Dask workers. - Use deployment pull steps to: 1) clone the repo 2) set the working directory to the clone 3) pip install your project (editable or regular) and its requirements Example prefect.yaml snippet:

deployments:
  - name: my-deployment
    entrypoint: my_package/flows.py:my_flow
    pull:
      - prefect.deployments.steps.git_clone:
          id: repo
          repository: <https://gitlab.com/my_package.git>
          branch: my_feature_branch
      - prefect.deployments.steps.set_working_directory:
          directory: "${{ repo.directory }}"
      - prefect.deployments.steps.utility.pip_install_requirements:
          directory: "${{ repo.directory }}"
          requirements_file: requirements.txt
      - prefect.deployments.steps.utility.run_shell_script:
          script: "pip install -e ."
    work_pool:
      name: my-pool

- After this, workers can import your code (e.g.,

from my_package.tasks import save_zip

) and Dask will be able to import modules by name. - Your transaction hooks should be plain, top-level functions:

from prefect.transactions import Transaction

@save_zip.on_rollback
def cleanup(txn: Transaction) -> None:
    ...

@save_zip.on_commit
def committed(txn: Transaction) -> None:
    ...

Alternative: add the clone to PYTHONPATH via job variables - If you can’t install the package, set PYTHONPATH to include the clone path emitted by

git_clone

. - Use step outputs with an id and reference

${{ repo.directory }}

. Example:

deployments:
  - name: my-deployment
    entrypoint: my_package/flows.py:my_flow
    pull:
      - prefect.deployments.steps.git_clone:
          id: repo
          repository: <https://gitlab.com/my_package.git>
          branch: my_feature_branch
      - prefect.deployments.steps.set_working_directory:
          directory: "${{ repo.directory }}"
    work_pool:
      name: my-pool
      job_variables:
        working_dir: "${{ repo.directory }}"
        env:
          PYTHONPATH: "${{ repo.directory }}:${{ env.PYTHONPATH }}"

- Note: use “;” instead of “:” on Windows. Important caveats - If your Dask cluster is external (address=...), those remote Dask workers do not inherit env from the Prefect worker. You must also ensure the repo is installed or available on those Dask workers (install the package there or mount the same path and set PYTHONPATH on the Dask worker side). - Keep all flows, tasks, and hooks at module top-level; avoid closures/non-picklable globals at import time. - If you’re using

write_on_commit=False

anywhere, note that commit hooks won’t fire for those transactions by design; rollback hooks still require the failure to occur inside the

with transaction(...)

block. Docs - Customize job variables: Customize job variables - Steps you can use:

git_clone

,

set_working_directory

,

pip_install_requirements

,

run_shell_script

If you can share whether your Dask cluster is local-in-worker vs external, I can suggest the exact setup (install vs PYTHONPATH) for that topology.

Thanks — since your Dask cluster is launched “inside” the Prefect worker and submits PBS jobs to nodes that all share NFS, you have two good options. The goal is to ensure the PBS-launched Dask workers can import your repo by module name so your plain function hooks run. Option A (recommended): install the repo after cloning - Works without fiddling with PYTHONPATH at runtime - Most robust if the Python environment is the same across scheduler and workers prefect.yaml

deployments:
  - name: my-deployment
    entrypoint: my_package/flows.py:my_flow
    pull:
      - prefect.deployments.steps.git_clone:
          id: repo
          repository: <https://gitlab.com/my_package.git>
          branch: my_feature_branch
      - prefect.deployments.steps.set_working_directory:
          directory: "${{ repo.directory }}"
      - prefect.deployments.steps.utility.pip_install_requirements:
          directory: "${{ repo.directory }}"
          requirements_file: requirements.txt
      - prefect.deployments.steps.utility.run_shell_script:
          # editable install so the installed package points to this branch/clone
          script: "pip install -e ."
    work_pool:
      name: my-pool
      job_variables:
        # optional but helpful for your code to know where the repo lives
        working_dir: "${{ repo.directory }}"
        env:
          REPO_DIR: "${{ repo.directory }}"

Then in code, just import your package normally and keep hooks as plain functions:

from prefect.transactions import Transaction

@save_zip.on_rollback
def cleanup(txn: Transaction) -> None:
    ...

@save_zip.on_commit
def committed(txn: Transaction) -> None:
    ...

Notes - Editable install creates a .pth pointing to the clone. If your Python environment is on a shared filesystem (common on HPC) the Dask workers will pick it up. Ensure workers use the same environment (modules/conda) as the Prefect process. Option B: set PYTHONPATH for PBS-launched Dask workers - Use repo directory from the git_clone step and export it in the Dask job script via dask-jobqueue’s env_extra. - Keep editable install if you also want entry points/console scripts; otherwise PYTHONPATH alone is enough for imports. prefect.yaml

deployments:
  - name: my-deployment
    entrypoint: my_package/flows.py:my_flow
    pull:
      - prefect.deployments.steps.git_clone:
          id: repo
          repository: <https://gitlab.com/my_package.git>
          branch: my_feature_branch
      - prefect.deployments.steps.set_working_directory:
          directory: "${{ repo.directory }}"
    work_pool:
      name: my-pool
      job_variables:
        working_dir: "${{ repo.directory }}"
        env:
          # expose the dynamic repo path to your flow code
          REPO_DIR: "${{ repo.directory }}"
          # also make it available to the Prefect process; helpful for LocalCluster cases
          PYTHONPATH: "${{ repo.directory }}:${{ env.PYTHONPATH }}"

Flow code: tell the PBSCluster to export PYTHONPATH on worker jobs

import os
from prefect import flow
from prefect_dask import DaskTaskRunner

REPO_DIR = os.environ["REPO_DIR"]

task_runner = DaskTaskRunner(
    cluster_class="dask_jobqueue.PBSCluster",
    cluster_kwargs={
        # ... your PBS settings: cores, memory, walltime, queue, etc.
        "env_extra": [f"export PYTHONPATH={REPO_DIR}:$PYTHONPATH"],
        # optionally:
        # "job_script_prologue": ["module load python/3.10", "source /path/to/venv/bin/activate"],
    },
)

@flow(task_runner=task_runner)
def my_flow(...):
    ...

Which should you choose? - If your Python environment is truly identical across nodes and on shared storage, Option A (pip install -e .) is simplest. - If you want to avoid touching site-packages or your environment isn’t shared, Option B (PYTHONPATH via env_extra) is safe and explicit. - You can also combine them: install requirements and your package, and still export PYTHONPATH for belt-and-suspenders. Final reminders

- Keep flow, tasks, and hook functions at module top-level; do not decorate hooks with @task. - Rollback hooks only fire if the failure occurs inside the with transaction(...) block. - On NFS, prefer atomic renames (os.replace) and keep the lock-held section minimal. Helpful docs: - Customize job variables (env, working_dir) - Prefect + Dask integration - Store/distribute flow code

Short answer: yes. In Prefect 3.x, both Docker and Kubernetes workers accept environment variables via the

env

field in

job_variables

. You can set

PYTHONPATH

to include multiple paths by joining them with your OS path separator (colon on Linux/macOS, semicolon on Windows). There isn’t a built-in

env_extra

field in the default worker templates—use

env

. Given your code lives under

repo.directory

, you can add that (and any other folders) explicitly. Examples: - prefect.yaml

deployments:
  - name: my-deployment
    entrypoint: repo/directory/flows.py:my_flow
    work_pool:
      name: my-docker-or-k8s-pool
      job_variables:
        env:
          PYTHONPATH: "./repo/directory:./another/path"

- Python API

flow.from_source("<https://github.com/org/repo>", entrypoint="repo/directory/flows.py:my_flow").deploy(
    name="my-deployment",
    work_pool_name="my-docker-or-k8s-pool",
    job_variables={
        "env": {
            "PYTHONPATH": "./repo/directory:./another/path"
        }
    },
)

Notes: - Relative paths are resolved from the flow’s working directory inside the run environment. If you clone your repo at runtime, ensure the working directory is the repo root (or set

entrypoint

relative to wherever you land). - If you’d like to inspect the default worker template to see the

env

key, you can:

prefect work-pool get-default-base-job-template --type docker
prefect work-pool get-default-base-job-template --type kubernetes

Docs: - Customize job variables - Docker workers - Kubernetes workers If you were relying on an

env_extra

pattern from an older example or a custom base job template, you can still use it—just make sure your base job template merges

env_extra

into

env

. By default, Prefect’s templates only expose

env

.

Got it — thanks for the context. - Yes, you can add extra paths via

job_variables.env.PYTHONPATH

. Join multiple paths with

:

on Linux/macOS (or

;

on Windows). - Since you’re cloning at runtime with a

pull

step, you cannot reference the clone path in

job_variables

directly. Instead, set the working directory after the clone, then use relative paths like

.

in

PYTHONPATH

. Example prefect.yaml pattern:

pull:
  - prefect.deployments.steps.git_clone:
      id: clone
      repository: <https://gitlab.com/my_package.git>
      branch: my_feature_branch
  - prefect.deployments.steps.set_working_directory:
      # point to the subdir that contains your flow code
      directory: "{{ clone.directory }}/repo/directory"

deployments:
  - name: my-dask-deployment
    entrypoint: repo/directory/flows.py:my_flow
    work_pool:
      name: my-pool
      job_variables:
        env:
          # include the current working dir plus any other shared libs
          PYTHONPATH: ".:/nfs/shared/other_libs"

A few important notes for your Dask + PBS setup: - The

env

you set in Prefect

job_variables

affects the flow-run process. It does not automatically propagate to Dask workers that PBS launches on other nodes. - Since you have NFS across nodes, make sure the repo clone happens into a shared path if you want workers to import from it. For a Process work pool, you can set a shared

working_dir

so the clone lands on NFS:

work_pool:
  name: my-process-pool
  job_variables:
    working_dir: "/nfs/shared/prefect-workdir"
    env:
      PYTHONPATH: ".:/nfs/shared/other_libs"

- Also configure your Dask cluster to export the same

PYTHONPATH

to workers. With dask-jobqueue’s PBSCluster, use `env_extra`:

from prefect_dask.task_runners import DaskTaskRunner

task_runner = DaskTaskRunner(
    cluster_class="dask_jobqueue.PBSCluster",
    cluster_kwargs={
        # ... your PBS options ...
        "env_extra": [
            "export PYTHONPATH=/nfs/shared/prefect-workdir/repo/directory:/nfs/shared/other_libs:$PYTHONPATH"
        ],
    },
)

Related docs: - Prefect: Customize job variables - Dask Jobqueue: dask-jobqueue docs Quick clarifier so I can tailor this precisely: - Which work pool type are you using (Process, Docker, or Kubernetes)? - Do you have a specific NFS path you want to clone into? If so, we can set that as the

working_dir

(Process) or rely on

set_working_directory

(Docker/K8s) and adjust

PYTHONPATH

accordingly.

You’re hitting that warning because

{{ clone.directory }}

is a pull-step output that only exists at flow run time on the worker. The

deployments

section is rendered at deploy time, so it can’t see outputs from

pull

steps. What works vs. not: - Works: Referencing

{{ clone.directory }}

in subsequent pull steps (e.g.,

set_working_directory

) because those run after

git_clone

at runtime. - Doesn’t work: Using

{{ clone.directory }}

in the

deployments:

section (including

job_variables

) because that’s evaluated during

prefect deploy

. Practical patterns you can use: Option 1: Set working directory to the cloned repo, then use relative paths in env

pull:
  - prefect.deployments.steps.git_clone:
      id: clone
      repository: <https://gitlab.com/myrepo.git>
      branch: my_branch
  - prefect.deployments.steps.set_working_directory:
      directory: "{{ clone.directory }}"

deployments:
  - name: my_deployment
    work_pool:
      name: my_pool
      job_variables:
        env:
          # now "." resolves to the cloned repo directory at runtime
          REPO_PATH: "."
          PYTHONPATH: ".:/nfs/shared/other_libs"

Option 2: Clone into a known NFS path, then reference that static path in env - If you can control the clone destination (either via a

directory

argument on the clone step, or by an additional step to move it), choose a stable NFS path that all nodes can see.

pull:
  - prefect.deployments.steps.git_clone:
      id: clone
      repository: <https://gitlab.com/myrepo.git>
      branch: my_branch
      # if supported, specify a target directory everyone can reach
      directory: "/nfs/shared/repos/myrepo"
  - prefect.deployments.steps.set_working_directory:
      directory: "/nfs/shared/repos/myrepo"

deployments:
  - name: my_deployment
    work_pool:
      name: my_pool
      job_variables:
        env:
          REPO_PATH: "/nfs/shared/repos/myrepo"

If

directory

isn’t supported in your setup, you can add a shell step to move the clone into your target path, then

set_working_directory

there. Dask + PBS note: - The

env

in

job_variables

affects the Prefect flow-run process only. It will not automatically propagate to the Dask workers PBS launches. If your tasks import from

REPO_PATH

, export the same value in your Dask/PBS worker environment (e.g., via

env_extra

in your PBSCluster config) and make sure it points to a shared NFS location. References: - prefect.yaml and templating stages - Pull steps (git_clone, set_working_directory) - Customize job variables If you share your work pool type (Process/Docker/Kubernetes), I can suggest the most robust way to make the clone path and env consistent across your Dask workers.

Short answer: you can’t access

{{ clone.directory }}

directly from your flow code. Pull step outputs only exist within the prefect.yaml templating context and only for other pull steps. Workable patterns to achieve the same outcome: 1) Set the working directory to the clone, then use Path.cwd() in your flow - In prefect.yaml:

pull:
  - prefect.deployments.steps.git_clone:
      id: clone
      repository: <https://gitlab.com/myrepo.git>
      branch: my_branch
  - prefect.deployments.steps.set_working_directory:
      directory: "{{ clone.directory }}"

- In your flow:

from pathlib import Path
from prefect import flow

@flow
def my_flow():
    repo_dir = Path.cwd()
    print(f"Repo cloned at: {repo_dir}")

2) Persist the clone path to a file during pull, then read it in your flow - In prefect.yaml:

pull:
  - prefect.deployments.steps.git_clone:
      id: clone
      repository: <https://gitlab.com/myrepo.git>
      branch: my_branch
  - prefect.deployments.steps.utility.run_shell_script:
      id: write-clone-dir
      script: |
        mkdir -p .prefect_meta
        echo "{{ clone.directory }}" > .prefect_meta/clone_dir.txt

- In your flow:

from pathlib import Path
from prefect import flow

def get_clone_dir():
    p = Path(".prefect_meta/clone_dir.txt")
    return p.read_text().strip() if p.exists() else str(Path.cwd())

@flow
def my_flow():
    clone_dir = get_clone_dir()
    print(f"Repo cloned at: {clone_dir}")

3) Clone to a known shared path and just use that in your code - If you need all workers to see the same path (e.g., NFS), move the clone into a stable directory in a pull step, then either

set_working_directory

there or reference that path directly in your flow. Why this is necessary: - Pull steps run on the worker before your flow starts. Their outputs (like

clone.directory

) are only available for templating within the pull section, not exposed via

prefect.runtime

or environment variables for the flow. References: - prefect.yaml and templating - Pull steps (git_clone, set_working_directory) - Runtime context

That FileNotFoundError means the placeholder wasn’t rendered — your second step received the literal string “{{ clone.directory }}” instead of the actual path. A few things to check and a minimal working pattern: What to check - Make sure the step id matches exactly: - You have

id: clone

and you reference

{{ clone.directory }}

— that’s correct. - Quote the template string: - Use

directory: "{{ clone.directory }}"

(double quotes). Avoid wrapping in angle brackets; if your YAML has

<{{ clone.directory }}>

, remove the angle brackets. - Verify you’re on a recent Prefect 3.x: - There have been fixes in the step templating flow; upgrading to the latest 3.x is recommended. - Keep steps in order: - Steps run sequentially. If you want to be extra explicit, you can add

requires: [clone]

to the second step. A minimal working example

pull:
  - prefect.deployments.steps.git_clone:
      id: clone
      repository: <https://gitlab.com/jrsrp/sys/workflows/overnight_prefect.git>
      branch: "25-s2import-implement-rollback-hook-to-avoid-incomplete-qvf-imports"

  - prefect.deployments.steps.set_working_directory:
      # This must be quoted; do not wrap with angle brackets
      directory: "{{ clone.directory }}"
      # Optional but explicit:
      requires: [clone]

If you want to confirm the value being rendered, add a quick debug step between them:

- prefect.deployments.steps.utility.run_shell_script:
      id: debug
      script: |
        echo "clone dir is: {{ clone.directory }}"
      requires: [clone]

If the debug echo shows the placeholder instead of a concrete path, templating isn’t happening. In that case: - Double-check your YAML file for stray angle brackets or escaping. - Ensure there’s no duplicate top-level

pull:

vs per-deployment

pull:

overriding each other. - Share your Prefect version (

prefect version

) — I can advise if an upgrade is needed. Docs for reference: - Prefect YAML and templating: https://docs-3.prefect.io/v3/how-to-guides/deployments/prefect-yaml - Pull steps (git_clone, set_working_directory): https://docs-3.prefect.io/v3/api-ref/python/prefect-deployments-steps-pull

Great progress! The hooks should work with Dask — they execute on the Dask worker in the same process as the task — but there are a few common gotchas that can make them look like they’re “not firing”. Checklist to make them work and visible: - Ensure a real commit/rollback occurs - In your code, if

txn.is_committed()

returns True, you return early. In that path no new commit happens and no hooks run. Test with a fresh

txn_key

or force a rollback (raise an exception inside the

with transaction(...)

block) to verify hooks. - Use Prefect’s logger inside hooks rather than print - Your prints run on the Dask worker and may not show up in Prefect logs. Use

get_run_logger()

so they land in the task logs. - Keep hook functions at module-level (importable by workers) - You’ve already addressed importability via PYTHONPATH; just make sure hooks aren’t nested/closures. Example update using Prefect logger:

from prefect import task, flow, get_run_logger, Transaction
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from pathlib import Path
import os

LOCK_DIR = Path("/nfs/locks")

@task(tags=["some-tag"], log_prints=True, timeout_seconds=4500)
def save_zip(zip_path: Path):
    logger = get_run_logger()

    txn_key = f"esa-id:{zip_path.stem}"
    final_path = Path("/nfs/out") / f"{zip_path.stem}.zip"
    tmp_path = final_path.with_suffix(".zip.tmp")

    # ... produce tmp_path ...

    with transaction(
        key=txn_key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=LOCK_DIR)),
    ) as txn:
        if txn.is_committed():
            logger.warning("Already committed for %s", txn_key)
            if tmp_path.exists():
                tmp_path.unlink(missing_ok=True)
            return

        os.replace(tmp_path, final_path)  # atomic rename on POSIX
        <http://logger.info|logger.info>("Prepared final file for %s; leaving transaction block to commit", txn_key)

@save_zip.on_commit
def committed(txn: Transaction):
    logger = get_run_logger()
    <http://logger.info|logger.info>(f"Committed transaction: {txn.key}")

@save_zip.on_rollback
def cleanup(txn: Transaction):
    logger = get_run_logger()
    logger.warning(f"Rolled back transaction: {txn.key}")

Tips if it still looks quiet: - Try a minimal, fresh-key test to prove hooks run:

@task(log_prints=True)
def t():
    with transaction(key="hook-test-123", store=ResultStore(lock_manager=FileSystemLockManager(lock_files_directory=LOCK_DIR))):
        pass

@t.on_commit
def _c(txn: Transaction):
    get_run_logger().info(f"Commit hook fired for {txn.key}")

- Make sure the Dask workers have the same Prefect version and can import your module (you’ve set PYTHONPATH — good). If you’re creating workers via PBS, ensure the PBS job passes that env (e.g., jobqueue

env_extra

). - Alternative registration: you can also embed hooks via options so they’re serialized with the task object:

def _commit(txn: Transaction):
    get_run_logger().info(f"Commit hook fired for {txn.key}")

def _rollback(txn: Transaction):
    get_run_logger().warning(f"Rollback hook fired for {txn.key}")

save_zip = save_zip.with_options(on_commit=[_commit], on_rollback=[_rollback])

Docs for deeper reference: - Transactions (commit/rollback hooks) - Runtime context and logging - Using DaskTaskRunner If you try the fresh

txn_key

and logger-based hooks and still don’t see logs, share a snippet of the task run logs (and your Prefect/Dask versions) and I’ll help dig in.

You’re right to question that —

with_options(on_commit=..., on_rollback=...)

is not supported in Prefect 3.x. There isn’t a version where those transaction hooks were added to `with_options`; only state hooks like

on_completion

and

on_failure

are supported there. Supported ways to register transaction hooks in 3.4.1 (and current 3.x): - Use the task’s decorator methods

@save_zip.on_commit
def _commit(txn: Transaction):
    get_run_logger().info(f"Commit hook fired for {txn.key}")

@save_zip.on_rollback
def _rollback(txn: Transaction):
    get_run_logger().warning(f"Rollback hook fired for {txn.key}")

- Keep these hook functions at module scope so Dask workers can import them. References: - Transactions and hooks: Transactions (commit/rollback hooks) - Task options and state hooks supported by with_options: State change hooks

Short answer: when you use the transaction context manager in Prefect 3.x, the lock is acquired on entry and automatically released on commit or rollback in a finally block. If your task raises an exception (including Prefect cancellations), the lock will be released. Longer answer and best practices: - Normal and error paths - The

with transaction(...) as txn:

context acquires the lock at entry and releases it on both

commit()

and

rollback()

. Any exception inside the block triggers a rollback which releases the lock. - You generally don’t need to manually release the lock in typical failure paths. - Docs: Transactions, API: prefect.transactions, Locking: FileSystemLockManager, Protocol: LockManager protocol - Fatal failures (process death, SIGKILL, OOM, hard crash) - If the Python process dies and the

__exit__

of the context can’t run, the lock file can be left behind. FileSystemLockManager does not have automatic cleanup, TTLs, or heartbeats. - Your options: 1) Catch graceful termination and roll back - If your environment sends SIGTERM on shutdown/scale-in, register a signal handler that calls

txn.rollback()

so the lock is released before exit:

import signal

       current_txn = {"obj": None}

       def _graceful_shutdown(signum, frame):
           txn = current_txn["obj"]
           if txn is not None:
               try:
                   txn.rollback()
               except Exception:
                   pass  # best effort

       signal.signal(signal.SIGTERM, _graceful_shutdown)
       signal.signal(signal.SIGINT, _graceful_shutdown)

       with transaction(
           key=f"{zip_path.stem}",
           isolation_level=IsolationLevel.SERIALIZABLE,
           write_on_commit=True,
           store=ResultStore(
               lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
           ),
       ) as txn:
           current_txn["obj"] = txn
           try:
               # do processing
               ...
           finally:
               current_txn["obj"] = None

- Note: SIGKILL (-9) and hard crashes cannot be intercepted. 2) Add a stale-lock cleanup step - On process start (or as a periodic maintenance job), scan the lock directory and remove lock files you deem stale. With FileSystemLockManager there’s no forced unlock API—you either: - Call

store.release_lock(key, holder)

if you can determine the holder; or - Carefully delete the stale file from the lock directory when you’re certain the holder process is gone. - A pragmatic approach is to: - Only use FileSystemLockManager on a single host - Encode the key deterministically (as you have with

zip_path.stem

) - If the holder in the file references a PID/hostname that no longer exists, treat it as stale and delete the file 3) Consider a lock manager with expirations for resilience - If you want locks to auto-expire, use a distributed lock manager that supports TTLs (e.g., Redis). The

prefect-redis

integration provides a

RedisLockManager

that can be configured with expiration. - Important caveat: transactions do not currently pass a lock TTL; with the built-in ResultStore/transaction path, FileSystem and Memory managers hold locks indefinitely. If you need TTL-based expiry, you may need to manage the lock explicitly with your lock manager before/around the transaction, or design a cleanup policy appropriate to your environment. - Integration docs: prefect-redis A couple of clarifying questions to help tailor the recommendation:

- By “fatally wrong,” do you mean clean shutdowns (SIGTERM), hard kills (SIGKILL), container OOM, or node crashes? - Are all tasks that contend for this key running on the same host (FileSystemLockManager is single-host), or do you need cross-host locking? - Are you seeing stale lock files today, or is this a preventative design question? If you share those details, I can suggest a concrete cleanup pattern (and code) that fits your deployment.

Thanks — that context helps a lot. Given: - Task run timeouts are occurring (TimedOut) - You need cross-host locking - You’re seeing stale lock files with FileSystemLockManager Key points: - FileSystemLockManager is single-host oriented. On timeouts that kill the process or hard failures, it can leave stale lock files. There’s no TTL/heartbeat to clean them up automatically. - For cross-host locking with automatic expiry, switch to a distributed lock manager that supports TTL (e.g., Redis via prefect-redis). With a TTL, the lock will auto-expire even if the process dies. Recommended patterns 1) Easiest and most robust: use Redis with a TTL-backed lock - Install: prefect-redis - Use RedisLockManager to acquire a lock with a hold_timeout (TTL) that’s a bit longer than your task’s worst-case runtime. - Because the transaction context does not pass a TTL to the lock manager, do the lock acquisition explicitly before the transaction, then use the transaction at READ_COMMITTED (so it won’t try to lock again). Release in a finally, and if the task dies, Redis will auto-expire the lock. Example:

from pathlib import Path
from prefect.results import ResultStore
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.results import ResultStore

# requires `pip install prefect-redis`
from prefect_redis.locking import RedisLockManager

def process_zip(zip_path: Path):
    key = zip_path.stem
    store = ResultStore(
        lock_manager=RedisLockManager(
            host="your-redis-host", port=6379, db=0,
            # configure auth/SSL as needed
        )
    )

    holder = ResultStore.generate_default_holder()

    # Acquire the distributed lock with TTL (hold_timeout)
    # Set acquire_timeout to how long you’re willing to wait for another run to finish
    acquired = store.lock_manager.acquire_lock(
        key=key,
        holder=holder,
        acquire_timeout=60,   # wait up to 60s to acquire
        hold_timeout=7200,    # lock auto-expires after 2h if we die
    )
    if not acquired:
        # choose your policy: skip, retry later, or wait_for_lock
        return

    try:
        # Use READ_COMMITTED so the transaction doesn’t try to acquire a second lock
        with transaction(
            key=key,
            isolation_level=IsolationLevel.READ_COMMITTED,
            write_on_commit=True,
            store=store,
        ) as txn:
            # do processing
            ...
    finally:
        # Best-effort release; if we crashed, Redis TTL will release later anyway
        try:
            store.lock_manager.release_lock(key, holder)
        except Exception:
            pass

2) If you must stay on filesystem locks - Catch graceful shutdown signals and roll back the transaction, so the lock is released before exit:

import signal
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager
from pathlib import Path

current_txn = {"obj": None}

def _graceful_shutdown(signum, frame):
    txn = current_txn["obj"]
    if txn is not None:
        try:
            txn.rollback()
        except Exception:
            pass

signal.signal(signal.SIGTERM, _graceful_shutdown)
signal.signal(signal.SIGINT, _graceful_shutdown)

with transaction(
    key=f"{zip_path.stem}",
    isolation_level=IsolationLevel.SERIALIZABLE,
    write_on_commit=True,
    store=ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
    ),
) as txn:
    current_txn["obj"] = txn
    try:
        # do processing
        ...
    finally:
        current_txn["obj"] = None

- And add a periodic stale-lock cleanup for the directory (e.g., delete .lock files older than your max runtime). This is inherently racy on a shared FS; Redis is a better fit for cross-host. Why Redis (or similar) is best here

- Cross-host safe, avoids shared filesystem quirks - TTL-based cleanup handles process death, OOMs, SIGKILL, node failures without manual intervention References - Transactions: Docs | API - Locking protocol: API - FileSystemLockManager: API - Redis integration: prefect-redis Open questions to tailor this further: - Can you run a Redis (or similar) service in your environment? - What’s your typical and worst-case runtime for these tasks? We can size the TTL and acquire timeouts accordingly. - Do you want contention to wait or to fail fast when another run holds the lock? We can tune

acquire_timeout

behavior.

That plan can work as a pragmatic stopgap without Redis. Here’s how I’d structure it and the tradeoffs to be aware of. What you’ll do - Add a “stale lock cleanup” task at the start of your flow that deletes old .lock files in your shared lock dir. - Run your transactional task(s) as-is; if any task times out or fails, let that failure bubble so the flow fails and retries at the flow level. That guarantees the cleanup task runs again on the next attempt. Important caveats - FileSystemLockManager is best on a single host. It can work cross-host only if all workers share the same directory and your filesystem provides reliable atomic file operations. Network filesystems (e.g., some NFS configs) can be flaky here. - You’re choosing time-based cleanup. Pick a max-age comfortably larger than your worst-case task runtime to avoid deleting valid locks. Example implementation - Cleanup task: remove stale .lock files by mtime

from pathlib import Path
import time
from prefect import task, get_run_logger

LOCK_DIR = "/scratch/tmp/my_path/lockdir"

@task
def cleanup_stale_locks(lock_dir: str = LOCK_DIR, max_age_seconds: int = 4 * 3600):
    """
    Deletes *.lock files older than max_age_seconds.
    Choose max_age_seconds > your worst-case task runtime.
    """
    logger = get_run_logger()
    now = time.time()
    lock_path = Path(lock_dir)
    removed = []
    for p in lock_path.glob("*.lock"):
        try:
            age = now - p.stat().st_mtime
            if age > max_age_seconds:
                p.unlink(missing_ok=True)
                removed.append(p.name)
        except FileNotFoundError:
            # concurrently removed
            pass
        except Exception as e:
            logger.warning(f"Could not evaluate/remove {p}: {e}")
    <http://logger.info|logger.info>(f"Removed {len(removed)} stale locks: {removed}")
    return removed

- Transactional task

from pathlib import Path
from prefect import task
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager

@task(timeout_seconds=4500)
def process_zip(zip_path: str, txn_lock_dir: str = LOCK_DIR):
    key = Path(zip_path).stem
    store = ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
    )
    with transaction(
        key=key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        write_on_commit=True,
        store=store,
    ) as txn:
        # do processing
        ...

- Flow: fail fast if any transactional task is not COMPLETED, and let flow retries handle another cleanup + rerun

from prefect import flow

@flow(retries=2, retry_delay_seconds=600)
def main(zip_paths: list[str], lock_dir: str = LOCK_DIR, lock_max_age_s: int = 4*3600):
    # 1) Cleanup first
    cleanup_stale_locks.submit(lock_dir, lock_max_age_s).result()

    # 2) Run transactional tasks
    futures = [process_zip.submit(p, lock_dir) for p in zip_paths]

    # 3) Enforce failure if any task failed/timed out
    #    .result() will raise on failure/timed out, failing the flow so it retries
    for f in futures:
        f.result()

Additional hardening (optional) - Add a SIGTERM handler so graceful shutdowns release the lock before exit. This won’t help on SIGKILL/OOM/node death, but it reduces stale files on normal terminations. See the signal handler pattern I shared earlier. - If you want to be more selective than “age-based” cleanup, you can read the lock file contents to see the recorded holder (typically includes hostname/pid). If the age is high and the hostname/pid is clearly dead, delete; otherwise, keep. I’m keeping the example simple and conservative with mtime-based deletion. - Consider running the cleanup as a separate periodic maintenance flow too, so you’re not only relying on retries. References

- Transactions: Docs | API - Locking protocol: API - FileSystemLockManager: API Questions to confirm - Is your lock directory on a shared filesystem visible to all workers, and does it provide reliable atomic file creation? If not, we should revisit the approach. - What’s your worst-case runtime? I’ll suggest a concrete

max_age_seconds

based on that.

Great — with a worst-case runtime of ~20 minutes and a shared filesystem: - Keep your “stale lock cleanup” task at the start. Set max_age to something comfortably larger than worst-case runtime. I’d use 60–90 minutes. Example: 3600 seconds. - To “carry forward” successful task outputs from a failed flow attempt to the next retry, use task result persistence + caching with a stable cache key. Then the retry will reuse results from previously COMPLETED tasks and only re-run the ones that failed. Your final successful retry can return the full, combined outputs without having to fetch anything from the failed attempt. Why this works - In Prefect 3.x, flow return values from FAILED runs are not persisted, but task results are persisted when

persist_result=True

. - With caching and a stable

cache_key_fn

, a retried flow will load previously completed task results from result storage instead of re-running those tasks. - Because you have a shared filesystem, point task result storage at a shared path so all workers/attempts can see prior results. Concrete setup 1) Stale-lock cleanup task (60 minutes max age)

from pathlib import Path
import time
from prefect import task, get_run_logger

LOCK_DIR = "/scratch/tmp/my_path/lockdir"

@task
def cleanup_stale_locks(lock_dir: str = LOCK_DIR, max_age_seconds: int = 3600):
    """
    Deletes *.lock files older than max_age_seconds.
    Choose > worst-case runtime (20m) for safety.
    """
    logger = get_run_logger()
    now = time.time()
    lock_path = Path(lock_dir)
    removed = []
    for p in lock_path.glob("*.lock"):
        try:
            age = now - p.stat().st_mtime
            if age > max_age_seconds:
                p.unlink(missing_ok=True)
                removed.append(p.name)
        except FileNotFoundError:
            pass
        except Exception as e:
            logger.warning(f"Could not evaluate/remove {p}: {e}")
    <http://logger.info|logger.info>(f"Removed {len(removed)} stale locks: {removed}")
    return removed

2) Task: persist results, cache with stable key, and use filesystem locks - Make sure the result storage is shared across hosts. For example, a shared LocalFileSystem path.

from pathlib import Path
from prefect import task
from prefect.cache_policies import CachePolicy
from prefect.filesystems import LocalFileSystem
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager

# Shared storage visible on all workers
SHARED_STORAGE = LocalFileSystem(basepath="/scratch/tmp/prefect_storage")

@task(
    timeout_seconds=4500,
    persist_result=True,
    result_storage=SHARED_STORAGE,
    cache_policy=CachePolicy.DEFAULT,
    cache_key_fn=lambda ctx, params: f"process_zip:{Path(params['zip_path']).name}",
)
def process_zip(zip_path: str, txn_lock_dir: str = LOCK_DIR):
    key = Path(zip_path).stem
    store = ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
    )
    with transaction(
        key=key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        write_on_commit=True,
        store=store,
    ) as txn:
        # do processing
        # return a serializable "output" for this zip
        ...
        return {"zip": key, "status": "ok", "data": "..."}

3) Flow: run cleanup, submit tasks, collect outputs; fail the flow if too many failures (so it retries). On the retry, cached results from successful tasks are reused and combined with new successes automatically. ``` from prefect import flow from prefect.futures import wait @flow(retries=2, retry_delay_seconds=600, persist_result=True) def main(zip_paths: list[str], lock_dir: str = LOCK_DIR, lock_max_age_s: int = 3600): # 1) Clean up stale locks first cleanup_stale_locks.submit(lock_dir, lock_max_age_s).result() # 2) Submit all tasks futures = [process_zip.submit(p, lock_dir) for p in zip_paths]

# 3) Wait for completion of submitted tasks finished, not_done = wait(futures) outputs = [] failures = 0 for f in finished: if f.state.is_completed(): # Will resolve from cache/storage if already computed in a prior attempt outputs.append(f.result()) else: failures += 1 # 4) Policy: if too many failures, fail the flow to trigger retry if failures > 1: # Optional: persist partial outputs somewhere durable (artifact or your own store) before raising # from prefect.artifacts import create_table_artifact # create_table_artifact(key="partial-outputs", table=outputs) raise RuntimeError(f"Too many failures: {failures}") return outputs ``` What you’ll get - Attempt 1 (FAILED): flow return value is not persisted, but any task that reached COMPLETED has its result in shared storage, keyed by your

cache_key_fn

. - Attempt 2 (retry): COMPLETED tasks from attempt 1 will be read from cache; only failed ones will re-run. The final successful flow return will include the full set of outputs (previous successes + new successes), so you don’t have to “pull” outputs from the failed attempt manually. Notes and gotchas - Make sure your outputs are serializable and small enough for your result storage backend. - If you still want access to partial outputs from a failed attempt for observability, create an artifact before raising the failure. Artifacts persist regardless of flow final state. See API: Artifacts API. - For cross-host filesystem locks: ensure your shared FS supports atomic file creation; otherwise, prefer a distributed lock manager when you’re ready. References - Results and persistence: Advanced results - Caching: Advanced caching - Transactions: Transactions If you’d like, share your exact shared storage path and I can tailor the

LocalFileSystem

and any serializer settings.

You’ve got two good options: A) Keep submitting all zips each attempt but avoid re-execution with task caching - Submit all

process_zip

tasks on each attempt. With a stable

cache_key_fn

and shared

result_storage

, Prefect will return cached results for already-successful zips without re-running them. - Your flow can still collect all outputs with `f.result()`; cached tasks resolve quickly and don’t re-execute. - This is the simplest approach if “unnecessary” means “don’t re-run work” (even if tasks are submitted). B) Don’t even submit already-successful zips, and still return all outputs - Persist each zip’s output to a per-zip file in shared storage (e.g.,

/scratch/tmp/prefect_storage/zip_outputs/{zip}.json

) when the task succeeds. Use an atomic write (temp file then rename). - At the start of the flow, list which zips already have output files and filter them out so you only submit tasks for the remaining zips. - At the end, combine: - existing outputs loaded from disk for already-successful zips - new outputs from the tasks that ran this attempt - If the flow fails and retries, the previously-successful outputs are already on disk; the retry will only run failed/missing zips and merge everything for the final return value. Below is a concrete pattern for option B that also includes your stale-lock cleanup. Stale lock cleanup (60 min > 20 min max runtime)

from pathlib import Path
import time
from prefect import task, get_run_logger

LOCK_DIR = "/scratch/tmp/my_path/lockdir"

@task
def cleanup_stale_locks(lock_dir: str = LOCK_DIR, max_age_seconds: int = 3600):
    logger = get_run_logger()
    now = time.time()
    lock_path = Path(lock_dir)
    removed = []
    for p in lock_path.glob("*.lock"):
        try:
            age = now - p.stat().st_mtime
            if age > max_age_seconds:
                p.unlink(missing_ok=True)
                removed.append(p.name)
        except FileNotFoundError:
            pass
        except Exception as e:
            logger.warning(f"Could not evaluate/remove {p}: {e}")
    <http://logger.info|logger.info>(f"Removed {len(removed)} stale locks: {removed}")
    return removed

Persist per-zip outputs and use transactions

import json
from pathlib import Path
from prefect import task
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager

OUTPUT_DIR = "/scratch/tmp/prefect_storage/zip_outputs"

def _output_path_for(zip_path: str, output_dir: str = OUTPUT_DIR) -> Path:
    return Path(output_dir) / f"{Path(zip_path).stem}.json"

@task(timeout_seconds=4500, persist_result=True)
def process_zip(zip_path: str, txn_lock_dir: str = LOCK_DIR, output_dir: str = OUTPUT_DIR):
    key = Path(zip_path).stem
    store = ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
    )
    # Ensure output dir exists
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    with transaction(
        key=key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        write_on_commit=True,
        store=store,
    ):
        # do processing
        result = {"zip": key, "status": "ok", "data": "..."}  # must be JSON-serializable

        # Atomically persist per-zip output so retries can detect/merge it
        out_path = _output_path_for(zip_path, output_dir)
        tmp_path = out_path.with_suffix(".json.tmp")
        tmp_path.write_text(json.dumps(result))
        tmp_path.replace(out_path)

        return result

Helpers to detect and load existing outputs ``` @task def list_completed_zips(output_dir: str = OUTPUT_DIR) -> set[str]: return {p.stem for p in Path(output_dir).glob("*.json")} @task def load_outputs_for(zip_paths: list[str], output_dir: str = OUTPUT_DIR): outputs = [] for zp in zip_paths: key = Path(zp).stem p = Path(output_dir) / f"{key}.json" if p.exists():

outputs.append(json.loads(p.read_text())) return outputs

Flow: only run missing zips; collect all outputs across attempts

from prefect import flow from prefect.futures import wait @flow(retries=2, retry_delay_seconds=600, persist_result=True) def main(zip_paths: list[str], lock_dir: str = LOCK_DIR, lock_max_age_s: int = 3600, output_dir: str = OUTPUT_DIR): # 1) Cleanup stale locks cleanup_stale_locks.submit(lock_dir, lock_max_age_s).result() # 2) Determine which zips already have outputs already_done = list_completed_zips.submit(output_dir).result() # 3) Submit only the zips that still need processing to_run = [zp for zp in zip_paths if Path(zp).stem not in already_done] futures = [process_zip.submit(zp, lock_dir, output_dir) for zp in to_run] # 4) Wait and collect new outputs finished, _ = wait(futures) new_outputs = [f.result() for f in finished if f.state.is_completed()] failures = sum(1 for f in finished if not f.state.is_completed()) # 5) Load existing outputs for previously completed zips prev_outputs = load_outputs_for.submit([zp for zp in zip_paths if Path(zp).stem in already_done], output_dir).result() # 6) Merge and decide whether to retry outputs = prev_outputs + new_outputs if failures > 1: # outputs are still safely persisted per-zip on disk for next retry raise RuntimeError(f"Too many failures: {failures}") return outputs ``` This pattern: - Avoids submitting tasks for already-completed zips - Persists per-zip outputs so a retry can merge previously successful outputs with newly produced ones - Ensures your final successful attempt returns the complete set If you prefer the simpler path with task caching (Option A), just submit all tasks each attempt with a stable

cache_key_fn

and shared `result_storage`; cached tasks won’t re-run, and your flow can return all outputs by collecting

f.result()

. Docs for reference: - Transactions: Transactions - Advanced results and storage: Results - Caching: Caching Want me to adapt this to a different storage target (e.g., S3) or to add a checksum-based cache key so renaming the zip won’t trigger reprocessing?

Yes — that’s a reasonable approach. Put a conservative, key-scoped stale-lock cleanup step at the top of your

process_zip

task, then proceed into the transaction. On task retries, that preflight cleanup will run again so a stale file won’t block the next attempt. Key points and guardrails - Only delete as a last resort and with a generous max age. Since your worst-case runtime is ~20 minutes, use 60+ minutes (e.g., 3600s) to minimize the risk of deleting a legitimate in-progress lock. - Prefer to “wait a little” first. If another run is legitimately finishing, a short

wait_for_lock()

avoids unnecessary deletes. - This is best-effort; FileSystemLockManager does not expose a “force release” or a standard way to read the holder safely across hosts. You’re relying on the file’s mtime and your operational knowledge. Example: preflight cleanup inside the task, then the transaction

import time
from pathlib import Path
from prefect import task, get_run_logger
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager

LOCK_DIR = "/scratch/tmp/my_path/lockdir"

def _lock_file_path(lock_dir: str, key: str) -> Path:
    # Based on your observed filenames: "<key>.lock"
    return Path(lock_dir) / f"{key}.lock"

def cleanup_stale_lock_for_key(
    key: str,
    lock_dir: str,
    max_age_seconds: int = 3600,         # >= 3x worst-case runtime
    wait_grace_seconds: float = 10.0,    # brief grace wait to let legitimate holders finish
) -> bool:
    """
    Best-effort cleanup of a stale lock for a specific key.
    Returns True if a stale file was removed.
    """
    logger = get_run_logger()
    lm = FileSystemLockManager(lock_files_directory=Path(lock_dir))

    # 1) Grace period: wait briefly for any legitimate holder to finish
    released = lm.wait_for_lock(key, timeout=wait_grace_seconds)
    if released:
        return False  # lock became available; nothing to delete

    # 2) Check on-disk file age; delete only if very old
    p = _lock_file_path(lock_dir, key)
    if p.exists():
        try:
            age = time.time() - p.stat().st_mtime
        except FileNotFoundError:
            return False  # concurrently deleted
        if age > max_age_seconds:
            try:
                p.unlink(missing_ok=True)
                logger.warning(f"Removed stale lock file for key={key}: {p} (age={age:.0f}s)")
                return True
            except Exception as e:
                logger.warning(f"Failed to remove stale lock file {p}: {e}")
    return False

@task(timeout_seconds=4500, retries=2, retry_delay_seconds=300, persist_result=True)
def process_zip(zip_path: str, txn_lock_dir: str = LOCK_DIR, max_lock_age_s: int = 3600):
    key = Path(zip_path).stem

    # Preflight stale-lock cleanup for this specific key
    cleanup_stale_lock_for_key(key, txn_lock_dir, max_age_seconds=max_lock_age_s)

    store = ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
    )

    # Now proceed with the transactional work
    with transaction(
        key=key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        write_on_commit=True,
        store=store,
    ):
        # do processing
        ...
        return {"zip": key, "status": "ok"}

Notes - If you prefer not to derive the file path yourself, you can rely solely on

wait_for_lock()

with a longer timeout and skip deletion entirely. But given you’re seeing stale files from timeouts/crashes, the age-based delete is the practical way to unblock retries. - Keep the per-flow “stale lock cleanup” you already added as a second line of defense for any locks not associated with the current attempt’s keys. References - Transactions: Docs | API

- Locking protocol (methods like

wait_for_lock

/

is_locked

): API - FileSystemLockManager: API If you want, I can adapt the cleanup to parse the lockfile contents and try to be more selective (e.g., check age plus host), but file format isn’t officially documented, so the age-based approach is the safest portable option.

Short answer - Transactions only have on_commit and on_rollback hooks. There is no transaction-level on_failure hook. - Task/flow decorators do have on_failure hooks, but they only run if Prefect’s engine is still executing your code when the state changes. If a worker/process is killed (OOM, Dask worker evicted, SIGKILL), hooks will not run. What this means for your cases - Processing errors inside the

with transaction(...)

block: the context exits,

rollback()

is called, and your on_rollback hooks run. - Prefect-managed timeouts that raise in-process (e.g., a TaskRunTimeoutError the engine can handle): task/flow

on_failure

hooks will run and can be used for logging or best-effort cleanup. - Hard failures (worker killed, node crash, OOM, SIGKILL): neither transaction hooks nor task/flow hooks will run because your process is gone. This is exactly where stale lock files come from. Practical recommendations - Keep using transaction on_rollback hooks for normal processing failures. - Add a preflight cleanup at the start of your task (key-scoped, age-based, with a short wait) so retries can acquire the lock even after a hard failure. This is the only reliably executed place after a hard crash. - Optionally add task/flow

on_failure

hooks for observability and best-effort cleanup when Prefect can still run your code (e.g., timeouts that don’t kill the process), but don’t rely on them for crash cleanup. Example: on_rollback hook plus preflight cleanup

from pathlib import Path
import time
from prefect import task, get_run_logger
from prefect.transactions import transaction, IsolationLevel
from prefect.results import ResultStore
from prefect.locking.filesystem import FileSystemLockManager

LOCK_DIR = "/scratch/tmp/my_path/lockdir"

def cleanup_stale_lock_for_key(
    key: str,
    lock_dir: str,
    max_age_seconds: int = 3600,
    wait_grace_seconds: float = 10.0,
) -> bool:
    logger = get_run_logger()
    lm = FileSystemLockManager(lock_files_directory=Path(lock_dir))

    # Wait briefly in case a legitimate holder is finishing
    if lm.wait_for_lock(key, timeout=wait_grace_seconds):
        return False

    # Conservative age-based cleanup
    lock_path = Path(lock_dir) / f"{key}.lock"
    if lock_path.exists():
        try:
            age = time.time() - lock_path.stat().st_mtime
        except FileNotFoundError:
            return False
        if age > max_age_seconds:
            try:
                lock_path.unlink(missing_ok=True)
                logger.warning(f"Removed stale lock file for key={key}: {lock_path} (age={age:.0f}s)")
                return True
            except Exception as e:
                logger.warning(f"Failed to remove stale lock file {lock_path}: {e}")
    return False

@task(timeout_seconds=4500, retries=2, retry_delay_seconds=300, persist_result=True)
def process_zip(zip_path: str, txn_lock_dir: str = LOCK_DIR, max_lock_age_s: int = 3600):
    logger = get_run_logger()
    key = Path(zip_path).stem

    # Preflight cleanup to unblock retries after hard failures
    cleanup_stale_lock_for_key(key, txn_lock_dir, max_age_seconds=max_lock_age_s)

    def on_rollback(txn):
        logger.warning(f"Rollback for key={key}; releasing any partial artifacts if needed")

    store = ResultStore(
        lock_manager=FileSystemLockManager(lock_files_directory=Path(txn_lock_dir))
    )

    with transaction(
        key=key,
        isolation_level=IsolationLevel.SERIALIZABLE,
        write_on_commit=True,
        store=store,
    ) as txn:
        txn.on_rollback_hooks.append(on_rollback)

        # do processing
        ...
        return {"zip": key, "status": "ok"}

When

on_failure

is useful - You can attach hooks on the task/flow for additional telemetry or best-effort cleanup when the engine is able to run them (timeouts that don’t fully kill the process, typical failures): - Task API:

@task(on_failure=[...])

- Flow API:

@flow(on_failure=[...], on_crashed=[...], on_cancellation=[...])

Docs: flows API, tasks API, transactions Docs. Bottom line - There’s no transaction on_failure hook; use on_rollback for in-process failures. - For the non-processing failures you described (Dask worker killed, OOM), preflight cleanup at task start (and/or a periodic directory cleanup) is the reliable mechanism. Hooks won’t run when the process is gone.