run_aws_listing.py 1.73 KB
Newer Older
1
import argparse
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
2
from collections import defaultdict
3

4
from boto3 import resource
5
6
7
8
9
10

parser = argparse.ArgumentParser()
parser.add_argument("--bucket", type=str, required=True)
parser.add_argument("--region", type=str, required=True)
args = parser.parse_args()

Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
11
12
13
ROOT_URL = f"https://{args.bucket}.s3.{args.region}.amazonaws.com/whl"
html = "<!DOCTYPE html>\n<html>\n<body>\n{}\n</body>\n</html>"
href = '  <a href="{}">{}</a><br/>'
14
html_args = {
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
15
16
17
18
    "ContentType": "text/html",
    "CacheControl": "max-age=300",
    "ACL": "public-read",
}
19

20
bucket = resource("s3").Bucket(name="nerfacc-bucket")
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
21
22
23
24
25
26
27
28
29
30
31
32

wheels_dict = defaultdict(list)
for obj in bucket.objects.filter(Prefix="whl"):
    if obj.key[-3:] != "whl":
        continue
    torch_version, wheel = obj.key.split("/")[-2:]
    wheel = f"{torch_version}/{wheel}"
    wheels_dict[torch_version].append(wheel)

index_html = html.format(
    "\n".join(
        [
33
34
35
36
            href.format(
                f"{torch_version}.html".replace("+", "%2B"), torch_version
            )
            for torch_version in wheels_dict
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
37
38
39
40
41
42
        ]
    )
)

with open("index.html", "w") as f:
    f.write(index_html)
43
bucket.Object("whl/index.html").upload_file("index.html", html_args)
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
44
45
46
47
48
49

for torch_version, wheel_names in wheels_dict.items():
    torch_version_html = html.format(
        "\n".join(
            [
                href.format(
50
51
                    f"{ROOT_URL}/{wheel_name}".replace("+", "%2B"),
                    wheel_name.split("/")[-1],
Ruilong Li(李瑞龙)'s avatar
Ruilong Li(李瑞龙) committed
52
53
54
55
56
57
58
59
60
                )
                for wheel_name in wheel_names
            ]
        )
    )

    with open(f"{torch_version}.html", "w") as f:
        f.write(torch_version_html)
    bucket.Object(f"whl/{torch_version}.html").upload_file(
61
        f"{torch_version}.html", html_args
62
    )