mirror of
https://github.com/vale981/ray
synced 2025-03-05 10:01:43 -05:00
76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
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)
|