Add a flag to enable detailed timestamped logging of subprocess commands.

This adds a new command-line flag, `--detailed_timestamped_log`, that enables detailed logging of Bazel build commands. When disabled (the default), logging mirrors the output you'd see when running the command directly in your terminal.

When this flag is enabled:
- Bazel's output is captured line by line.
- Each line is timestamped for improved traceability.
- The complete log is stored for potential use as an artifact.

The flag is disabled by default and only enabled in the CI builds. If you're running locally and enable `detailed_timestamped_log`, you might notice that Bazel's output is not colored. To force a color output, include `--bazel_options=--color=yes` in your command.

PiperOrigin-RevId: 703581368

Author: nitins17

build/build.py

@@ -123,6 +123,15 @@ def add_global_arguments(parser: argparse.ArgumentParser):
       help="Produce verbose output for debugging.",
   )
 
+  parser.add_argument(
+      "--detailed_timestamped_log",
+      action="store_true",
+      help="""
+        Enable detailed logging of the Bazel command with timestamps. The logs
+        will be stored and can be accessed as artifacts.
+        """,
+  )
+
 
 def add_artifact_subcommand_arguments(parser: argparse.ArgumentParser):
   """Adds all the arguments that applies to the artifact subcommands."""
@@ -399,7 +408,7 @@ async def main():
     else:
       requirements_command.append("//build:requirements.update")
 
-    result = await executor.run(requirements_command.get_command_as_string(), args.dry_run)
+    result = await executor.run(requirements_command.get_command_as_string(), args.dry_run, args.detailed_timestamped_log)
     if result.return_code != 0:
       raise RuntimeError(f"Command failed with return code {result.return_code}")
     else:
@@ -597,7 +606,7 @@ async def main():
 
       wheel_build_command.append(f"--jaxlib_git_hash={git_hash}")
 
-      result = await executor.run(wheel_build_command.get_command_as_string(), args.dry_run)
+      result = await executor.run(wheel_build_command.get_command_as_string(), args.dry_run, args.detailed_timestamped_log)
       # Exit with error if any wheel build fails.
       if result.return_code != 0:
         raise RuntimeError(f"Command failed with return code {result.return_code}")

build/tools/command.py

@@ -75,7 +75,7 @@ def __init__(self, environment: Dict[str, str] = None):
     """
     self.environment = environment or dict(os.environ)
 
-  async def run(self, cmd: str, dry_run: bool = False) -> CommandResult:
+  async def run(self, cmd: str, dry_run: bool = False, detailed_timestamped_log: bool = False) -> CommandResult:
     """
     Executes a subprocess command.
 
@@ -96,14 +96,15 @@ async def run(self, cmd: str, dry_run: bool = False) -> CommandResult:
 
     process = await asyncio.create_subprocess_shell(
       cmd,
-      stdout=asyncio.subprocess.PIPE,
-      stderr=asyncio.subprocess.PIPE,
+      stdout=asyncio.subprocess.PIPE if detailed_timestamped_log else None,
+      stderr=asyncio.subprocess.PIPE if detailed_timestamped_log else None,
       env=self.environment,
     )
 
-    await asyncio.gather(
-      _process_log_stream(process.stdout, result), _process_log_stream(process.stderr, result)
-    )
+    if detailed_timestamped_log:
+      await asyncio.gather(
+        _process_log_stream(process.stdout, result), _process_log_stream(process.stderr, result)
+      )
 
     result.return_code = await process.wait()
     result.end_time = datetime.datetime.now()

ci/build_artifacts.sh

@@ -69,7 +69,7 @@ if [[ "${allowed_artifacts[@]}" =~ "${artifact}" ]]; then
     fi
 
     # Build the artifact.
-    python build/build.py build --wheels="$artifact" --bazel_options=--config="$bazelrc_config" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose
+    python build/build.py build --wheels="$artifact" --bazel_options=--config="$bazelrc_config" --python_version=$JAXCI_HERMETIC_PYTHON_VERSION --verbose --detailed_timestamped_log
 
     # If building `jaxlib` or `jax-cuda-plugin` or `jax-cuda-pjrt` for Linux, we
     # run `auditwheel show` to verify manylinux compliance.