Venv Isolation Update (#42772)

* checks can now properly isolate themselves without a recursive call to azure-sdk-tools from a venv to another venv

Author: scbedd

eng/tools/azure-sdk-tools/README.md

@@ -217,6 +217,9 @@ class my_check(Check):
         for parsed in targeted:
             pkg_dir = parsed.folder
             pkg_name = parsed.name
+            executable, staging_directory = self.get_executable(args.isolate, args.command, sys.executable, pkg)
+            # the rest of the check should use executable, not sys.executable
+            # if a staging area is needed use staging_directory
             print(f"Processing {pkg_name} for my_check")
 ```
 

eng/tools/azure-sdk-tools/azpysdk/Check.py

@@ -2,9 +2,20 @@
 import os
 import argparse
 import traceback
-from typing import Sequence, Optional, List, Any
+import sys
+
+from typing import Sequence, Optional, List, Any, Tuple
+from subprocess import check_call
+
 from ci_tools.parsing import ParsedSetup
-from ci_tools.functions import discover_targeted_packages
+from ci_tools.functions import discover_targeted_packages, get_venv_call
+from ci_tools.variables import discover_repo_root
+from ci_tools.scenario import install_into_venv, get_venv_python
+
+# right now, we are assuming you HAVE to be in the azure-sdk-tools repo
+# we assume this because we don't know how a dev has installed this package, and might be
+# being called from within a site-packages folder. Due to that, we can't trust the location of __file__
+REPO_ROOT = discover_repo_root()
 
 class Check(abc.ABC):
     """
@@ -34,13 +45,41 @@ def run(self, args: argparse.Namespace) -> int:
         """
         return 0
 
+    def create_venv(self, isolate: bool, venv_location: str) -> str:
+        """Abstraction for creating a virtual environment."""
+        if (isolate):
+            venv_cmd = get_venv_call(sys.executable)
+            check_call(venv_cmd + [venv_location])
+
+            # TODO: we should reuse part of build_whl_for_req to integrate with PREBUILT_WHL_DIR so that we don't have to fresh build for each
+            # venv
+            install_into_venv(venv_location, os.path.join(REPO_ROOT, "eng/tools/azure-sdk-tools"), False, "build")
+            venv_python_exe = get_venv_python(venv_location)
+
+            return venv_python_exe
+
+        # if we don't need to isolate, just return the python executable that we're invoking
+        return sys.executable
+
+    def get_executable(self, isolate: bool, check_name: str, executable: str, package_folder: str) -> Tuple[str, str]:
+        """Get the Python executable that should be used for this check."""
+        venv_location = os.path.join(package_folder, f".venv_{check_name}")
+
+        # if isolation is required, the executable we get back will align with the venv
+        # otherwise we'll just get sys.executable and install in current
+        executable = self.create_venv(isolate, venv_location)
+        staging_directory = os.path.join(venv_location, ".staging")
+        os.makedirs(staging_directory, exist_ok=True)
+        return executable, staging_directory
+
+
     def get_targeted_directories(self, args: argparse.Namespace) -> List[ParsedSetup]:
         """
         Get the directories that are targeted for the check.
         """
         targeted: List[ParsedSetup] = []
         targeted_dir = os.getcwd()
- 
+
         if args.target == ".":
             try:
                 targeted.append(ParsedSetup.from_path(targeted_dir))

eng/tools/azure-sdk-tools/azpysdk/import_all.py

@@ -53,18 +53,18 @@ def run(self, args: argparse.Namespace) -> int:
 
         for parsed in targeted:
             pkg = parsed.folder
-                
+            executable, staging_directory = self.get_executable(args.isolate, args.command, sys.executable, pkg)
 
-            staging_area = tempfile.mkdtemp()
             create_package_and_install(
-                distribution_directory=staging_area,
+                distribution_directory=staging_directory,
                 target_setup=pkg,
                 skip_install=False,
                 cache_dir=None,
-                work_dir=staging_area,
+                work_dir=staging_directory,
                 force_create=False,
                 package_type="wheel",
                 pre_download_disabled=False,
+                python_executable=executable
             )
 
             if should_run_import_all(parsed.name):
@@ -76,7 +76,7 @@ def run(self, args: argparse.Namespace) -> int:
                 )
                 import_script_all = "from {0} import *".format(parsed.namespace)
                 commands = [
-                    sys.executable,
+                    executable,
                     "-c",
                     import_script_all
                 ]

eng/tools/azure-sdk-tools/azpysdk/main.py

@@ -11,21 +11,11 @@
 import sys
 import os
 from typing import Sequence, Optional
-from subprocess import check_call
 
 from .whl import whl
 from .import_all import import_all
 from .mypy import mypy
 
-from ci_tools.scenario import install_into_venv, get_venv_python
-from ci_tools.functions import get_venv_call
-from ci_tools.variables import discover_repo_root
-
-# right now, we are assuming you HAVE to be in the azure-sdk-tools repo
-# we assume this because we don't know how a dev has installed this package, and might be
-# being called from within a site-packages folder. Due to that, we can't trust the location of __file__
-REPO_ROOT = discover_repo_root()
-
 __all__ = ["main", "build_parser"]
 __version__ = "0.0.0"
 
@@ -63,35 +53,7 @@ def build_parser() -> argparse.ArgumentParser:
 
     return parser
 
-def handle_venv(isolate: bool, args: argparse.Namespace) -> None:
-    """Handle virtual environment commands."""
-    # we are already in an isolated venv and so do not need to recurse
-    if(os.getenv("AZURE_SDK_TOOLS_VENV", None)):
-        return
-
-    # however, if we are not already in an isolated venv, and should be, then we need to
-    # call
-    if (isolate):
-        os.environ["AZURE_SDK_TOOLS_VENV"] = "1"
-
-        venv_cmd = get_venv_call()
-        venv_location = os.path.join(REPO_ROOT, f".venv_{args.command}")
-        # todo, make this a consistent directory based on the command
-        # I'm seriously thinking we should move handle_venv within each check's main(),
-        # which will mean that we will _know_ what folder we're in.
-        # however, that comes at the cost of not having every check be able to handle one or multiple packages
-        # I don't want to get into an isolation loop where every time we need a new venv, we create it, call it,
-        # and now as we foreach across the targeted packages we've lost our spot.
-        check_call(venv_cmd + [venv_location])
-
-        # now use the current virtual environment to install os.path.join(REPO_ROOT, eng/tools/azure-sdk-tools[build])
-        # into the NEW virtual env
-        install_into_venv(venv_location, os.path.join(REPO_ROOT, "eng/tools/azure-sdk-tools"), False, "build")
-        venv_python_exe = get_venv_python(venv_location)
-        command_args = [venv_python_exe, "-m", "azpysdk.main"] + sys.argv[1:]
-        check_call(command_args)
-
-def main(argv: Optional[Sequence[str]] = None) -> int:#
+def main(argv: Optional[Sequence[str]] = None) -> int:
     """CLI entrypoint.
 
     Args:
@@ -108,9 +70,6 @@ def main(argv: Optional[Sequence[str]] = None) -> int:#
         return 1
 
     try:
-        # scbedd 8/25 I'm betting that this would be best placed within the check itself,
-        # but leaving this for now
-        handle_venv(args.isolate, args)
         result = args.func(args)
         return int(result or 0)
     except KeyboardInterrupt:

eng/tools/azure-sdk-tools/azpysdk/mypy.py

@@ -9,7 +9,7 @@
 
 from .Check import Check
 from ci_tools.parsing import ParsedSetup
-from ci_tools.functions import discover_targeted_packages
+from ci_tools.functions import pip_install
 from ci_tools.scenario.generation import create_package_and_install
 from ci_tools.variables import in_ci, set_envvar_defaults
 from ci_tools.environment_exclusions import (
@@ -33,8 +33,8 @@ def register(self, subparsers: "argparse._SubParsersAction", parent_parsers: Opt
         p.set_defaults(func=self.run)
 
         p.add_argument(
-            "--next", 
-            default=False, 
+            "--next",
+            default=False,
             help="Next version of mypy is being tested",
             required=False
         )
@@ -52,27 +52,27 @@ def run(self, args: argparse.Namespace) -> int:
         for parsed in targeted:
             package_dir = parsed.folder
             package_name = parsed.name
+            executable, staging_directory = self.get_executable(args.isolate, args.command, sys.executable, package_dir)
             print(f"Processing {package_name} for mypy check")
-
-            staging_area = tempfile.mkdtemp()
             create_package_and_install(
-                distribution_directory=staging_area,
+                distribution_directory=staging_directory,
                 target_setup=package_dir,
                 skip_install=False,
                 cache_dir=None,
-                work_dir=staging_area,
+                work_dir=staging_directory,
                 force_create=False,
                 package_type="wheel",
                 pre_download_disabled=False,
+                python_executable=executable
             )
 
             # install mypy
             try:
                 if (args.next):
                     # use latest version of mypy
-                    check_call([sys.executable, "-m", "pip", "install", "mypy"])
+                    pip_install(["mypy"], True, executable, package_dir)
                 else:
-                    check_call([sys.executable, "-m", "pip", "install", f"mypy=={MYPY_VERSION}"])
+                    pip_install([f"mypy=={MYPY_VERSION}"], True, executable, package_dir)
             except CalledProcessError as e:
                 print("Failed to install mypy:", e)
                 return e.returncode
@@ -89,7 +89,7 @@ def run(self, args: argparse.Namespace) -> int:
             top_level_module = parsed.namespace.split(".")[0]
 
             commands = [
-                sys.executable,
+                executable,
                 "-m",
                 "mypy",
                 "--python-version",
@@ -107,9 +107,9 @@ def run(self, args: argparse.Namespace) -> int:
                 results.append(check_call(src_code))
                 logging.info("Verified mypy, no issues found")
             except CalledProcessError as src_error:
-                src_code_error = src_error 
+                src_code_error = src_error
                 results.append(src_error.returncode)
-            
+
             if not args.next and in_ci() and not is_check_enabled(package_dir, "type_check_samples", True):
                 logging.info(
                     f"Package {package_name} opts-out of mypy check on samples."
@@ -145,6 +145,5 @@ def run(self, args: argparse.Namespace) -> int:
                     create_vnext_issue(package_dir, "mypy")
                 else:
                     close_vnext_issue(package_name, "mypy")
-            
+
         return max(results) if results else 0
-        

eng/tools/azure-sdk-tools/azpysdk/whl.py

@@ -3,12 +3,12 @@
 import tempfile
 import os
 from typing import Optional, List, Any
-from pytest import main as pytest_main
 import sys
+from subprocess import run
 
 from .Check import Check
 
-from ci_tools.functions import discover_targeted_packages, is_error_code_5_allowed, pip_install
+from ci_tools.functions import is_error_code_5_allowed, pip_install
 from ci_tools.variables import set_envvar_defaults
 from ci_tools.parsing import ParsedSetup
 from ci_tools.scenario.generation import create_package_and_install
@@ -25,7 +25,7 @@ def register(self, subparsers: "argparse._SubParsersAction", parent_parsers: Opt
         parents = parent_parsers or []
         p = subparsers.add_parser("whl", parents=parents, help="Run the whl check")
         p.set_defaults(func=self.run)
-        # Add any additional arguments specific to WhlCheck here (do not re-add common handled by parents)
+        # TODO add mark_args, and other parameters
 
     def run(self, args: argparse.Namespace) -> int:
         """Run the whl check command."""
@@ -39,30 +39,47 @@ def run(self, args: argparse.Namespace) -> int:
 
         for parsed in targeted:
             pkg = parsed.folder
+            executable, staging_directory = self.get_executable(args.isolate, args.command, sys.executable, pkg)
 
+            print(f"Invoking check with {executable}")
             dev_requirements = os.path.join(pkg, "dev_requirements.txt")
 
             if os.path.exists(dev_requirements):
-                pip_install([f"-r", f"{dev_requirements}"], True, sys.executable)
+                print(f"Installing dev_requirements at {dev_requirements}")
+                pip_install([f"-r", f"{dev_requirements}"], True, executable, pkg)
+            else:
+                print("Skipping installing dev_requirements")
 
-            staging_area = tempfile.mkdtemp()
             create_package_and_install(
-                distribution_directory=staging_area,
+                distribution_directory=staging_directory,
                 target_setup=pkg,
                 skip_install=False,
                 cache_dir=None,
-                work_dir=staging_area,
+                work_dir=staging_directory,
                 force_create=False,
                 package_type="wheel",
                 pre_download_disabled=False,
+                python_executable=executable
             )
 
-            # todo, come up with a good pattern for passing all the additional args after -- to pytest
+            # TODO: split sys.argv[1:] on -- and pass in everything after the -- as additional arguments
+            # TODO: handle mark_args
             logging.info(f"Invoke pytest for {pkg}")
-
-            exit_code = pytest_main(
-                [pkg]
-            )
+            exit_code = run(
+                [executable, "-m", "pytest", "."] + [
+                    "-rsfE",
+                    f"--junitxml={pkg}/test-junit-{args.command}.xml",
+                    "--verbose",
+                    "--cov-branch",
+                    "--durations=10",
+                    "--ignore=azure",
+                    "--ignore-glob=.venv*",
+                    "--ignore=build",
+                    "--ignore=.eggs",
+                    "--ignore=samples"
+                ]
+                , cwd=pkg
+            ).returncode
 
             if exit_code != 0:
                 if exit_code == 5 and is_error_code_5_allowed(parsed.folder, parsed.name):

eng/tools/azure-sdk-tools/ci_tools/functions.py

@@ -438,7 +438,7 @@ def find_sdist(dist_dir: str, pkg_name: str, pkg_version: str) -> Optional[str]:
 
 
 def pip_install(
-    requirements: List[str], include_dependencies: bool = True, python_executable: Optional[str] = None
+    requirements: List[str], include_dependencies: bool = True, python_executable: Optional[str] = None, cwd: Optional[str] = None
 ) -> bool:
     """
     Attempts to invoke an install operation using the invoking python's pip. Empty requirements are auto-success.
@@ -454,7 +454,10 @@ def pip_install(
         return True
 
     try:
-        subprocess.check_call(command)
+        if cwd:
+            subprocess.check_call(command, cwd=cwd)
+        else:
+            subprocess.check_call(command)
     except subprocess.CalledProcessError as f:
         return False
 
@@ -488,8 +491,10 @@ def get_pip_list_output(python_executable: Optional[str] = None):
     """Uses the invoking python executable to get the output from pip list."""
     exe = python_executable or sys.executable
 
+    pip_cmd = get_pip_command(exe)
+
     out = subprocess.Popen(
-        [exe, "-m", "pip", "list", "--disable-pip-version-check", "--format", "freeze"],
+        pip_cmd + ["list", "--disable-pip-version-check", "--format", "freeze"],
         stdout=subprocess.PIPE,
         stderr=subprocess.STDOUT,
     )

eng/tools/azure-sdk-tools/ci_tools/scenario/generation.py

@@ -316,7 +316,7 @@ def build_whl_for_req(req: str, package_path: str, wheel_dir: Optional[str]) ->
     """Builds a whl from the dev_requirements file.
 
     :param str req: a requirement from the dev_requirements.txt
-    :param str package_path: the absolute path to the package's root
+    :param str package_path: the absolute path to the package's root (used as relative path root)
     :param Optional[str] wheel_dir: the absolute path to the prebuilt wheel directory
     :return: The absolute path to the whl built or the requirement if a third-party package
     """