diff --git a/.github/scripts/trymerge.py b/.github/scripts/trymerge.py index bf83154708..84f4400758 100644 --- a/.github/scripts/trymerge.py +++ b/.github/scripts/trymerge.py @@ -1163,7 +1163,6 @@ def merge_into( # Finally, upload the record to Rockset. The list of pending and failed # checks are at the time of the merge save_merge_record( - collection=ROCKSET_MERGES_COLLECTION, comment_id=comment_id, pr_num=self.pr_num, owner=self.org, @@ -1179,10 +1178,8 @@ def merge_into( merge_base_sha=self.get_merge_base(), merge_commit_sha=merge_commit_sha, is_failed=False, - dry_run=dry_run, skip_mandatory_checks=skip_mandatory_checks, ignore_current=bool(ignore_current_checks), - workspace=ROCKSET_MERGES_WORKSPACE, ) else: print("Missing comment ID or PR number, couldn't upload to Rockset") @@ -1489,7 +1486,6 @@ def checks_to_markdown_bullets( @retries_decorator() def save_merge_record( - collection: str, comment_id: int, pr_num: int, owner: str, @@ -1505,59 +1501,44 @@ def save_merge_record( merge_base_sha: str, merge_commit_sha: str = "", is_failed: bool = False, - dry_run: bool = False, skip_mandatory_checks: bool = False, ignore_current: bool = False, error: str = "", - workspace: str = "commons", ) -> None: """ - This saves the merge records into Rockset, so we can query them (for fun and profit) + This saves the merge records as a json, which can later be uploaded to s3 """ - if dry_run: - # Decide not to save the record to Rockset if dry-run is set to not pollute - # the collection - return - - try: - import rockset # type: ignore[import] - - # Prepare the record to be written into Rockset - data = [ - { - "comment_id": comment_id, - "pr_num": pr_num, - "owner": owner, - "project": project, - "author": author, - "pending_checks": pending_checks, - "failed_checks": failed_checks, - "ignore_current_checks": ignore_current_checks, - "broken_trunk_checks": broken_trunk_checks, - "flaky_checks": flaky_checks, - "unstable_checks": unstable_checks, - "last_commit_sha": last_commit_sha, - "merge_base_sha": merge_base_sha, - "merge_commit_sha": merge_commit_sha, - "is_failed": is_failed, - "skip_mandatory_checks": skip_mandatory_checks, - "ignore_current": ignore_current, - "error": error, - } - ] - client = rockset.RocksetClient( - host="api.usw2a1.rockset.com", api_key=os.environ["ROCKSET_API_KEY"] - ) - client.Documents.add_documents( - collection=collection, - data=data, - workspace=workspace, - ) + # Prepare the record to be written into Rockset + data = [ + { + "comment_id": comment_id, + "pr_num": pr_num, + "owner": owner, + "project": project, + "author": author, + "pending_checks": pending_checks, + "failed_checks": failed_checks, + "ignore_current_checks": ignore_current_checks, + "broken_trunk_checks": broken_trunk_checks, + "flaky_checks": flaky_checks, + "unstable_checks": unstable_checks, + "last_commit_sha": last_commit_sha, + "merge_base_sha": merge_base_sha, + "merge_commit_sha": merge_commit_sha, + "is_failed": is_failed, + "skip_mandatory_checks": skip_mandatory_checks, + "ignore_current": ignore_current, + "error": error, + # This is a unique identifier for the record for deduping purposes + # in rockset. Any unique string would work + "_id": f"{project}-{pr_num}-{comment_id}-{os.environ.get('GITHUB_RUN_ID')}", + } + ] + repo_root = Path(__file__).resolve().parent.parent.parent - except ModuleNotFoundError: - print("Rockset is missing, no record will be saved") - return + with open(repo_root / "merge_record.json", "w") as f: + json.dump(data, f) @retries_decorator(rc=[]) @@ -2374,7 +2355,6 @@ def handle_exception(e: Exception, title: str = "Merge failed") -> None: # list of pending and failed checks here, but they are not really # needed at the moment save_merge_record( - collection=ROCKSET_MERGES_COLLECTION, comment_id=args.comment_id, pr_num=args.pr_num, owner=org, @@ -2389,11 +2369,9 @@ def handle_exception(e: Exception, title: str = "Merge failed") -> None: last_commit_sha=pr.last_commit().get("oid", ""), merge_base_sha=pr.get_merge_base(), is_failed=True, - dry_run=args.dry_run, skip_mandatory_checks=args.force, ignore_current=args.ignore_current, error=str(e), - workspace=ROCKSET_MERGES_WORKSPACE, ) else: print("Missing comment ID or PR number, couldn't upload to Rockset") diff --git a/.github/workflows/trymerge.yml b/.github/workflows/trymerge.yml index ab8a58c2e7..1db4926f9d 100644 --- a/.github/workflows/trymerge.yml +++ b/.github/workflows/trymerge.yml @@ -9,6 +9,8 @@ jobs: name: try_merge_pr_${{ github.event.client_payload.pr_num }} runs-on: ubuntu-latest environment: pytorchbot-env + permissions: + id-token: write env: GH_RUN_URL: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }} steps: @@ -45,6 +47,7 @@ jobs: IGNORE_CURRENT: ${{ github.event.client_payload.ignore_current }} ROCKSET_API_KEY: ${{ secrets.ROCKSET_API_KEY }} DRCI_BOT_KEY: ${{ secrets.DRCI_BOT_KEY }} + GITHUB_RUN_ID: ${{ github.run_id }} run: | set -x if [ -n "${FORCE}" ]; then @@ -65,6 +68,22 @@ jobs: python3 .github/scripts/trymerge.py "${PR_NUM}" fi + - name: configure aws credentials + uses: aws-actions/configure-aws-credentials@v3 + continue-on-error: true + with: + role-to-assume: arn:aws:iam::308535385114:role/upload_to_ossci_raw_job_status + aws-region: us-east-1 + + - name: Upload merge record to s3 + if: always() + continue-on-error: true + uses: seemethere/upload-artifact-s3@v5 + with: + s3-bucket: ossci-raw-job-status + s3-prefix: merges/${{ github.repository }}/${{ github.event.client_payload.pr_num }}/${{ github.event.client_payload.comment_id }}/${{ github.run_id }} + path: merge_record.json + # We want newer merge commands to supercede old ones concurrency: group: try-merge-${{ github.event.client_payload.pr_num }} diff --git a/.gitignore b/.gitignore index cba6feab0d..5c201d1b34 100644 --- a/.gitignore +++ b/.gitignore @@ -127,6 +127,7 @@ env .circleci/scripts/COMMIT_MSG scripts/release_notes/*.json sccache-stats*.json +merge_record.json # These files get copied over on invoking setup.py torchgen/packaged/*