warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/thirdparty/unittest_parallel.py

@@ -1,11 +1,22 @@
 # Licensed under the MIT License
 # https://github.com/craigahobbs/unittest-parallel/blob/main/LICENSE
 
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
 """
 unittest-parallel command-line script main module
 """
 
 import argparse
+import concurrent.futures  # NVIDIA Modification
 import multiprocessing
 import os
 import sys
@@ -15,7 +26,20 @@ import unittest
 from contextlib import contextmanager
 from io import StringIO
 
-import coverage
+import warp.tests.unittest_suites  # NVIDIA Modification
+
+try:
+    import coverage
+
+    COVERAGE_AVAILABLE = True  # NVIDIA Modification
+except ImportError:
+    COVERAGE_AVAILABLE = False  # NVIDIA Modification
+
+
+# The following variables are NVIDIA Modifications
+RUNNING_IN_TEAMCITY = os.environ.get("TEAMCITY_VERSION") is not None
+TEST_SUITE_NAME = "WarpTests"
+START_DIRECTORY = os.path.dirname(__file__)  # The directory to start test discovery
 
 
 def main(argv=None):
@@ -25,8 +49,8 @@ def main(argv=None):
 
     # Command line arguments
     parser = argparse.ArgumentParser(prog="unittest-parallel")
-    parser.add_argument("-v", "--verbose", action="store_const", const=2, default=1, help="Verbose output")
-    parser.add_argument("-q", "--quiet", dest="verbose", action="store_const", const=0, default=1, help="Quiet output")
+    # parser.add_argument("-v", "--verbose", action="store_const", const=2, default=1, help="Verbose output")
+    parser.add_argument("-q", "--quiet", dest="verbose", action="store_const", const=0, default=2, help="Quiet output")
     parser.add_argument("-f", "--failfast", action="store_true", default=False, help="Stop on first fail or error")
     parser.add_argument(
         "-b", "--buffer", action="store_true", default=False, help="Buffer stdout and stderr during tests"
@@ -38,9 +62,6 @@ def main(argv=None):
         type=_convert_select_pattern,
         help="Only run tests which match the given substring",
     )
-    parser.add_argument(
-        "-s", "--start-directory", metavar="START", default=".", help="Directory to start discovery ('.' default)"
-    )
     parser.add_argument(
         "-p", "--pattern", metavar="PATTERN", default="test*.py", help="Pattern to match tests ('test*.py' default)"
     )
@@ -59,11 +80,19 @@ def main(argv=None):
         default=0,
         help="The number of test processes (default is 0, all cores)",
     )
+    group_parallel.add_argument(
+        "-m",
+        "--maxjobs",
+        metavar="MAXCOUNT",
+        type=int,
+        default=8,
+        help="The maximum number of test processes (default is 8)",
+    )  # NVIDIA Modification
     group_parallel.add_argument(
         "--level",
         choices=["module", "class", "test"],
-        default="module",
-        help="Set the test parallelism level (default is 'module')",
+        default="class",
+        help="Set the test parallelism level (default is 'class')",
     )
     group_parallel.add_argument(
         "--disable-process-pooling",
@@ -71,40 +100,45 @@ def main(argv=None):
         default=False,
         help="Do not reuse processes used to run test suites",
     )
+    group_parallel.add_argument(
+        "--disable-concurrent-futures",
+        action="store_true",
+        default=False,
+        help="Use multiprocessing instead of concurrent.futures.",
+    )  # NVIDIA Modification
+    group_parallel.add_argument(
+        "--serial-fallback",
+        action="store_true",
+        default=False,
+        help="Run in a single-process (no spawning) mode without multiprocessing or concurrent.futures.",
+    )  # NVIDIA Modification
     group_coverage = parser.add_argument_group("coverage options")
     group_coverage.add_argument("--coverage", action="store_true", help="Run tests with coverage")
     group_coverage.add_argument("--coverage-branch", action="store_true", help="Run tests with branch coverage")
-    group_coverage.add_argument("--coverage-rcfile", metavar="RCFILE", help="Specify coverage configuration file")
-    group_coverage.add_argument(
-        "--coverage-include",
-        metavar="PAT",
-        action="append",
-        help="Include only files matching one of these patterns. Accepts shell-style (quoted) wildcards.",
-    )
-    group_coverage.add_argument(
-        "--coverage-omit",
-        metavar="PAT",
-        action="append",
-        help="Omit files matching one of these patterns. Accepts shell-style (quoted) wildcards.",
-    )
     group_coverage.add_argument(
-        "--coverage-source",
-        metavar="SRC",
-        action="append",
-        help="A list of packages or directories of code to be measured",
+        "--coverage-html",
+        metavar="DIR",
+        help="Generate coverage HTML report",
+        default=os.path.join(START_DIRECTORY, "..", "..", "htmlcov"),
     )
-    group_coverage.add_argument("--coverage-html", metavar="DIR", help="Generate coverage HTML report")
     group_coverage.add_argument("--coverage-xml", metavar="FILE", help="Generate coverage XML report")
     group_coverage.add_argument(
         "--coverage-fail-under", metavar="MIN", type=float, help="Fail if coverage percentage under min"
     )
     args = parser.parse_args(args=argv)
+
     if args.coverage_branch:
         args.coverage = args.coverage_branch
 
+    if args.coverage and not COVERAGE_AVAILABLE:
+        parser.exit(
+            status=2, message="--coverage was used, but coverage was not found. Is it installed?\n"
+        )  # NVIDIA Modification
+
     process_count = max(0, args.jobs)
     if process_count == 0:
         process_count = multiprocessing.cpu_count()
+    process_count = min(process_count, args.maxjobs)  # NVIDIA Modification
 
     # Create the temporary directory (for coverage files)
     with tempfile.TemporaryDirectory() as temp_dir:
@@ -113,9 +147,10 @@ def main(argv=None):
             test_loader = unittest.TestLoader()
             if args.testNamePatterns:
                 test_loader.testNamePatterns = args.testNamePatterns
-            discover_suite = test_loader.discover(
-                args.start_directory, pattern=args.pattern, top_level_dir=args.top_level_directory
-            )
+            discover_suite = warp.tests.unittest_suites.auto_discover_suite(
+                test_loader, args.pattern
+            )  # NVIDIA Modification
+            # discover_suite = warp.tests.unittest_suites.explicit_suite()
 
         # Get the parallelizable test suites
         if args.level == "test":
@@ -128,26 +163,69 @@ def main(argv=None):
         # Don't use more processes than test suites
         process_count = max(1, min(len(test_suites), process_count))
 
-        # Report test suites and processes
-        print(
-            f"Running {len(test_suites)} test suites ({discover_suite.countTestCases()} total tests) across {process_count} processes",
-            file=sys.stderr,
-        )
-        if args.verbose > 1:
-            print(file=sys.stderr)
+        if RUNNING_IN_TEAMCITY:
+            print(f"##teamcity[testSuiteStarted name='{TEST_SUITE_NAME}']")  # NVIDIA Modification for TC
+
+        if not args.serial_fallback:
+            # Report test suites and processes
+            print(
+                f"Running {len(test_suites)} test suites ({discover_suite.countTestCases()} total tests) across {process_count} processes",
+                file=sys.stderr,
+            )
+            if args.verbose > 1:
+                print(file=sys.stderr)
+
+            # Run the tests in parallel
+            start_time = time.perf_counter()
+
+            if args.disable_concurrent_futures:
+                multiprocessing_context = multiprocessing.get_context(method="spawn")
+                maxtasksperchild = 1 if args.disable_process_pooling else None
+                with multiprocessing_context.Pool(
+                    process_count,
+                    maxtasksperchild=maxtasksperchild,
+                    initializer=set_worker_cache,
+                    initargs=(args, temp_dir),
+                ) as pool, multiprocessing.Manager() as manager:
+                    test_manager = ParallelTestManager(manager, args, temp_dir)
+                    results = pool.map(test_manager.run_tests, test_suites)
+            else:
+                # NVIDIA Modification added concurrent.futures
+                with concurrent.futures.ProcessPoolExecutor(
+                    max_workers=process_count,
+                    mp_context=multiprocessing.get_context(method="spawn"),
+                    initializer=set_worker_cache,
+                    initargs=(args, temp_dir),
+                ) as executor, multiprocessing.Manager() as manager:
+                    test_manager = ParallelTestManager(manager, args, temp_dir)
+                    results = list(executor.map(test_manager.run_tests, test_suites, timeout=3600))
+        else:
+            # This entire path is an NVIDIA Modification
+
+            # Report test suites and processes
+            print(f"Running {discover_suite.countTestCases()} total tests (serial fallback)", file=sys.stderr)
+            if args.verbose > 1:
+                print(file=sys.stderr)
+
+            import warp as wp
+
+            # force rebuild of all kernels
+            wp.build.clear_kernel_cache()
+            print("Cleared Warp kernel cache")
+
+            # Run the tests in serial
+            start_time = time.perf_counter()
+
+            with multiprocessing.Manager() as manager:
+                test_manager = ParallelTestManager(manager, args, temp_dir)
+                results = [test_manager.run_tests(discover_suite)]
 
-        # Run the tests in parallel
-        start_time = time.perf_counter()
-        multiprocessing_context = multiprocessing.get_context(method="spawn")
-        maxtasksperchild = 1 if args.disable_process_pooling else None
-        with multiprocessing_context.Pool(
-            process_count, maxtasksperchild=maxtasksperchild
-        ) as pool, multiprocessing.Manager() as manager:
-            test_manager = ParallelTestManager(manager, args, temp_dir)
-            results = pool.map(test_manager.run_tests, test_suites)
         stop_time = time.perf_counter()
         test_duration = stop_time - start_time
 
+        if RUNNING_IN_TEAMCITY:
+            print(f"##teamcity[testSuiteFinished name='{TEST_SUITE_NAME}']")  # NVIDIA Modification for TC
+
         # Aggregate parallel test run results
         tests_run = 0
         errors = []
@@ -195,14 +273,15 @@ def main(argv=None):
 
         # Return an error status on failure
         if not is_success:
+            if RUNNING_IN_TEAMCITY:
+                print("##teamcity[buildStatus status='FAILURE']")  # NVIDIA Modification for TC
             parser.exit(status=len(errors) + len(failures) + unexpected_successes)
 
         # Coverage?
         if args.coverage:
             # Combine the coverage files
             cov_options = {}
-            if args.coverage_rcfile is not None:
-                cov_options["config_file"] = args.coverage_rcfile
+            cov_options["config_file"] = True  # Grab configuration from pyproject.toml (must install coverage[toml])
             cov = coverage.Coverage(**cov_options)
             cov.combine(data_paths=[os.path.join(temp_dir, x) for x in os.listdir(temp_dir)])
 
@@ -225,7 +304,7 @@ def main(argv=None):
 
 
 def _convert_select_pattern(pattern):
-    if not "*" in pattern:
+    if "*" not in pattern:
         return f"*{pattern}*"
     return pattern
 
@@ -242,12 +321,9 @@ def _coverage(args, temp_dir):
         cov_options = {
             "branch": args.coverage_branch,
             "data_file": coverage_file.name,
-            "include": args.coverage_include,
-            "omit": (args.coverage_omit if args.coverage_omit else []) + [__file__],
-            "source": args.coverage_source,
+            # NVIDIA Modification removed unneeded options
         }
-        if args.coverage_rcfile is not None:
-            cov_options["config_file"] = args.coverage_rcfile
+        cov_options["config_file"] = True  # Grab configuration from pyproject.toml (must install coverage[toml])
         cov = coverage.Coverage(**cov_options)
         try:
             # Start measuring code coverage
@@ -307,7 +383,9 @@ class ParallelTestManager:
         with _coverage(self.args, self.temp_dir):
             runner = unittest.TextTestRunner(
                 stream=StringIO(),
-                resultclass=ParallelTextTestResult,
+                resultclass=ParallelTeamCityTestResult
+                if RUNNING_IN_TEAMCITY
+                else ParallelTextTestResult,  # NVIDIA Modification for TC
                 verbosity=self.args.verbose,
                 failfast=self.args.failfast,
                 buffer=self.args.buffer,
@@ -384,3 +462,128 @@ class ParallelTextTestResult(unittest.TextTestResult):
 
     def printErrors(self):
         pass
+
+
+# NVIDIA Modifications from here until end of file
+
+
+def set_worker_cache(args, temp_dir):
+    """This function is run at the start of ever new process. It changes the Warp cache to avoid conflicts."""
+
+    with _coverage(args, temp_dir):
+        import warp as wp
+        from warp.thirdparty import appdirs
+
+        pid = os.getpid()
+        cache_root_dir = appdirs.user_cache_dir(
+            appname="warp", appauthor="NVIDIA Corporation", version=f"{wp.config.version}-{pid}"
+        )
+
+        wp.config.kernel_cache_dir = cache_root_dir
+
+        wp.build.clear_kernel_cache()
+
+
+def _tc_escape(s):
+    s = s.replace("|", "||")
+    s = s.replace("\n", "|n")
+    s = s.replace("\r", "|r")
+    s = s.replace("'", "|'")
+    s = s.replace("[", "|[")
+    s = s.replace("]", "|]")
+    return s
+
+
+class ParallelTeamCityTestResult(unittest.TextTestResult):
+    def __init__(self, stream, descriptions, verbosity):
+        stream = type(stream)(sys.stderr)
+        super().__init__(stream, descriptions, verbosity)
+
+    def startTest(self, test):
+        if self.showAll:
+            self.stream.writeln(f"{self.getDescription(test)} ...")
+            self.stream.flush()
+        self.start_time = time.perf_counter_ns()
+        super(unittest.TextTestResult, self).startTest(test)
+
+    def _add_helper(self, test, dots_message, show_all_message):
+        if self.showAll:
+            self.stream.writeln(f"{self.getDescription(test)} ... {show_all_message}")
+        elif self.dots:
+            self.stream.write(dots_message)
+        self.stream.flush()
+
+    def addSuccess(self, test):
+        super(unittest.TextTestResult, self).addSuccess(test)
+        self._add_helper(test, ".", "ok")
+        self.reportSuccess(test)
+
+    def addError(self, test, err):
+        super(unittest.TextTestResult, self).addError(test, err)
+        self._add_helper(test, "E", "ERROR")
+        self.reportFailure(test, err)
+
+    def addFailure(self, test, err):
+        super(unittest.TextTestResult, self).addFailure(test, err)
+        self._add_helper(test, "F", "FAIL")
+        self.reportFailure(test, err)
+
+    def addSkip(self, test, reason):
+        super(unittest.TextTestResult, self).addSkip(test, reason)
+        self._add_helper(test, "s", f"skipped {reason!r}")
+        self.reportIgnored(test, reason)
+
+    def addExpectedFailure(self, test, err):
+        super(unittest.TextTestResult, self).addExpectedFailure(test, err)
+        self._add_helper(test, "x", "expected failure")
+        self.reportSuccess(test)
+
+    def addUnexpectedSuccess(self, test):
+        super(unittest.TextTestResult, self).addUnexpectedSuccess(test)
+        self._add_helper(test, "u", "unexpected success")
+        self.reportFailure(test, "unexpected success")
+
+    def addSubTest(self, test, subtest, err):
+        super(unittest.TextTestResult, self).addSubTest(test, subtest, err)
+        if err is not None:
+            self._add_helper(test, "E", "ERROR")
+            self.reportSubTestFailure(test, err)
+
+    def printErrors(self):
+        pass
+
+    def reportIgnored(self, test, reason):
+        test_id = test.id()
+        reason_str = str(reason)
+        print(reason_str)
+        self.stream.writeln(f"##teamcity[testIgnored name='{test_id}' message='{_tc_escape(str(reason))}']")
+        self.stream.flush()
+
+    def reportSuccess(self, test):
+        duration = round((time.perf_counter_ns() - self.start_time) / 1e6)  # [ms]
+        test_id = test.id()
+        self.stream.writeln(f"##teamcity[testStarted name='{test_id}']")
+        self.stream.writeln(f"##teamcity[testFinished name='{test_id}' duration='{duration}']")
+        self.stream.flush()
+
+    def reportFailure(self, test, err):
+        test_id = test.id()
+        self.stream.writeln(f"##teamcity[testStarted name='{test_id}']")
+        self.stream.writeln(
+            f"##teamcity[testFailed name='{test_id}' message='{_tc_escape(str(err[1]))}' details='{_tc_escape(str(err[2]))}']"
+        )
+        self.stream.writeln(f"##teamcity[testFinished name='{test_id}']")
+        self.stream.flush()
+
+    def reportSubTestFailure(self, test, err):
+        test_id = test.id()
+        self.stream.writeln(f"##teamcity[testStarted name='{test_id}']")
+        self.stream.writeln(
+            f"##teamcity[testFailed name='{test_id}' message='{_tc_escape(str(err[1]))}' details='{_tc_escape(self._exc_info_to_string(err, test))}']"
+        )
+        self.stream.writeln(f"##teamcity[testFinished name='{test_id}']")
+        self.stream.flush()
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()