Return the list of artifacts from mobile job

.github/workflows/mobile_job.yml

@@ -92,14 +92,22 @@ on:
         type: string
         default: ''
 
+    outputs:
+      artifacts:
+        description: |
+          The list of artifacts from AWS in JSON format returned to the caller
+        value: ${{ jobs.mobile.outputs.artifacts }}
+
 jobs:
-  job:
+  mobile:
     name: ${{ inputs.job-name }} (${{ inputs.device-type }})
     runs-on: ${{ inputs.runner }}
     timeout-minutes: ${{ inputs.timeout }}
     permissions:
       id-token: write
       contents: read
+    outputs:
+      artifacts: ${{ inputs.device-type == 'ios' && steps.ios-test.outputs.artifacts || inputs.device-type == 'android' && steps.android-test.outputs.artifacts || '[]' }}
     steps:
       - name: Clean workspace
         run: |
@@ -112,7 +120,7 @@ jobs:
         uses: aws-actions/configure-aws-credentials@v3
         with:
           role-to-assume: arn:aws:iam::308535385114:role/gha_workflow_mobile_job
-          # This could go up to 18000, the max duration enforced by the server side
+          # The max duration enforced by the server side
           role-duration-seconds: 18000
           aws-region: us-east-1
 
@@ -254,7 +262,15 @@ jobs:
 
           echo "extra-data-output=${EXTRA_DATA_OUTPUT}" >> "${GITHUB_OUTPUT}"
 
+      - name: Get workflow job id
+        id: get-job-id
+        uses: ./test-infra/.github/actions/get-workflow-job-id
+        if: always()
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
       - name: Run iOS tests on devices
+        id: ios-test
         if: ${{ inputs.device-type == 'ios' }}
         shell: bash
         working-directory: test-infra/tools/device-farm-runner
@@ -270,6 +286,7 @@ jobs:
           DEVICE_TYPE: ${{ inputs.device-type }}
           RUN_ID: ${{ github.run_id }}
           RUN_ATTEMPT: ${{ github.run_attempt }}
+          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
         run: |
           set -ex
 
@@ -282,9 +299,23 @@ jobs:
             --test-spec "${TEST_SPEC}" \
             --name-prefix "${JOB_NAME}-${DEVICE_TYPE}" \
             --workflow-id "${RUN_ID}" \
-            --workflow-attempt "${RUN_ATTEMPT}"
+            --workflow-attempt "${RUN_ATTEMPT}" \
+            --output "ios-artifacts-${JOB_ID}.json"
+
+      - name: Upload iOS artifacts to S3
+        uses: seemethere/upload-artifact-s3@v5
+        if: ${{ inputs.device-type == 'ios' }}
+        with:
+          name: ios-artifacts
+          retention-days: 14
+          s3-bucket: gha-artifacts
+          s3-prefix: |
+            device_farm/${{ github.run_id }}/${{ github.run_attempt }}/artifacts
+          path: |
+            test-infra/tools/device-farm-runner/ios-artifacts-${{ steps.get-job-id.outputs.job-id }}.json
 
       - name: Run Android tests on devices
+        id: android-test
         if: ${{ inputs.device-type == 'android' }}
         shell: bash
         working-directory: test-infra/tools/device-farm-runner
@@ -300,6 +331,7 @@ jobs:
           DEVICE_TYPE: ${{ inputs.device-type }}
           RUN_ID: ${{ github.run_id }}
           RUN_ATTEMPT: ${{ github.run_attempt }}
+          JOB_ID: ${{ steps.get-job-id.outputs.job-id }}
         run: |
           set -ex
 
@@ -312,4 +344,17 @@ jobs:
             --test-spec "${TEST_SPEC}" \
             --name-prefix "${JOB_NAME}-${DEVICE_TYPE}" \
             --workflow-id "${RUN_ID}" \
-            --workflow-attempt "${RUN_ATTEMPT}"
+            --workflow-attempt "${RUN_ATTEMPT}" \
+            --output "android-artifacts-${JOB_ID}.json"
+
+      - name: Upload Android artifacts to S3
+        uses: seemethere/upload-artifact-s3@v5
+        if: ${{ inputs.device-type == 'android' }}
+        with:
+          name: android-artifacts
+          retention-days: 14
+          s3-bucket: gha-artifacts
+          s3-prefix: |
+            device_farm/${{ github.run_id }}/${{ github.run_attempt }}/artifacts
+          path: |
+            test-infra/tools/device-farm-runner/android-artifacts-${{ steps.get-job-id.outputs.job-id }}.json

tools/device-farm-runner/run_on_aws_devicefarm.py

@@ -1,17 +1,17 @@
 #!/usr/bin/env python3
 
+import json
 import datetime
 import logging
 import os
 import random
-import re
 import string
 import sys
 import time
-from argparse import Action, ArgumentParser
+from argparse import Action, ArgumentParser, Namespace
 from enum import Enum
 from logging import info
-from typing import Any, Dict, List, Optional, Pattern
+from typing import Any, Dict, List, Optional
 from warnings import warn
 
 import boto3
@@ -27,7 +27,6 @@
 AWS_GUID = "082d10e5-d7d7-48a5-ba5c-b33d66efa1f5"
 AWS_ARN_PREFIX = "arn:aws:devicefarm:"
 DEFAULT_DEVICE_POOL_ARN = f"{AWS_ARN_PREFIX}{AWS_REGION}::devicepool:{AWS_GUID}"
-TEST_SPEC_OUTPUT_HIGHLIGHT_HINT = re.compile(r"^\[PyTorch\] .*$")
 
 
 # Device Farm report type
@@ -44,7 +43,13 @@ class ReportType(Enum):
 
 
 class ValidateArchive(Action):
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: Any,
+        option_string: Optional[str] = None,
+    ) -> None:
         if values.startswith(AWS_ARN_PREFIX) or (
             os.path.isfile(values) and values.endswith(".zip")
         ):
@@ -55,7 +60,13 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 
 class ValidateExtraDataArchive(Action):
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: Any,
+        option_string: Optional[str] = None,
+    ) -> None:
         # This parameter is optional and can accept an empty string, or it can be
         # an existing ARN, or a local zip archive to be uploaded to AWS
         if (
@@ -72,7 +83,13 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 
 class ValidateApp(Action):
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: Any,
+        option_string: Optional[str] = None,
+    ) -> None:
         # This can be a local file or an existing app that has previously been uploaded
         # to AWS
         if values.startswith(AWS_ARN_PREFIX) or (
@@ -88,7 +105,13 @@ def __call__(self, parser, namespace, values, option_string=None):
 
 
 class ValidateTestSpec(Action):
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(
+        self,
+        parser: ArgumentParser,
+        namespace: Namespace,
+        values: Any,
+        option_string: Optional[str] = None,
+    ) -> None:
         if values.startswith(AWS_ARN_PREFIX) or (
             os.path.isfile(values)
             and (values.endswith(".yml") or values.endswith(".yaml"))
@@ -167,6 +190,11 @@ def parse_args() -> Any:
         default=0,
         help="the workflow run attempt",
     )
+    parser.add_argument(
+        "--output",
+        type=str,
+        help="an optional file to write the list of artifacts from AWS in JSON format",
+    )
 
     return parser.parse_args()
 
@@ -249,24 +277,40 @@ def upload_file_to_s3(
     )
 
 
-def print_highlights(
+def set_output(val: Any, gh_var_name: str, filename: Optional[str]) -> None:
+    if os.getenv("GITHUB_OUTPUT"):
+        with open(str(os.getenv("GITHUB_OUTPUT")), "a") as env:
+            print(f"{gh_var_name}={val}", file=env)
+    else:
+        print(f"::set-output name={gh_var_name}::{val}")
+
+    # Also write the value to file if it exists
+    if filename:
+        with open(filename, "w") as f:
+            print(val, file=f)
+
+
+def print_testspec(
+    report_name: Optional[str],
     file_name: str,
-    highlight: Pattern,
     indent: int = 0,
-):
+) -> None:
     """
     The test spec output from AWS Device Farm is the main output of the test job.
-    It's a bit too verbose, so this is a workaround to print only the relevant
-    parts
     """
+    print(f"::group::{report_name} test output")
     with open(file_name) as f:
-        for line in f:
-            if re.match(highlight, line):
-                info(f"{' ' * indent}{line.strip()}")
+        print(f.read())
+    print("::endgroup::")
 
 
 def print_test_artifacts(
-    client: Any, test_arn: str, workflow_id: str, workflow_attempt: int, indent: int = 0
+    client: Any,
+    test_arn: str,
+    workflow_id: str,
+    workflow_attempt: int,
+    report_name: Optional[str],
+    indent: int = 0,
 ) -> List[Dict[str, str]]:
     """
     Return all artifacts from this specific test. There are 3 types of artifacts
@@ -293,25 +337,30 @@ def print_test_artifacts(
                 s3_key,
             )
 
+            s3_url = f"https://{DEVICE_FARM_BUCKET}.s3.amazonaws.com/{s3_key}"
+            artifact["s3_url"] = s3_url
+
             info(
                 f"{' ' * indent}Saving {artifact_type} {filename}.{extension} ({filetype}) "
-                + f"at https://{DEVICE_FARM_BUCKET}.s3.amazonaws.com/{s3_key}"
+                + f"at {s3_url}"
             )
+
+            # Some more metadata to identify where the artifact comes from
+            artifact["report_name"] = report_name
             gathered_artifacts.append(artifact)
 
-            # Additional step to print highlights from the test output
+            # Additional step to print the test output
             if filetype == "TESTSPEC_OUTPUT":
-                print_highlights(
-                    local_filename, TEST_SPEC_OUTPUT_HIGHLIGHT_HINT, indent + 2
-                )
+                print_testspec(report_name, local_filename, indent + 2)
 
     return gathered_artifacts
 
 
 def print_report(
     client: Any,
     report: Dict[str, Any],
-    rtype: ReportType,
+    report_name: Optional[str],
+    report_type: ReportType,
     workflow_id: str,
     workflow_attempt: int,
     indent: int = 0,
@@ -325,37 +374,42 @@ def print_report(
         return []
 
     name = report["name"]
+    # Keep the top-level report name as the name of the whole test report, this
+    # is used to connect all artifacts from one report together
+    if not report_name:
+        report_name = name
     result = report["result"]
 
     extra_msg = ""
-    if rtype == ReportType.SUITE or is_success(result):
+    if report_type == ReportType.SUITE or is_success(result):
         counters = report["counters"]
         extra_msg = f"with stats {counters}"
 
     info(f"{' ' * indent}{name} {result} {extra_msg}")
 
     arn = report["arn"]
-    if rtype == ReportType.RUN:
+    if report_type == ReportType.RUN:
         more_reports = client.list_jobs(arn=arn)
-        next_rtype = ReportType.JOB
-    elif rtype == ReportType.JOB:
+        next_report_type = ReportType.JOB
+    elif report_type == ReportType.JOB:
         more_reports = client.list_suites(arn=arn)
-        next_rtype = ReportType.SUITE
-    elif rtype == ReportType.SUITE:
+        next_report_type = ReportType.SUITE
+    elif report_type == ReportType.SUITE:
         more_reports = client.list_tests(arn=arn)
-        next_rtype = ReportType.TEST
-    elif rtype == ReportType.TEST:
+        next_report_type = ReportType.TEST
+    elif report_type == ReportType.TEST:
         return print_test_artifacts(
-            client, arn, workflow_id, workflow_attempt, indent + 2
+            client, arn, workflow_id, workflow_attempt, report_name, indent + 2
         )
 
     artifacts = []
-    for more_report in more_reports.get(f"{next_rtype.value}s", []):
+    for more_report in more_reports.get(f"{next_report_type.value}s", []):
         artifacts.extend(
             print_report(
                 client,
                 more_report,
-                next_rtype,
+                report_name,
+                next_report_type,
                 workflow_id,
                 workflow_attempt,
                 indent + 2,
@@ -556,9 +610,10 @@ def main() -> None:
         warn(f"Failed to run {unique_prefix}: {error}")
         sys.exit(1)
     finally:
-        print_report(
-            client, r.get("run"), ReportType.RUN, workflow_id, workflow_attempt
+        artifacts = print_report(
+            client, r.get("run"), None, ReportType.RUN, workflow_id, workflow_attempt
         )
+        set_output(json.dumps(artifacts), "artifacts", args.output)
 
     if not is_success(result):
         sys.exit(1)