[TRTLLMINF-191][infra] Use native pytest capture for S3 logs

jenkins/L0_Test.groovy

@@ -47,7 +47,6 @@ ARTIFACT_PATH = env.artifactPath ? env.artifactPath : "sw-tensorrt-generic/llm-a
 UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}"
 URM_ARTIFACTORY_BASE = "https://urm.nvidia.com/artifactory"
 ENABLE_UPLOAD_TEST_RESULTS = params.enableUploadTestResults != null ? params.enableUploadTestResults : true
-ENABLE_S3_ECHO_STDOUT = params.enableS3EchoStdout != null ? params.enableS3EchoStdout : false
 
 X86_64_TRIPLE = "x86_64-linux-gnu"
 AARCH64_TRIPLE = "aarch64-linux-gnu"
@@ -1480,15 +1479,10 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG
                 ]
                 if (ENABLE_UPLOAD_TEST_RESULTS) {
                     extraArgs += [
-                        "-s",
+                        "--capture=fd",
                         "--s3-upload-path=${uploadPath}/${stageName}",
+                        "--s3-upload-mode=deferred",
                     ]
-                    if (ENABLE_S3_ECHO_STDOUT) {
-                        extraArgs += [
-                            "--s3-echo-stdout",
-                            "--s3-capture-mode=timestamped",
-                        ]
-                    }
                 }
                 def pytestCommand = getPytestBaseCommandLine(
                     llmSrcNode,
@@ -4012,15 +4006,10 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
         def extraArgs = [*clusterDurationsArgs]
         if (ENABLE_UPLOAD_TEST_RESULTS) {
             extraArgs += [
-                "-s",
+                "--capture=fd",
                 "--s3-upload-path=${uploadPath}/${stageName}",
+                "--s3-upload-mode=deferred",
             ]
-            if (ENABLE_S3_ECHO_STDOUT) {
-                extraArgs += [
-                    "--s3-echo-stdout",
-                    "--s3-capture-mode=timestamped",
-                ]
-            }
         }
         def pytestCommand = getPytestBaseCommandLine(
             llmSrc,

jenkins/scripts/slurm_run.sh

@@ -103,6 +103,8 @@ perf_report_exit_code=0
 eval $pytestCommand
 pytest_exit_code=$?
 echo "Rank${SLURM_PROCID} Pytest finished execution with exit code $pytest_exit_code"
+python3 "$llmSrcNode/tests/test_common/s3_output.py" \
+    --drain-spool "$jobWorkspace" || true
 
 # DEBUG: Diagnose intermittent "unrecognized arguments" failure (Exit Code 4)
 # Remove this after the issue is resolved

tests/integration/defs/test_unittests.py

@@ -175,12 +175,15 @@ def test_unittests_v2(llm_root, llm_venv, case: str, output_dir, request):
         command += ["-m", unittest_markexpr]
 
     s3_secret_key = None
+    s3_output_module = None
     s3_upload_path = request.config.getoption("--s3-upload-path", default=None)
     if s3_upload_path:
+        from test_common import s3_output as s3_output_module
+
         inner_output_dir = os.path.join(output_dir, "inner-s3", case_fn)
         inner_upload_path = os.path.join(s3_upload_path, "inner", case_fn)
         command += [
-            "-s",
+            "--capture=fd",
             f"--output-dir={inner_output_dir}",
             f"--s3-upload-path={inner_upload_path}",
             "--s3-upload-mode=deferred",
@@ -257,6 +260,12 @@ def run_command(cmd, num_workers=1):
                 )
             print(f"{'='*60}\n")
             return False
+        finally:
+            if s3_output_module is not None:
+                s3_output_module.drain_pending_uploads(
+                    inner_output_dir,
+                    secret_key=s3_secret_key,
+                )
         return True
 
     if num_workers == 1: