add flag to enable detailed timestamped logging

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}")
@@ -607,4 +616,4 @@ async def main():
 
 
 if __name__ == "__main__":
-  asyncio.run(main())
+  asyncio.run(main())

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.