pytorch

1
import argparse
2
import json
3
import multiprocessing as mp
4
import os
5
import re
6
import tempfile
7
from typing import Any, Dict, List, Optional, Tuple
8

9
import requests
10
import rockset  # type: ignore[import]
11
from gitutils import retries_decorator
12

13
LOGS_QUERY = """
14
with
15
    shas as (
16
        SELECT
17
            push.head_commit.id as sha,
18
        FROM
19
            commons.push
20
        WHERE
21
            push.ref = 'refs/heads/viable/strict'
22
            AND push.repository.full_name = 'pytorch/pytorch'
23
        ORDER BY
24
            push._event_time DESC
25
        LIMIT
26
            5
27
    )
28
select
29
    id,
30
    name
31
from
32
    workflow_job j
33
    join shas on shas.sha = j.head_sha
34
where
35
    j.name like '% / test%'
36
    and j.name not like '%rerun_disabled_tests%'
37
    and j.name not like '%mem_leak_check%'
38
"""
39

40
TEST_EXISTS_QUERY = """
41
select
42
    count(*) as c
43
from
44
    test_run_s3
45
where
46
    cast(name as string) like :name
47
    and classname like :classname
48
    and _event_time > CURRENT_TIMESTAMP() - DAYS(7)
49
"""
50

51
CLOSING_COMMENT = (
52
    "I cannot find any mention of this test in rockset for the past 7 days "
53
    "or in the logs for the past 5 commits on viable/strict.  Closing this "
54
    "issue as it is highly likely that this test has either been renamed or "
55
    "removed.  If you think this is a false positive, please feel free to "
56
    "re-open this issue."
57
)
58

59
DISABLED_TESTS_JSON = (
60
    "https://ossci-metrics.s3.amazonaws.com/disabled-tests-condensed.json"
61
)
62

63

64
def parse_args() -> Any:
65
    parser = argparse.ArgumentParser()
66
    parser.add_argument(
67
        "--dry-run",
68
        action="store_true",
69
        help="Only list the tests.",
70
    )
71
    return parser.parse_args()
72

73

74
@retries_decorator()
75
def query_rockset(
76
    query: str, params: Optional[Dict[str, Any]] = None
77
) -> List[Dict[str, Any]]:
78
    res = rockset.RocksetClient(
79
        host="api.rs2.usw2.rockset.com", api_key=os.environ["ROCKSET_API_KEY"]
80
    ).sql(query, params)
81
    results: List[Dict[str, Any]] = res.results
82
    return results
83

84

85
def download_log_worker(temp_dir: str, id: int, name: str) -> None:
86
    url = f"https://ossci-raw-job-status.s3.amazonaws.com/log/{id}"
87
    data = requests.get(url).text
88
    with open(f"{temp_dir}/{name.replace('/', '_')} {id}.txt", "x") as f:
89
        f.write(data)
90

91

92
def printer(item: Tuple[str, Tuple[int, str, List[Any]]], extra: str) -> None:
93
    test, (_, link, _) = item
94
    print(f"{link:<55} {test:<120} {extra}")
95

96

97
def close_issue(num: int) -> None:
98
    headers = {
99
        "Accept": "application/vnd.github.v3+json",
100
        "Authorization": f"token {os.environ['GITHUB_TOKEN']}",
101
    }
102
    requests.post(
103
        f"https://api.github.com/repos/pytorch/pytorch/issues/{num}/comments",
104
        data=json.dumps({"body": CLOSING_COMMENT}),
105
        headers=headers,
106
    )
107
    requests.patch(
108
        f"https://api.github.com/repos/pytorch/pytorch/issues/{num}",
109
        data=json.dumps({"state": "closed"}),
110
        headers=headers,
111
    )
112

113

114
def check_if_exists(
115
    item: Tuple[str, Tuple[int, str, List[str]]], all_logs: List[str]
116
) -> Tuple[bool, str]:
117
    test, (_, link, _) = item
118
    # Test names should look like `test_a (module.path.classname)`
119
    reg = re.match(r"(\S+) \((\S*)\)", test)
120
    if reg is None:
121
        return False, "poorly formed"
122

123
    name = reg[1]
124
    classname = reg[2].split(".")[-1]
125

126
    # Check if there is any mention of the link or the test name in the logs.
127
    # The link usually shows up in the skip reason.
128
    present = False
129
    for log in all_logs:
130
        if link in log:
131
            present = True
132
            break
133
        if f"{classname}::{name}" in log:
134
            present = True
135
            break
136
    if present:
137
        return True, "found in logs"
138

139
    # Query rockset to see if the test is there
140
    count = query_rockset(
141
        TEST_EXISTS_QUERY, {"name": f"{name}%", "classname": f"{classname}%"}
142
    )
143
    if count[0]["c"] == 0:
144
        return False, "not found"
145
    return True, "found in rockset"
146

147

148
if __name__ == "__main__":
149
    args = parse_args()
150
    disabled_tests_json = json.loads(requests.get(DISABLED_TESTS_JSON).text)
151

152
    all_logs = []
153
    jobs = query_rockset(LOGS_QUERY)
154
    with tempfile.TemporaryDirectory() as temp_dir:
155
        pool = mp.Pool(20)
156
        for job in jobs:
157
            id = job["id"]
158
            name = job["name"]
159
            pool.apply_async(download_log_worker, args=(temp_dir, id, name))
160
        pool.close()
161
        pool.join()
162

163
        for filename in os.listdir(temp_dir):
164
            with open(f"{temp_dir}/{filename}") as f:
165
                all_logs.append(f.read())
166

167
    # If its less than 200 something definitely went wrong.
168
    assert len(all_logs) > 200
169
    assert len(all_logs) == len(jobs)
170

171
    to_be_closed = []
172
    for item in disabled_tests_json.items():
173
        exists, reason = check_if_exists(item, all_logs)
174
        printer(item, reason)
175
        if not exists:
176
            to_be_closed.append(item)
177

178
    print(f"There are {len(to_be_closed)} issues that will be closed:")
179
    for item in to_be_closed:
180
        printer(item, "")
181

182
    if args.dry_run:
183
        print("dry run, not actually closing")
184
    else:
185
        for item in to_be_closed:
186
            _, (num, _, _) = item
187
            close_issue(num)
188