from airflow import DAG
from airflow.decorators import task
import pendulum

TARGET_HOST = "airflow-k8s-demo-0164-worker-gpu-0.airflow-k8s-demo-0164-worker-gpu"

with DAG(
    dag_id='debug_gpu_network_check',
    # UPDATED: New argument name
    schedule=None,
    # UPDATED: Explicit pendulum date
    start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
    catchup=False
) as dag:

    @task(queue='gpu')
    def check_network():
        import subprocess
        import socket
        
        print(f"I am running on hostname: {socket.gethostname()}")
        print(f"Attempting to resolve and ping: {TARGET_HOST}")
        
        # 1. DNS Check
        try:
            ip = socket.gethostbyname(TARGET_HOST)
            print(f"DNS Resolution Success: {TARGET_HOST} -> {ip}")
        except Exception as e:
            print(f"DNS Resolution Failed: {e}")
            raise

        # 2. Ping Check
        cmd = ["ping", "-c", "4", TARGET_HOST]
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("Ping Success!")
            print(result.stdout)
        else:
            print("Ping Failed (might be firewall, check if TCP port 29500 works instead)")
            print(result.stderr)

    check_network()