PythonVenv tool (#3094)

* Add Lisa Tool for Python Virtualenv

* GPU Pytorch case: Replace Python and Pip tool with PythonVenv

* PythonVenv: Accept path in init

* Apply lint

* PythonVenv: Change default name for venv

Change default name from 'venv' to timestamp. This is to
avoid the side effects of env being resued by multiple tests

* PythonVenv: Remove default venv path

* PythonVenv: Move Public methods above pvt methods

* PythonVenv: Fix PurePath for node

* GPU Pytorch: Increase required size to 20GB

* PythonVenv: Use local cache and build path

Use venv path as cache and build path, this is to
avoid OS disk out of space issue.

* Implement package cleanup

Implement package cleanup and invoke during GPU Pytorch test.

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
+        )

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)

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: