mirror of
https://github.com/vale981/ray
synced 2025-03-06 10:31:39 -05:00
71 lines
2.3 KiB
Python
71 lines
2.3 KiB
Python
import argparse
|
|
import os
|
|
from collections import OrderedDict
|
|
|
|
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"}
|
|
assert "BUILDKITE_JOB_ID" in os.environ
|
|
assert "BUILDKITE_COMMIT" in os.environ
|
|
|
|
is_dir = os.path.isdir(args.path)
|
|
|
|
# # Assume the caller role from the instance
|
|
# sts_client = boto3.client('sts')
|
|
# assumed_role_object = sts_client.assume_role(
|
|
# RoleArn="arn:aws:iam::029272617770:role/presigner_caller_role",
|
|
# RoleSessionName="ProdSessionFromBK")
|
|
# credentials = assumed_role_object["Credentials"]
|
|
|
|
# # Construct the HTTP auth to call the API gateway
|
|
# auth = AWSRequestsAuth(
|
|
# aws_host="vop4ss7n22.execute-api.us-west-2.amazonaws.com",
|
|
# aws_region="us-west-2",
|
|
# aws_service="execute-api",
|
|
# aws_access_key=credentials["AccessKeyId"],
|
|
# aws_secret_access_key=credentials["SecretAccessKey"],
|
|
# aws_token=credentials["SessionToken"],
|
|
# )
|
|
auth = BotoAWSRequestsAuth(
|
|
aws_host="vop4ss7n22.execute-api.us-west-2.amazonaws.com",
|
|
aws_region="us-west-2",
|
|
aws_service="execute-api",
|
|
)
|
|
|
|
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", resp.status_code)
|
|
|
|
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}"
|
|
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)
|