[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"
@@ -711,6 +710,7 @@ def isValidSlurmJobId(def slurmJobID) {
 def cleanUpSlurmResources(def pipeline, SlurmCluster cluster, String clusterName, String jobUID){
     CloudManager.withSlurmFrontendFailover(pipeline, clusterName, cluster) { remote ->
         def jobWorkspace = "/home/svc_tensorrt/bloom/scripts/${jobUID}"
+        def s3SpoolRoot = "/home/svc_tensorrt/bloom/scripts/.s3-spool-${jobUID}"
 
         Utils.exec(pipeline, script: "echo Sleeping to allow Slurm job completion; sleep 30")
 
@@ -746,7 +746,7 @@ def cleanUpSlurmResources(def pipeline, SlurmCluster cluster, String clusterName
             // instead of deleting per job; reused images keep a refreshed mtime.
             "find ${cluster.scratchPath}/users/svc_tensorrt/containers -maxdepth 1 -name 'container-*.sqsh' -mtime +3 -delete 2>/dev/null || true",
             "find ${cluster.scratchPath}/users/svc_tensorrt/containers -maxdepth 1 \\( -name 'container-*.tmp' -o -name 'container-*.lock' \\) -mtime +1 -delete 2>/dev/null || true",
-            "rm -rf ${jobWorkspace} || true",
+            "rm -rf ${jobWorkspace} ${s3SpoolRoot} || true",
         ].join(" ; ")
         Utils.exec(
             pipeline,
@@ -1529,9 +1529,6 @@ def getPytestBaseCommandLine(
     }
     def unittestMarkExpr = (stageName.startsWith("CPU-")) ? "cpu_only" : "not cpu_only"
     testCmdLine += ["--unittest-markexpr='${unittestMarkExpr}'"]
-    if (ENABLE_UPLOAD_TEST_RESULTS) {
-        testCmdLine += ["-o console_output_style=progress-even-when-capture-no"]
-    }
     if (extraArgs) {
         testCmdLine += extraArgs
     }
@@ -1796,17 +1793,12 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG
                     "--group $splitId",
                     *clusterDurationsArgsNode,
                 ]
-                if (ENABLE_UPLOAD_TEST_RESULTS) {
+                if (ENABLE_UPLOAD_TEST_RESULTS && !testFilter[(DETAILED_LOG)]) {
                     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,
@@ -4488,17 +4480,12 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
         // Temporarily disable to reduce the log size
         // sh 'if [ "$(id -u)" -eq 0 ]; then dmesg -C || true; fi'
         def extraArgs = [*clusterDurationsArgs]
-        if (ENABLE_UPLOAD_TEST_RESULTS) {
+        if (ENABLE_UPLOAD_TEST_RESULTS && !testFilter[(DETAILED_LOG)]) {
             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,
@@ -4533,57 +4520,66 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
             ]) {
                 sh "env | sort"
                 try {
-                    if (preprocessedLists.regularCount > 0) {
-                        sh """
-                            rm -rf ${stageName}/ && \
-                            cd ${llmSrc}/tests/integration/defs && \
-                            ${pytestCommand.join(" ")}
-                        """
-                    } else {
-                        echo "No regular tests to run for stage ${stageName}"
-                        noRegularTests = true
-                        sh "mkdir -p ${stageName}"
-                        // Create an empty results.xml file for consistency
-                        sh """
-                            echo '<?xml version="1.0" encoding="UTF-8"?>' > ${stageName}/results.xml
-                            echo '<testsuites>' >> ${stageName}/results.xml
-                            echo '<testsuite name="${stageName}" errors="0" failures="0" skipped="0" tests="0" time="0.0">' >> ${stageName}/results.xml
-                            echo '</testsuite>' >> ${stageName}/results.xml
-                            echo '</testsuites>' >> ${stageName}/results.xml
-                        """
-                    }
-                } catch (InterruptedException e) {
-                    throw e
-                } catch (Exception e) {
-                    def isRerunFailed = rerunFailedTests(stageName, llmSrc, pytestCommand, "results.xml", "regular")
-                    if (isRerunFailed) {
-                        catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') {
-                            error "Regular tests failed after rerun attempt"
+                    try {
+                        if (preprocessedLists.regularCount > 0) {
+                            sh """
+                                rm -rf ${stageName}/ && \
+                                cd ${llmSrc}/tests/integration/defs && \
+                                ${pytestCommand.join(" ")}
+                            """
+                        } else {
+                            echo "No regular tests to run for stage ${stageName}"
+                            noRegularTests = true
+                            sh "mkdir -p ${stageName}"
+                            // Create an empty results.xml file for consistency
+                            sh """
+                                echo '<?xml version="1.0" encoding="UTF-8"?>' > ${stageName}/results.xml
+                                echo '<testsuites>' >> ${stageName}/results.xml
+                                echo '<testsuite name="${stageName}" errors="0" failures="0" skipped="0" tests="0" time="0.0">' >> ${stageName}/results.xml
+                                echo '</testsuite>' >> ${stageName}/results.xml
+                                echo '</testsuites>' >> ${stageName}/results.xml
+                            """
                         }
-                        rerunFailed = true
-                    } else if (generateTimeoutTestResultXml(pipeline, stageName)) {
-                        // Rerun passed but the first run had a timeout: mark this
-                        // stage FAILURE so "[${stageName}] Run Pytest" turns red,
-                        // not just the enclosing parent stage.
-                        catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') {
-                            error "Some tests terminated unexpectedly, please check the test report."
+                    } catch (InterruptedException e) {
+                        throw e
+                    } catch (Exception e) {
+                        def isRerunFailed = rerunFailedTests(stageName, llmSrc, pytestCommand, "results.xml", "regular")
+                        if (isRerunFailed) {
+                            catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') {
+                                error "Regular tests failed after rerun attempt"
+                            }
+                            rerunFailed = true
+                        } else if (generateTimeoutTestResultXml(pipeline, stageName)) {
+                            // Rerun passed but the first run had a timeout: mark this
+                            // stage FAILURE so "[${stageName}] Run Pytest" turns red,
+                            // not just the enclosing parent stage.
+                            catchError(buildResult: 'SUCCESS', stageResult: 'FAILURE') {
+                                error "Some tests terminated unexpectedly, please check the test report."
+                            }
                         }
                     }
-                }
 
-                // Run the isolated tests if exists
-                if (preprocessedLists.isolateCount > 0) {
-                    stage ("[${stageName}] Run Pytest (Isolated)") {
-                        echo "There are ${preprocessedLists.isolateCount} isolated tests to run"
-                        rerunFailed = runIsolatedTests(preprocessedLists, pytestCommand, llmSrc, stageName) || rerunFailed
+                    // Run the isolated tests if exists
+                    if (preprocessedLists.isolateCount > 0) {
+                        stage ("[${stageName}] Run Pytest (Isolated)") {
+                            echo "There are ${preprocessedLists.isolateCount} isolated tests to run"
+                            rerunFailed = runIsolatedTests(preprocessedLists, pytestCommand, llmSrc, stageName) || rerunFailed
+                        }
+                    } else {
+                        echo "No isolated tests to run for stage ${stageName}"
+                        noIsolateTests = true
                     }
-                } else {
-                    echo "No isolated tests to run for stage ${stageName}"
-                    noIsolateTests = true
-                }
 
-                if (noRegularTests && noIsolateTests) {
-                    error "No tests were executed for stage ${stageName}, please check the test list and test-db rendering result."
+                    if (noRegularTests && noIsolateTests) {
+                        error "No tests were executed for stage ${stageName}, please check the test list and test-db rendering result."
+                    }
+                } finally {
+                    if (ENABLE_UPLOAD_TEST_RESULTS && !testFilter[(DETAILED_LOG)]) {
+                        sh """
+                            python3 ${llmSrc}/tests/test_common/s3_output.py \
+                                --drain-spool "${WORKSPACE}/${stageName}" || true
+                        """
+                    }
                 }
             }
 

jenkins/scripts/slurm_run.sh

@@ -82,6 +82,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/pytest.ini

@@ -1,5 +1,7 @@
 [pytest]
 asyncio_default_fixture_loop_scope = module
+log_format = %(asctime)s.%(msecs)03d %(levelname)-8s %(name)s:%(filename)s:%(lineno)d %(message)s
+log_date_format = %Y-%m-%d %H:%M:%S
 threadleak = True
 # Thread-\d+ \(_manager_spawn\) / session-reuse-* / session-prefetch-* belong to
 # a pool cached for reuse by the NEXT test (tests/test_common/session_reuse.py)

tests/integration/defs/test_unittests.py

@@ -180,12 +180,15 @@ def test_unittests_v2(llm_root, llm_venv, case: str, output_dir, request):
         command += ["--run-ray"]
 
     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",
@@ -262,6 +265,18 @@ def run_command(cmd, num_workers=1):
                 )
             print(f"{'='*60}\n")
             return False
+        finally:
+            if s3_output_module is not None:
+                try:
+                    drained = s3_output_module.drain_pending_uploads(
+                        inner_output_dir,
+                        secret_key=s3_secret_key,
+                    )
+                except Exception as e:
+                    print(f"Failed to drain pending S3 test logs: {e}")
+                else:
+                    if not drained:
+                        print("Some pending S3 test logs could not be drained")
         return True
 
     if num_workers == 1: