diff --git a/lisa/tools/python.py b/lisa/tools/python.py index 901d3f0c0..04c203f01 100644 --- a/lisa/tools/python.py +++ b/lisa/tools/python.py @@ -1,11 +1,16 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. import re -from typing import List, Type +from pathlib import PurePath +from typing import TYPE_CHECKING, List, Type from assertpy import assert_that from lisa.executable import Tool + +if TYPE_CHECKING: + from lisa.node import Node + from lisa.operating_system import Posix from lisa.tools.mkdir import Mkdir from lisa.util import UnsupportedDistroException, get_matched_str @@ -89,3 +94,79 @@ class Pip(Tool): def uninstall_package(self, package_name: str) -> bool: result = self.run(f"uninstall {package_name} -y", force_run=True, sudo=True) return result.exit_code == 0 + + +class PythonVenv(Tool): + @property + def command(self) -> str: + path = self.get_venv_path() / "bin" / self._python.command + return str(path) + + @property + def can_install(self) -> bool: + return True + + @property + def dependencies(self) -> List[Type[Tool]]: + return [Python] + + def __init__(self, node: "Node", venv_path: str) -> None: + super().__init__(node) + self._python: Python = self.node.tools[Python] + self._venv_installation_path = venv_path + + def _install(self) -> bool: + if isinstance(self.node.os, Posix): + self.node.os.install_packages("python3-venv") + return self._check_exists() + + def get_venv_path(self) -> PurePath: + if not hasattr(self, "_venv_path"): + self._venv_path = self._create_venv(self._venv_installation_path) + return self._venv_path + + def install_packages(self, packages_name: str) -> None: + venv_path = self.get_venv_path() + cache_dir = venv_path.joinpath("cache") + self.node.tools[Mkdir].create_directory(str(cache_dir)) + envs = {"TMPDIR": str(cache_dir)} + cmd_result = self.run( + f"-m pip install -q {packages_name} --cache-dir={cache_dir}", + force_run=True, + update_envs=envs, + ) + assert_that( + cmd_result.exit_code, f"fail to install {packages_name}" + ).is_equal_to(0) + + def exists_package(self, package_name: str) -> bool: + result = self.run(f"-m pip show {package_name}") + return result.exit_code == 0 + + def uninstall_package(self, package_name: str) -> bool: + result = self.run(f"-m pip uninstall {package_name} -y", force_run=True) + return result.exit_code == 0 + + def delete_venv(self) -> None: + if hasattr(self, "_venv_path"): + self.node.execute(f"rm -rf {self._venv_path}") + delattr(self, "_venv_path") + else: + self._log.info("venv path not found, nothing to delete") + + def _create_venv(self, venv_path: str) -> PurePath: + cmd_result = self._python.run( + f"-m venv {venv_path}", force_run=True, shell=True + ) + assert_that( + cmd_result.exit_code, f"fail to create venv: {venv_path}" + ).is_equal_to(0) + self._venv_path = self.node.get_pure_path(venv_path) + return self._venv_path + + def _check_exists(self) -> bool: + venv = self._python.run("-m venv --help", force_run=True) + ensurepip = self._python.run("-m ensurepip", force_run=True) + return ( + venv.exit_code == 0 and "No module named ensurepip" not in ensurepip.stdout + ) diff --git a/microsoft/testsuites/gpu/gpusuite.py b/microsoft/testsuites/gpu/gpusuite.py index b9354e83b..7556114d6 100644 --- a/microsoft/testsuites/gpu/gpusuite.py +++ b/microsoft/testsuites/gpu/gpusuite.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -import os import re from pathlib import Path from typing import Any, List @@ -20,9 +19,19 @@ from lisa import ( ) from lisa.features import Gpu, GpuEnabled, SerialConsole, StartStop from lisa.features.gpu import ComputeSDK -from lisa.operating_system import BSD, AlmaLinux, Debian, Oracle, Suse, Ubuntu, Windows +from lisa.operating_system import ( + BSD, + AlmaLinux, + Debian, + Linux, + Oracle, + Suse, + Ubuntu, + Windows, +) from lisa.sut_orchestrator.azure.features import AzureExtension -from lisa.tools import Lspci, Mkdir, NvidiaSmi, Pip, Python, Reboot, Service, Tar, Wget +from lisa.tools import Lspci, Mkdir, NvidiaSmi, Reboot, Service, Tar, Wget +from lisa.tools.python import PythonVenv from lisa.util import UnsupportedOperationException, get_matched_str _cudnn_location = ( @@ -262,49 +271,40 @@ class GpuTestSuite(TestSuite): _install_driver(node, log_path, log) _check_driver_installed(node, log) - # Step 1, pytorch/CUDA needs 8GB to download & install, increase to 16GB - torch_required_space = 16 + # Step 1, pytorch/CUDA needs 8GB to download & install, increase to 20GB + torch_required_space = 20 work_path = node.get_working_path_with_required_space(torch_required_space) - use_new_path = work_path != str(node.working_path) # Step 2, Install cudnn and pyTorch _install_cudnn(node, log, work_path) + pythonvenv_path = work_path + "/gpu_pytorch" + pythonvenv = node.tools.create(PythonVenv, venv_path=pythonvenv_path) - pip = node.tools[Pip] - if not pip.exists_package("torch"): - if use_new_path: - pip.install_packages("torch", work_path) - else: - pip.install_packages("torch") + # Pip downloads .whl and other tmp files to root disk. + # Clean package cache to avoid disk full issue. + if isinstance(node.os, Linux): + node.os.clean_package_cache() + + pythonvenv.install_packages("torch") # Step 3, verification gpu = node.features[Gpu] gpu_script = "import torch;print(f'gpu count: {torch.cuda.device_count()}')" - python = node.tools[Python] expected_count = gpu.get_gpu_count_with_lspci() - if use_new_path: - python_path = os.environ.get("PYTHONPATH", "") - python_path += f":{work_path}/python_packages" - python_envs = {"PYTHONPATH": python_path} - else: - python_envs = {} - - script_result = python.run( + script_result = pythonvenv.run( f'-c "{gpu_script}"', force_run=True, - update_envs=python_envs, ) if script_result.exit_code != 0 and self._numpy_error_pattern.findall( script_result.stdout ): - if pip.uninstall_package("numpy"): - pip.install_packages("numpy") - script_result = python.run( + if pythonvenv.uninstall_package("numpy"): + pythonvenv.install_packages("numpy") + script_result = pythonvenv.run( f'-c "{gpu_script}"', force_run=True, - update_envs=python_envs, ) gpu_count_str = get_matched_str(script_result.stdout, self._pytorch_pattern) diff --git a/microsoft/testsuites/vm_extensions/scripts/handle.txt b/microsoft/testsuites/vm_extensions/scripts/handle.txt index 261fe4af5..ec31aec48 100644 --- a/microsoft/testsuites/vm_extensions/scripts/handle.txt +++ b/microsoft/testsuites/vm_extensions/scripts/handle.txt @@ -16,12 +16,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import fnmatch import os import os.path import string import subprocess import sys -import fnmatch + def get_folders_with_string(root_path, folder_name_contains): matching_folders = [] @@ -41,11 +42,11 @@ for folder in matching_folders: break os.chdir(folder) -from Utils import HandlerUtil -from patch import GetMyPatching from backuplogger import Backuplogger from common import CommonVariables -from Utils import SizeCalculation +from patch import GetMyPatching +from Utils import HandlerUtil, SizeCalculation + try: from unittest.mock import MagicMock except ImportError: