lablink_client_service.check_gpu

lablink_client_service.check_gpu

Attributes

MAX_REPORT_RETRIES = 5 module-attribute

REPORT_RETRY_DELAY = 10 module-attribute

logger = logging.getLogger(__name__) module-attribute

Classes

Functions:

check_gpu_health(allocator_url, interval=20, client_secret='')

Check the health of the GPU.

Parameters:

Name	Type	Description	Default
allocator_url	str	The base URL of the allocator service.	required
interval	int	The interval in seconds to check the GPU health.	20
client_secret	str	Per-client secret for Bearer auth.	''

Source code in packages/client/src/lablink_client_service/check_gpu.py

def check_gpu_health(allocator_url: str, interval: int = 20, client_secret: str = ""):
    """Check the health of the GPU.

    Args:
        allocator_url (str): The base URL of the allocator service.
        interval (int, optional): The interval in seconds to check the GPU health.
        client_secret (str, optional): Per-client secret for Bearer auth.
    """
    logger.info("Starting GPU health monitoring")
    last_status = None
    base_url = sanitize_url(allocator_url)

    headers = {"Content-Type": "application/json"}
    headers.update(get_auth_headers(client_secret))

    while True:
        curr_status = None
        break_now = False

        try:
            # Run the nvidia-smi command to check GPU health
            subprocess.run(
                ["nvidia-smi"],
                capture_output=True,
                text=True,
                check=True,
            )
            curr_status = "Healthy"

        except subprocess.CalledProcessError as e:
            # Command not found -> likely nvidia-smi is not installed
            if e.returncode == 127:
                logger.warning("nvidia-smi not found - NVIDIA drivers not installed")
                curr_status = "N/A"
                break_now = True
            else:
                logger.error(f"nvidia-smi failed: {e.stderr.strip()}")
                curr_status = "Unhealthy"

        except FileNotFoundError:
            logger.warning("nvidia-smi not found - NVIDIA drivers not installed")
            curr_status = "N/A"
            break_now = True

        except Exception as e:
            curr_status = "Unhealthy"
            logger.error(f"GPU health check failed: {e}")

        if curr_status != last_status or break_now:
            logger.info(f"GPU status: {curr_status}")
            report_retry_count = 0

            while report_retry_count < MAX_REPORT_RETRIES:
                try:
                    response = requests.post(
                        f"{base_url}/api/gpu_health",
                        json={
                            "hostname": os.getenv("VM_NAME"),
                            "gpu_status": curr_status,
                        },
                        headers=headers,
                        timeout=(10, 20),
                    )
                    response.raise_for_status()
                    logger.info(f"Reported GPU status: {curr_status}")
                    last_status = curr_status
                    break  # Break out of the report retry loop on success
                except requests.exceptions.Timeout:
                    logger.warning(
                        f"GPU health report timed out "
                        f"(attempt {report_retry_count + 1}/{MAX_REPORT_RETRIES})"
                    )
                except requests.exceptions.RequestException as e:
                    logger.warning(
                        f"Failed to report GPU health: {e} "
                        f"(attempt {report_retry_count + 1}/{MAX_REPORT_RETRIES})"
                    )
                except Exception as e:
                    logger.error(
                        f"Unexpected error reporting GPU health: {e} "
                        f"(attempt {report_retry_count + 1}/{MAX_REPORT_RETRIES})"
                    )

                report_retry_count += 1
                if report_retry_count < MAX_REPORT_RETRIES:
                    jitter = random.uniform(0, 5)
                    time.sleep(REPORT_RETRY_DELAY + jitter)
            else:
                logger.error(
                    f"Failed to report GPU health after {MAX_REPORT_RETRIES} attempts"
                )

                # Fallback to avoid repeated failed reports
                last_status = curr_status

        if break_now:
            break
        time.sleep(interval)

main(cfg)

Source code in packages/client/src/lablink_client_service/check_gpu.py

@hydra.main(version_base=None, config_name="config")
def main(cfg: Config) -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    logging.getLogger("lablink_client_service").setLevel(logging.DEBUG)

    # Check GPU health
    base_url, client_secret, _ = get_client_env(cfg)
    check_gpu_health(allocator_url=base_url, client_secret=client_secret)