Best Python code snippet using lisa_python
gpusuite.py
Source:gpusuite.py  
...51        ),52        priority=1,53    )54    def verify_load_gpu_driver(self, node: Node, log_path: Path, log: Logger) -> None:55        _check_driver_installed(node)56    @TestCaseMetadata(57        description="""58            This test case verifies the gpu adapter count.59            Steps:60            1. Assert that node supports GPU.61            2. If GPU modules are not loaded, install and load the module first.62            3. Find the expected gpu count for the node.63            4. Validate expected and actual gpu count using lsvmbus output.64            5. Validate expected and actual gpu count using lspci output.65            6. Validate expected and actual gpu count using gpu vendor commands66                 example - nvidia-smi67        """,68        timeout=TIMEOUT,69        requirement=simple_requirement(70            supported_features=[GpuEnabled()],71        ),72        priority=2,73    )74    def verify_gpu_adapter_count(self, node: Node, log_path: Path, log: Logger) -> None:75        _check_driver_installed(node)76        gpu_feature = node.features[Gpu]77        assert isinstance(node.capability.gpu_count, int)78        expected_count = node.capability.gpu_count79        lsvmbus_device_count = gpu_feature.get_gpu_count_with_lsvmbus()80        assert_that(81            lsvmbus_device_count,82            "Expected device count didn't match Actual device count from lsvmbus",83        ).is_equal_to(expected_count)84        lspci_device_count = gpu_feature.get_gpu_count_with_lspci()85        assert_that(86            lspci_device_count,87            "Expected device count didn't match Actual device count from lspci",88        ).is_equal_to(expected_count)89        _check_driver_installed(node)90        vendor_cmd_device_count = gpu_feature.get_gpu_count_with_vendor_cmd()91        assert_that(92            vendor_cmd_device_count,93            "Expected device count didn't match Actual device count"94            " from vendor command",95        ).is_equal_to(expected_count)96    @TestCaseMetadata(97        description="""98        This test case will99        1. Validate disabling GPU devices.100        2. Validate enable back the disabled GPU devices.101        """,102        priority=2,103        requirement=simple_requirement(104            supported_features=[GpuEnabled()],105        ),106    )107    def verify_gpu_rescind_validation(self, node: Node) -> None:108        _check_driver_installed(node)109        lspci = node.tools[Lspci]110        gpu = node.features[Gpu]111        # 1. Disable GPU devices.112        gpu_devices = lspci.get_devices_by_type(device_type=constants.DEVICE_TYPE_GPU)113        gpu_devices = gpu.remove_virtual_gpus(gpu_devices)114        # stop the service which uses nvidia module115        service = node.tools[Service]116        service.stop_service("nvidia-persistenced")117        for device in gpu_devices:118            lspci.disable_device(device)119        # 2. Enable GPU devices.120        lspci.enable_devices()121    @TestCaseMetadata(122        description="""123        This test case will run PyTorch to check CUDA driver installed correctly.124        1. Install PyTorch.125        2. Check GPU count by torch.cuda.device_count()126        3. Compare with PCI result127        """,128        priority=3,129        requirement=simple_requirement(130            supported_features=[GpuEnabled()],131        ),132    )133    def verify_gpu_cuda_with_pytorch(self, node: Node) -> None:134        _check_driver_installed(node)135        _install_cudnn(node)136        gpu = node.features[Gpu]137        pip = node.tools[Pip]138        if not pip.exists_package("torch"):139            pip.install_packages("torch")140        gpu_script = "import torch;print(f'gpu count: {torch.cuda.device_count()}')"141        python = node.tools[Python]142        expected_count = gpu.get_gpu_count_with_lspci()143        script_result = python.run(144            f'-c "{gpu_script}"',145            force_run=True,146        )147        gpu_count_str = get_matched_str(script_result.stdout, self._pytorch_pattern)148        script_result.assert_exit_code(149            message=f"failed on run gpu script: {gpu_script}, "150            f"output: {script_result.stdout}"151        )152        assert_that(gpu_count_str).described_as(153            f"gpu count is not in result: {script_result.stdout}"154        ).is_not_empty()155        gpu_count = int(gpu_count_str)156        assert_that(gpu_count).described_as(157            "GPU must be greater than zero."158        ).is_greater_than(0)159        assert_that(gpu_count).described_as(160            "cannot detect GPU from PyTorch"161        ).is_equal_to(expected_count)162def _check_driver_installed(node: Node) -> None:163    gpu = node.features[Gpu]164    if not gpu.is_supported():165        raise SkippedException(f"GPU is not supported with distro {node.os.name}")166    if ComputeSDK.AMD in gpu.get_supported_driver():167        raise SkippedException("AMD vm sizes is not supported")168    try:169        _ = node.tools[NvidiaSmi]170    except Exception as identifier:171        raise LisaException(172            f"Cannot find nvidia-smi, make sure the driver installed correctly. "173            f"Inner exception: {identifier}"174        )175def _install_cudnn(node: Node) -> None:176    wget = node.tools[Wget]177    tar = node.tools[Tar]178    path = wget.get_tool_path(use_global=True)179    extracted_path = tar.get_tool_path(use_global=True)180    if node.shell.exists(path / _cudnn_file_name):181        return182    download_path = wget.get(183        url=_cudnn_location, filename=str(_cudnn_file_name), file_path=str(path)184    )185    tar.extract(download_path, dest_dir=str(extracted_path))186    if isinstance(node.os, Debian):187        target_path = "/usr/lib/x86_64-linux-gnu/"188    else:189        target_path = "/usr/lib64/"190    node.execute(191        f"cp -p {extracted_path}/cuda/lib64/libcudnn* {target_path}",192        shell=True,193        sudo=True,194    )195    return196# We use platform to install the driver by default. If in future, it needs to197# install independently, this logic can be reused.198def _ensure_driver_installed(199    node: Node, gpu_feature: Gpu, log_path: Path, log: Logger200) -> None:201    if gpu_feature.is_module_loaded():202        return203    gpu_feature.install_compute_sdk()204    log.debug(205        f"{gpu_feature.get_supported_driver()} sdk installed. "206        "Will reboot to load driver."207    )208    reboot_tool = node.tools[Reboot]209    reboot_tool.reboot_and_check_panic(log_path)...Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.
You could also refer to video tutorials over LambdaTest YouTube channel to get step by step demonstration from industry experts.
Get 100 minutes of automation test minutes FREE!!
