#!/usr/bin/env python3
"""Forza il timeout di default di PyTorch per i nuovi process group (usato quando
timeout=None, es. i gruppi TP/PP/DP/EP creati da vLLM nel worker Ray prima che il
contesto vllm_config sia disponibile - vedi indagine con debug.py) a un valore piu'
alto dei 600s hardcoded, invece dei 600s di default per NCCL.
Uso: patch-nccl-default-timeout.py <path-di-distributed_c10d.py> <secondi>"""
import sys

path = sys.argv[1]
seconds = sys.argv[2]

src = open(path).read()

marker = "def _get_default_timeout(backend: Backend) -> timedelta:\n"
override = (
    marker
    + f"    from datetime import timedelta as _td\n"
    + f"    return _td(seconds={seconds})  # override: vedi patch-nccl-default-timeout.py\n"
)

count = src.count(marker)
if count != 1:
    print(f"ERRORE: marker trovato {count} volte (atteso 1) in {path}", file=sys.stderr)
    sys.exit(1)

src = src.replace(marker, override, 1)
open(path, "w").write(src)
print(f"Patchato {path}: default timeout forzato a {seconds}s")