def get_slurm_piston_endpoints()

in src/open_r1/utils/competitive_programming/piston_client.py [0:0]


def get_slurm_piston_endpoints():
    """Get list of active piston worker endpoints from squeue output"""
    # Run squeue command to get job name, hostname and status, filtering for RUNNING state
    result = subprocess.run(
        ["squeue", '--format="%j %N %T"', "--noheader", "--states=RUNNING"], capture_output=True, text=True
    )

    # Split output into lines and skip header
    lines = result.stdout.strip().split("\n")

    endpoints = []
    for line in lines:
        # Parse job name from squeue output
        fields = line.split()
        job_name = fields[0].strip('"')  # Remove quotes
        hostname = fields[1]

        # Extract port if job name matches pattern
        match = re.match(r"piston-worker-(\d+)", job_name)
        if match:
            port = match.group(1)
            endpoints.append(f"http://{hostname}:{port}/api/v2")

    return endpoints