Return the list of artifacts from mobile job

.github/workflows/mobile_job.yml

@@ -92,14 +92,23 @@
         type: string
         default: ''
 
+    outputs:
+      artifacts:
+        description: |
+          The list of artifacts from AWS in JSON format returned to the caller
+        type: string
+        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 || input.device-types == 'android' && steps.android-test.outputs.artifacts || '[]' }}
     steps:
       - name: Clean workspace
         run: |
@@ -255,6 +264,7 @@
           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 +295,7 @@
             --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
@@ -39,12 +38,20 @@ class ReportType(Enum):
 
 
 DEVICE_FARM_BUCKET = "gha-artifacts"
+DYNAMODB_BENCHMARK_DB = "benchmark"
+DYNAMODB_BENCHMARK_TABLE = "oss_ci_benchmark_v2"
 
 logging.basicConfig(level=logging.INFO)
 
 
 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 +62,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 +85,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 +107,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 +274,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: 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: str,
+    indent: int = 0,
 ) -> List[Dict[str, str]]:
     """
     Return all artifacts from this specific test. There are 3 types of artifacts
@@ -293,17 +329,21 @@ 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
 
@@ -346,7 +386,7 @@ def print_report(
         next_rtype = ReportType.TEST
     elif rtype == ReportType.TEST:
         return print_test_artifacts(
-            client, arn, workflow_id, workflow_attempt, indent + 2
+            client, arn, workflow_id, workflow_attempt, name, indent + 2
         )
 
     artifacts = []
@@ -556,9 +596,10 @@ def main() -> None:
         warn(f"Failed to run {unique_prefix}: {error}")
         sys.exit(1)
     finally:
-        print_report(
+        artifacts = print_report(
             client, r.get("run"), ReportType.RUN, workflow_id, workflow_attempt
         )
+        set_output("artifacts", json.dumps(artifacts))
 
     if not is_success(result):
         sys.exit(1)