Stream dispatch_checks output when running a single check (#47214)

* test

* test2

* add header in stream output for vnext checks

* minor typing

* more specific exception and logging

* no try except

Author: JennyPng

eng/scripts/dispatch_checks.py

@@ -77,6 +77,48 @@ def _normalize_newlines(text: str) -> str:
     return text.replace("\r\n", "\n").replace("\r", "\n")
 
 
+async def _tee_stream(
+    proc: "asyncio.subprocess.Process", package: str, check: str
+) -> tuple:
+    """Read the child's stdout and stderr concurrently, mirroring each line to
+    this process's stdout/stderr while also accumulating the full text for the
+    final :class:`CheckResult`.
+
+    This exists because the default ``proc.communicate()`` path buffers all
+    output until the child exits. When a check hangs and the pipeline agent
+    cancels the parent at the job timeout, none of that buffered output is
+    ever printed -- the log shows ``CMD: ...`` followed by silence until the
+    cancellation marker. Streaming line-by-line ensures the last bytes the
+    child produced before getting stuck are visible in the pipeline log.
+
+    Should only be used when checks are running sequentially (``max_parallel
+    == 1``) -- otherwise output from concurrent children will interleave.
+    """
+    prefix = f"[{os.path.basename(os.path.normpath(package))} :: {check}] "
+
+    async def _pump(stream: Optional[asyncio.StreamReader], sink: IO[str]) -> str:
+        if stream is None:
+            return ""
+        chunks: List[str] = []
+        while True:
+            line_b = await stream.readline()
+            if not line_b:
+                break
+            line = line_b.decode(errors="replace")
+            chunks.append(line)
+            sink.write(prefix + line)
+            sink.flush()
+        return "".join(chunks)
+
+    stdout_text, stderr_text = await asyncio.gather(
+        _pump(proc.stdout, sys.stdout),
+        _pump(proc.stderr, sys.stderr),
+    )
+    # Make sure the process is reaped so returncode is populated.
+    await proc.wait()
+    return stdout_text, stderr_text
+
+
 def _checks_require_recording_restore(checks: List[str]) -> bool:
     return any(check in INSTALL_AND_TEST_CHECKS for check in checks)
 
@@ -87,7 +129,9 @@ def _compare_req_to_injected_reqs(parsed_req, injected_packages: List[str]) -> b
     return any(parsed_req.name in req for req in injected_packages)
 
 
-def _inject_custom_reqs(req_file: str, injected_packages: str, package_dir: str) -> None:
+def _inject_custom_reqs(
+    req_file: str, injected_packages: str, package_dir: str
+) -> None:
     req_lines = []
     injected_list = [p for p in re.split(r"[\s,]", injected_packages) if p]
 
@@ -116,7 +160,8 @@ def _inject_custom_reqs(req_file: str, injected_packages: str, package_dir: str)
         all_adjustments = installable + [
             line_tuple[0].strip()
             for line_tuple in req_lines
-            if line_tuple[0].strip() and not _compare_req_to_injected_reqs(line_tuple[1], all_filter_names)
+            if line_tuple[0].strip()
+            and not _compare_req_to_injected_reqs(line_tuple[1], all_filter_names)
         ]
     else:
         all_adjustments = installable
@@ -139,6 +184,7 @@ async def run_check(
     dest_dir: Optional[str] = None,
     service: Optional[str] = None,
     python_version: Optional[str] = None,
+    stream_live: bool = False,
 ) -> CheckResult:
     """Run a single check (subprocess) within a concurrency semaphore, capturing output and timing.
 
@@ -156,6 +202,12 @@ async def run_check(
     :type total: int
     :param proxy_port: Dedicated proxy port assigned to this check instance.
     :type proxy_port: int
+    :param stream_live: If True, tee the child's stdout/stderr to this process's
+        stdout/stderr line-by-line while also capturing them. Use only when
+        checks run sequentially (``max_parallel == 1``) to avoid interleaved
+        output. This prevents silent hangs from hiding all diagnostics when the
+        agent kills the parent before the child exits.
+    :type stream_live: bool
     :returns: A :class:`CheckResult` describing exit code, duration and captured output.
     :rtype: CheckResult
     """
@@ -173,9 +225,16 @@ async def run_check(
         logger.info(f"[START {idx}/{total}] {check} :: {package}\nCMD: {' '.join(cmd)}")
         env = os.environ.copy()
         env["PROXY_URL"] = f"http://localhost:{proxy_port}"
+        # Force the child Python process to use unbuffered stdio. Without this,
+        # the child's output can sit in its own internal buffers for minutes
+        # (or never appear at all if the agent kills us at the job timeout),
+        # making hung checks look like total silence in the pipeline log.
+        env["PYTHONUNBUFFERED"] = "1"
 
         if in_ci():
-            env["PROXY_ASSETS_FOLDER"] = os.path.join(root_dir, ".assets_distributed", str(proxy_port))
+            env["PROXY_ASSETS_FOLDER"] = os.path.join(
+                root_dir, ".assets_distributed", str(proxy_port)
+            )
         try:
             logger.info(" ".join(cmd))
             proc = await asyncio.create_subprocess_exec(
@@ -189,30 +248,43 @@ async def run_check(
             logger.error(f"Failed to start check {check} for {package}: {ex}")
             return CheckResult(package, check, 127, 0.0, "", str(ex))
 
-        stdout_b, stderr_b = await proc.communicate()
+        if stream_live:
+            if in_ci():
+                print(f"##[group]{package} :: {check} :: ?")
+            stdout, stderr = await _tee_stream(proc, package, check)
+            if in_ci():
+                print("##[endgroup]")
+        else:
+            stdout_b, stderr_b = await proc.communicate()
+            stdout = stdout_b.decode(errors="replace")
+            stderr = stderr_b.decode(errors="replace")
         duration = time.time() - start
-        stdout = stdout_b.decode(errors="replace")
-        stderr = stderr_b.decode(errors="replace")
         exit_code = proc.returncode or 0
         status = "OK" if exit_code == 0 else f"FAIL({exit_code})"
-        logger.info(f"[END   {idx}/{total}] {check} :: {package} -> {status} in {duration:.2f}s")
-        # Print captured output after completion to avoid interleaving
-        header = f"===== OUTPUT: {check} :: {package} (exit {exit_code}) ====="
-        trailer = "=" * len(header)
-        if in_ci():
-            print(f"##[group]{package} :: {check} :: {exit_code}")
-
-        if stdout:
-            print(header)
-            print(_normalize_newlines(stdout).rstrip())
-            print(trailer)
-        if stderr:
-            print(header.replace("OUTPUT", "STDERR"))
-            print(_normalize_newlines(stderr).rstrip())
-            print(trailer)
-
-        if in_ci():
-            print("##[endgroup]")
+        logger.info(
+            f"[END   {idx}/{total}] {check} :: {package} -> {status} in {duration:.2f}s"
+        )
+        # When streaming live we've already mirrored every line to the parent's
+        # stdout/stderr as it was produced, so skip the post-hoc grouped dump
+        # to avoid duplicating every line.
+        if not stream_live:
+            # Print captured output after completion to avoid interleaving
+            header = f"===== OUTPUT: {check} :: {package} (exit {exit_code}) ====="
+            trailer = "=" * len(header)
+            if in_ci():
+                print(f"##[group]{package} :: {check} :: {exit_code}")
+
+            if stdout:
+                print(header)
+                print(_normalize_newlines(stdout).rstrip())
+                print(trailer)
+            if stderr:
+                print(header.replace("OUTPUT", "STDERR"))
+                print(_normalize_newlines(stderr).rstrip())
+                print(trailer)
+
+            if in_ci():
+                print("##[endgroup]")
 
         # if we have any output collections to complete, do so now here
 
@@ -246,10 +318,14 @@ def summarize(results: List[CheckResult]) -> int:
     print("-" * len(header))
     for r in sorted(results, key=lambda x: (x.exit_code != 0, x.package, x.check)):
         status = "OK" if r.exit_code == 0 else f"FAIL({r.exit_code})"
-        print(f"{r.package.ljust(pkg_w)}  {r.check.ljust(chk_w)}  {status.ljust(8)}  {r.duration:>10.2f}")
+        print(
+            f"{r.package.ljust(pkg_w)}  {r.check.ljust(chk_w)}  {status.ljust(8)}  {r.duration:>10.2f}"
+        )
     worst = max((r.exit_code for r in results), default=0)
     failed = [r for r in results if r.exit_code != 0]
-    print(f"\nTotal checks: {len(results)} | Failed: {len(failed)} | Worst exit code: {worst}")
+    print(
+        f"\nTotal checks: {len(results)} | Failed: {len(failed)} | Worst exit code: {worst}"
+    )
     return worst
 
 
@@ -287,10 +363,14 @@ async def run_all_checks(
     dependency_tools_path = os.path.join(root_dir, "eng", "dependency_tools.txt")
 
     if in_ci():
-        logger.info("Replacing relative requirements in eng/test_tools.txt with prebuilt wheels.")
+        logger.info(
+            "Replacing relative requirements in eng/test_tools.txt with prebuilt wheels."
+        )
         replace_dev_reqs(test_tools_path, root_dir, wheel_dir)
 
-        logger.info("Replacing relative requirements in eng/dependency_tools.txt with prebuilt wheels.")
+        logger.info(
+            "Replacing relative requirements in eng/dependency_tools.txt with prebuilt wheels."
+        )
         replace_dev_reqs(dependency_tools_path, root_dir, wheel_dir)
 
     for pkg in packages:
@@ -312,19 +392,34 @@ async def run_all_checks(
         if not is_check_enabled(package, check, CHECK_DEFAULTS.get(check, True)):
             logger.warning(f"Skipping disabled check {check} for package {package}")
             continue
-        logger.info(f"Assigning proxy port {next_proxy_port} to check {check} for package {package}")
+        logger.info(
+            f"Assigning proxy port {next_proxy_port} to check {check} for package {package}"
+        )
 
         # Check if this package overrides the Python version for analysis
         pkg_python_version = get_config_setting(package, "analyze_python_version", None)
         if pkg_python_version:
-            logger.info(f"Package {package} overrides analyze Python version to {pkg_python_version}")
+            logger.info(
+                f"Package {package} overrides analyze Python version to {pkg_python_version}"
+            )
 
         scheduled.append((package, check, next_proxy_port, pkg_python_version))
         next_proxy_port += 1
 
     total = len(scheduled)
 
-    for idx, (package, check, proxy_port, pkg_python_version) in enumerate(scheduled, start=1):
+    # Mirror the child's stdio live to the parent's stdout/stderr when no
+    # concurrent children could interleave output. This makes hangs visible
+    # in real time instead of being hidden behind a post-hoc grouped dump
+    # that never prints if the agent cancels us at the job timeout. We check
+    # both the semaphore cap AND the actual number of tasks because cosmos
+    # (and similar live test legs) ship a single check per matrix leg but
+    # leave --max-parallel at its CPU-count default.
+    stream_live = max_parallel == 1 or total <= 1
+
+    for idx, (package, check, proxy_port, pkg_python_version) in enumerate(
+        scheduled, start=1
+    ):
         tasks.append(
             asyncio.create_task(
                 run_check(
@@ -339,6 +434,7 @@ async def run_all_checks(
                     dest_dir,
                     service,
                     pkg_python_version,
+                    stream_live=stream_live,
                 )
             )
         )
@@ -360,7 +456,9 @@ async def run_all_checks(
         elif isinstance(res, Exception):
             norm_results.append(CheckResult(package, check, 99, 0.0, "", str(res)))
         else:
-            norm_results.append(CheckResult(package, check, 98, 0.0, "", f"Unknown result type: {res}"))
+            norm_results.append(
+                CheckResult(package, check, 98, 0.0, "", f"Unknown result type: {res}")
+            )
     return summarize(norm_results)
 
 
@@ -436,11 +534,15 @@ def handler(signum, frame):
         ),
     )
 
-    parser.add_argument("--disablecov", help=("Flag. Disables code coverage."), action="store_true")
+    parser.add_argument(
+        "--disablecov", help=("Flag. Disables code coverage."), action="store_true"
+    )
 
     parser.add_argument(
         "--service",
-        help=("Name of service directory (under sdk/) to test. Example: --service applicationinsights"),
+        help=(
+            "Name of service directory (under sdk/) to test. Example: --service applicationinsights"
+        ),
     )
 
     parser.add_argument(
@@ -507,7 +609,9 @@ def handler(signum, frame):
     else:
         target_dir = root_dir
 
-    logger.info(f"Beginning discovery for {args.service} and root dir {root_dir}. Resolving to {target_dir}.")
+    logger.info(
+        f"Beginning discovery for {args.service} and root dir {root_dir}. Resolving to {target_dir}."
+    )
 
     # ensure that recursive virtual envs aren't messed with by this call
     os.environ.pop("VIRTUAL_ENV", None)
@@ -524,7 +628,9 @@ def handler(signum, frame):
     )
 
     if len(targeted_packages) == 0:
-        logger.info(f"No packages collected for targeting string {args.glob_string} and root dir {root_dir}. Exit 0.")
+        logger.info(
+            f"No packages collected for targeting string {args.glob_string} and root dir {root_dir}. Exit 0."
+        )
         exit(0)
 
     logger.info(f"Executing checks with the executable {sys.executable}.")
@@ -556,7 +662,9 @@ def handler(signum, frame):
         try:
             proxy_executable = prepare_local_tool(root_dir)
         except Exception as exc:
-            logger.error(f"Unable to prepare test proxy executable for recording restore: {exc}")
+            logger.error(
+                f"Unable to prepare test proxy executable for recording restore: {exc}"
+            )
             sys.exit(1)
 
     logger.info(
@@ -567,9 +675,13 @@ def handler(signum, frame):
     proxy_processes: List[ProxyProcess] = []
     try:
         if in_ci():
-            logger.info(f"Ensuring {len(checks)} test proxies are running for requested checks...")
+            logger.info(
+                f"Ensuring {len(checks)} test proxies are running for requested checks..."
+            )
         # Pass through service if set and not "auto"
-        effective_service = args.service if (args.service and args.service != "auto") else None
+        effective_service = (
+            args.service if (args.service and args.service != "auto") else None
+        )
         exit_code = asyncio.run(
             run_all_checks(
                 targeted_packages,