"""Optional: enforce the same GPU mutex from inside a Python dispatcher.

If your GPU jobs are fired over SSH to a remote box (or from a long-running Python

loop rather than a shell), you can hold the *same* flock in-process so the wrapper

and the dispatcher share one lock. The reentrancy trick is identical to gpu_run.sh:

honor an inherited GPU_LOCK_HELD and skip re-locking, and respect a skip flag for

non-GPU work (builds, uploads) that happens to run on the same box.

This is a generic skeleton - replace the dispatch body with your own call.

"""

import fcntl

import os

import subprocess

import sys

GPU_LOCK_PATH = os.environ.get("GPU_LOCK_PATH", "/tmp/gpu.lock")

def gpu_dispatch(argv, is_gpu_job=True, timeout=600):

    need_lock = (

        is_gpu_job

        and not os.environ.get("GPU_LOCK_HELD")

        and not os.environ.get("GPU_SKIP_LOCK")

    )

    lockf = None

    try:

        if need_lock:

            lockf = open(GPU_LOCK_PATH, "w")

            sys.stderr.write("[dispatch] waiting for GPU lock...\n")

            fcntl.flock(lockf, fcntl.LOCK_EX)

            sys.stderr.write("[dispatch] acquired GPU lock\n")

            os.environ["GPU_LOCK_HELD"] = "1"

        return subprocess.run(argv, text=True, timeout=timeout).returncode

    finally:

        if lockf is not None:

            fcntl.flock(lockf, fcntl.LOCK_UN)

            lockf.close()

            os.environ.pop("GPU_LOCK_HELD", None)

if __name__ == "__main__":

    sys.exit(gpu_dispatch(sys.argv[1:]))