def get_piston_client_from_env()

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

• 17 lines of code
• 4 McCabe index (conditional complexity)

def get_piston_client_from_env(session=None):
    piston_endpoints = os.getenv("PISTON_ENDPOINTS")
    if piston_endpoints is None:
        raise ValueError(
            "For IOI/CF problems Piston endpoints running our IOI package are required. Please add a list of valid Piston endpoints to a PISTON_ENDPOINTS variable in a `.env` file."
        )
    piston_endpoints = sorted(
        piston_endpoints.split(",") if piston_endpoints != "slurm" else get_slurm_piston_endpoints()
    )
    gpu_nb = int(os.getenv("LOCAL_RANK", 0))  # per‑GPU index
    world = int(os.getenv("WORLD_SIZE", 1))  # total GPUs
    if world > 1:
        print(f"Using a subset of piston endpoints for GPU#{gpu_nb}")
        piston_endpoints = piston_endpoints[gpu_nb::world]
    random.shuffle(piston_endpoints)
    max_requests_per_endpoint = os.getenv("PISTON_MAX_REQUESTS_PER_ENDPOINT", "1")
    return PistonClient(piston_endpoints, session, max_requests_per_endpoint=int(max_requests_per_endpoint))