12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576 |
- import argparse
- import os
- from collections import OrderedDict
- import sys
- import time
- from aws_requests_auth.boto_utils import BotoAWSRequestsAuth
- import requests
- parser = argparse.ArgumentParser(
- description="Helper script to upload files to S3 bucket")
- parser.add_argument("--path", type=str)
- parser.add_argument("--destination", type=str)
- args = parser.parse_args()
- assert os.path.exists(args.path)
- assert args.destination in {"wheels", "containers", "logs"}
- assert "BUILDKITE_JOB_ID" in os.environ
- assert "BUILDKITE_COMMIT" in os.environ
- is_dir = os.path.isdir(args.path)
- auth = BotoAWSRequestsAuth(
- aws_host="vop4ss7n22.execute-api.us-west-2.amazonaws.com",
- aws_region="us-west-2",
- aws_service="execute-api",
- )
- for _ in range(5):
- resp = requests.get(
- "https://vop4ss7n22.execute-api.us-west-2.amazonaws.com/endpoint/",
- auth=auth,
- params={"job_id": os.environ["BUILDKITE_JOB_ID"]})
- print("Getting Presigned URL, status_code", resp.status_code)
- if resp.status_code >= 500:
- print("errored, retrying...")
- print(resp.text)
- time.sleep(5)
- else:
- break
- if resp.status_code >= 500:
- print("still errorred after many retries")
- sys.exit(1)
- sha = os.environ["BUILDKITE_COMMIT"]
- if is_dir:
- paths = [os.path.join(args.path, f) for f in os.listdir(args.path)]
- else:
- paths = [args.path]
- print("Planning to upload", paths)
- for path in paths:
- fn = os.path.split(path)[-1]
- if args.destination == "wheels":
- c = resp.json()["presigned_wheels"]
- of = OrderedDict(c["fields"])
- of["key"] = f"scratch/bk/{sha}/{fn}"
- elif args.destination == "containers":
- c = resp.json()["presigned_containers"]
- of = OrderedDict(c["fields"])
- of["key"] = f"{sha}/{fn}"
- elif args.destination == "logs":
- c = resp.json()["presigned_logs"]
- of = OrderedDict(c["fields"])
- branch = os.environ["BUILDKITE_BRANCH"]
- bk_job_id = os.environ["BUILDKITE_JOB_ID"]
- of["key"] = f"bazel_events/{branch}/{sha}/{bk_job_id}/{fn}"
- else:
- raise ValueError("Unknown destination")
- of["file"] = open(path, "rb")
- r = requests.post(c["url"], files=of)
- print(f"Uploaded {path} to {of['key']}", r.status_code)
|