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: |
@@ -255,6 +263,7 @@ jobs:
           echo "extra-data-output=${EXTRA_DATA_OUTPUT}" >> "${GITHUB_OUTPUT}"
 
       - name: Run iOS tests on devices
+        id: ios-test
         if: ${{ inputs.device-type == 'ios' }}
         shell: bash
         working-directory: test-infra/tools/device-farm-runner
@@ -285,6 +294,7 @@ jobs:
             --workflow-attempt "${RUN_ATTEMPT}"
 
       - name: Run Android tests on devices
+        id: android-test
         if: ${{ inputs.device-type == 'android' }}
         shell: bash
         working-directory: test-infra/tools/device-farm-runner

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"))
@@ -249,24 +272,35 @@ def upload_file_to_s3(
     )
 
 
-def print_highlights(
+def set_output(name: str, val: Any) -> None:
+    if os.getenv("GITHUB_OUTPUT"):
+        with open(str(os.getenv("GITHUB_OUTPUT")), "a") as env:
+            print(f"{name}={val}", file=env)
+    else:
+        print(f"::set-output name={name}::{val}")
+
+
+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 +327,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 +364,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 +600,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("artifacts", json.dumps(artifacts))
 
     if not is_success(result):
         sys.exit(1)