tests_fetcher.py 32.3 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# coding=utf-8
# Copyright 2021 The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Sylvain Gugger's avatar
Sylvain Gugger committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
"""
Welcome to tests_fetcher V2.
This util is designed to fetch tests to run on a PR so that only the tests impacted by the modifications are run, and
when too many models are being impacted, only run the tests of a subset of core models. It works like this.

Stage 1: Identify the modified files. This takes all the files from the branching point to the current commit (so
all modifications in a PR, not just the last commit) but excludes modifications that are on docstrings or comments
only.

Stage 2: Extract the tests to run. This is done by looking at the imports in each module and test file: if module A
imports module B, then changing module B impacts module A, so the tests using module A should be run. We thus get the
dependencies of each model and then recursively builds the 'reverse' map of dependencies to get all modules and tests
impacted by a given file. We then only keep the tests (and only the code models tests if there are too many modules).

Caveats:
  - This module only filters tests by files (not individual tests) so it's better to have tests for different things
    in different files.
  - This module assumes inits are just importing things, not really building objects, so it's better to structure
    them this way and move objects building in separate submodules.
"""

37
38
import argparse
import collections
Yih-Dar's avatar
Yih-Dar committed
39
import json
40
41
42
43
44
45
46
47
import os
import re
from contextlib import contextmanager
from pathlib import Path

from git import Repo


Sylvain Gugger's avatar
Sylvain Gugger committed
48
49
50
51
52
53
PATH_TO_REPO = Path(__file__).parent.parent.resolve()
PATH_TO_TRANFORMERS = PATH_TO_REPO / "src/transformers"
PATH_TO_TESTS = PATH_TO_REPO / "tests"

# List here the models to always test.
IMPORTANT_MODELS = [
54
    "auto",
Sylvain Gugger's avatar
Sylvain Gugger committed
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
    # Most downloaded models
    "bert",
    "clip",
    "t5",
    "xlm-roberta",
    "gpt2",
    "bart",
    "mpnet",
    "gpt-j",
    "wav2vec2",
    "deberta-v2",
    "layoutlm",
    "opt",
    "longformer",
    "vit",
    # Pipeline-specific model (to be sure each pipeline has one model in this list)
    "tapas",
    "vilt",
    "clap",
    "detr",
    "owlvit",
    "dpt",
    "videomae",
]
79

80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

@contextmanager
def checkout_commit(repo, commit_id):
    """
    Context manager that checks out a commit in the repo.
    """
    current_head = repo.head.commit if repo.head.is_detached else repo.head.ref

    try:
        repo.git.checkout(commit_id)
        yield

    finally:
        repo.git.checkout(current_head)


def clean_code(content):
    """
    Remove docstrings, empty line or comments from `content`.
    """
    # fmt: off
    # Remove docstrings by splitting on triple " then triple ':
    splits = content.split('\"\"\"')
    content = "".join(splits[::2])
    splits = content.split("\'\'\'")
    # fmt: on
    content = "".join(splits[::2])

    # Remove empty lines and comments
    lines_to_keep = []
    for line in content.split("\n"):
        # remove anything that is after a # sign.
        line = re.sub("#.*$", "", line)
        if len(line) == 0 or line.isspace():
            continue
        lines_to_keep.append(line)
    return "\n".join(lines_to_keep)


Yih-Dar's avatar
Yih-Dar committed
119
120
121
122
123
124
125
126
127
128
def get_all_tests():
    """
    Return a list of paths to all test folders and files under `tests`. All paths are rooted at `tests`.

    - folders under `tests`: `tokenization`, `pipelines`, etc. The folder `models` is excluded.
    - folders under `tests/models`: `bert`, `gpt2`, etc.
    - test files under `tests`: `test_modeling_common.py`, `test_tokenization_common.py`, etc.
    """

    # test folders/files directly under `tests` folder
Sylvain Gugger's avatar
Sylvain Gugger committed
129
130
131
    tests = os.listdir(PATH_TO_TESTS)
    tests = [f"tests/{f}" for f in tests if "__pycache__" not in f]
    tests = sorted([f for f in tests if (PATH_TO_REPO / f).is_dir() or f.startswith("tests/test_")])
Yih-Dar's avatar
Yih-Dar committed
132
133

    # model specific test folders
Sylvain Gugger's avatar
Sylvain Gugger committed
134
135
136
    model_test_folders = os.listdir(PATH_TO_TESTS / "models")
    model_test_folders = [f"tests/models/{f}" for f in model_test_folders if "__pycache__" not in f]
    model_test_folders = sorted([f for f in model_test_folders if (PATH_TO_REPO / f).is_dir()])
Yih-Dar's avatar
Yih-Dar committed
137
138

    tests.remove("tests/models")
Sylvain Gugger's avatar
Sylvain Gugger committed
139
140
141
    # Sagemaker tests are not meant to be run on the CI.
    if "tests/sagemaker" in tests:
        tests.remove("tests/sagemaker")
Yih-Dar's avatar
Yih-Dar committed
142
143
144
145
146
    tests = model_test_folders + tests

    return tests


147
148
149
150
def diff_is_docstring_only(repo, branching_point, filename):
    """
    Check if the diff is only in docstrings in a filename.
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
151
    folder = Path(repo.working_dir)
152
    with checkout_commit(repo, branching_point):
Sylvain Gugger's avatar
Sylvain Gugger committed
153
        with open(folder / filename, "r", encoding="utf-8") as f:
154
155
            old_content = f.read()

Sylvain Gugger's avatar
Sylvain Gugger committed
156
    with open(folder / filename, "r", encoding="utf-8") as f:
157
158
159
160
161
162
163
164
        new_content = f.read()

    old_content_clean = clean_code(old_content)
    new_content_clean = clean_code(new_content)

    return old_content_clean == new_content_clean


165
166
167
168
def get_diff(repo, base_commit, commits):
    """
    Get's the diff between one or several commits and the head of the repository.
    """
169
170
    print("\n### DIFF ###\n")
    code_diff = []
171
172
    for commit in commits:
        for diff_obj in commit.diff(base_commit):
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
            # We always add new python files
            if diff_obj.change_type == "A" and diff_obj.b_path.endswith(".py"):
                code_diff.append(diff_obj.b_path)
            # We check that deleted python files won't break corresponding tests.
            elif diff_obj.change_type == "D" and diff_obj.a_path.endswith(".py"):
                code_diff.append(diff_obj.a_path)
            # Now for modified files
            elif diff_obj.change_type in ["M", "R"] and diff_obj.b_path.endswith(".py"):
                # In case of renames, we'll look at the tests using both the old and new name.
                if diff_obj.a_path != diff_obj.b_path:
                    code_diff.extend([diff_obj.a_path, diff_obj.b_path])
                else:
                    # Otherwise, we check modifications are in code and not docstrings.
                    if diff_is_docstring_only(repo, commit, diff_obj.b_path):
                        print(f"Ignoring diff in {diff_obj.b_path} as it only concerns docstrings or comments.")
                    else:
                        code_diff.append(diff_obj.a_path)

    return code_diff


Sylvain Gugger's avatar
Sylvain Gugger committed
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def get_modified_python_files(diff_with_last_commit=False):
    """
    Return a list of python files that have been modified between:

    - the current head and the main branch if `diff_with_last_commit=False` (default)
    - the current head and its parent commit otherwise.
    """
    repo = Repo(PATH_TO_REPO)

    if not diff_with_last_commit:
        print(f"main is at {repo.refs.main.commit}")
        print(f"Current head is at {repo.head.commit}")

        branching_commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in branching_commits:
            print(f"Branching commit: {commit}")
        return get_diff(repo, repo.head.commit, branching_commits)
    else:
        print(f"main is at {repo.head.commit}")
        parent_commits = repo.head.commit.parents
        for commit in parent_commits:
            print(f"Parent commit: {commit}")
        return get_diff(repo, repo.head.commit, parent_commits)


219
220
221
222
223
224
225
226
227
def get_diff_for_py_and_mdx_files(repo, base_commit, commits):
    """
    Get's the diff between one or several commits and the head of the repository.
    """
    print("\n### DIFF ###\n")
    code_diff = []
    for commit in commits:
        for diff_obj in commit.diff(base_commit):
            # We always add new python files
228
            if diff_obj.change_type in ["A"] and (diff_obj.b_path.endswith(".py") or diff_obj.b_path.endswith(".mdx")):
229
                code_diff.append(diff_obj.b_path)
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
            # Now for modified files
            elif (
                diff_obj.change_type in ["M", "R"]
                and diff_obj.b_path.endswith(".py")
                or diff_obj.b_path.endswith(".mdx")
            ):
                # In case of renames, we'll look at the tests using both the old and new name.
                if diff_obj.a_path != diff_obj.b_path:
                    code_diff.extend([diff_obj.a_path, diff_obj.b_path])
                else:
                    # Otherwise, we check modifications are in code and not docstrings.
                    if diff_is_docstring_only(repo, commit, diff_obj.b_path):
                        print(f"Ignoring diff in {diff_obj.b_path} as it only concerns docstrings or comments.")
                    else:
                        code_diff.append(diff_obj.a_path)
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291

    return code_diff


def get_modified_python_and_mdx_files(diff_with_last_commit=False):
    """
    Return a list of python and mdx files that have been modified between:

    - the current head and the main branch if `diff_with_last_commit=False` (default)
    - the current head and its parent commit otherwise.
    """
    repo = Repo(PATH_TO_REPO)

    if not diff_with_last_commit:
        print(f"main is at {repo.refs.main.commit}")
        print(f"Current head is at {repo.head.commit}")

        branching_commits = repo.merge_base(repo.refs.main, repo.head)
        for commit in branching_commits:
            print(f"Branching commit: {commit}")
        return get_diff_for_py_and_mdx_files(repo, repo.head.commit, branching_commits)
    else:
        print(f"main is at {repo.head.commit}")
        parent_commits = repo.head.commit.parents
        for commit in parent_commits:
            print(f"Parent commit: {commit}")
        return get_diff_for_py_and_mdx_files(repo, repo.head.commit, parent_commits)


def get_doctest_files(diff_with_last_commit=False):
    """
    Return a list of python and mdx files that have been modified between:

    - the current head and the main branch if `diff_with_last_commit=False` (default)
    - the current head and its parent commit otherwise.
    """
    test_files_to_run = get_modified_python_and_mdx_files(diff_with_last_commit)
    with open("utils/documentation_tests.txt") as fp:
        documentation_tests = set(fp.read().strip().split("\n"))
    # So far we don't have 100% coverage for doctest. This line will be removed once we achieve 100%.
    test_files_to_run = [x for x in test_files_to_run if x in documentation_tests]
    # Make sure we did not end up with a test file that was removed
    test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()]

    return test_files_to_run


Sylvain Gugger's avatar
Sylvain Gugger committed
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+(\.+\S+)\s+import\s+([^\n]+) -> Line only contains from .xxx import yyy and we catch .xxx and yyy
# (?=\n) -> Look-ahead to a new line. We can't just put \n here or using find_all on this re will only catch every
#           other import.
_re_single_line_relative_imports = re.compile(r"(?:^|\n)\s*from\s+(\.+\S+)\s+import\s+([^\n]+)(?=\n)")
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+(\.+\S+)\s+import\s+\(([^\)]+)\) -> Line continues with from .xxx import (yyy) and we catch .xxx and yyy
# yyy will take multiple lines otherwise there wouldn't be parenthesis.
_re_multi_line_relative_imports = re.compile(r"(?:^|\n)\s*from\s+(\.+\S+)\s+import\s+\(([^\)]+)\)")
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+transformers(\S*)\s+import\s+([^\n]+) -> Line only contains from transformers.xxx import yyy and we catch
#           .xxx and yyy
# (?=\n) -> Look-ahead to a new line. We can't just put \n here or using find_all on this re will only catch every
#           other import.
_re_single_line_direct_imports = re.compile(r"(?:^|\n)\s*from\s+transformers(\S*)\s+import\s+([^\n]+)(?=\n)")
# (:?^|\n) -> Non-catching group for the beginning of the doc or a new line.
# \s*from\s+transformers(\S*)\s+import\s+\(([^\)]+)\) -> Line continues with from transformers.xxx import (yyy) and we
# catch .xxx and yyy. yyy will take multiple lines otherwise there wouldn't be parenthesis.
_re_multi_line_direct_imports = re.compile(r"(?:^|\n)\s*from\s+transformers(\S*)\s+import\s+\(([^\)]+)\)")


def extract_imports(module_fname, cache=None):
314
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
315
316
    Get the imports a given module makes. This takes a module filename and returns the list of module filenames
    imported in the module with the objects imported in that module filename.
317
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
318
319
320
321
    if cache is not None and module_fname in cache:
        return cache[module_fname]

    with open(PATH_TO_REPO / module_fname, "r", encoding="utf-8") as f:
322
323
        content = f.read()

Sylvain Gugger's avatar
Sylvain Gugger committed
324
325
326
327
328
    # Filter out all docstrings to not get imports in code examples.
    splits = content.split('"""')
    content = "".join(splits[::2])

    module_parts = str(module_fname).split(os.path.sep)
329
330
331
    imported_modules = []

    # Let's start with relative imports
Sylvain Gugger's avatar
Sylvain Gugger committed
332
333
334
335
336
337
338
339
    relative_imports = _re_single_line_relative_imports.findall(content)
    relative_imports = [
        (mod, imp) for mod, imp in relative_imports if "# tests_ignore" not in imp and imp.strip() != "("
    ]
    multiline_relative_imports = _re_multi_line_relative_imports.findall(content)
    relative_imports += [(mod, imp) for mod, imp in multiline_relative_imports if "# tests_ignore" not in imp]

    for module, imports in relative_imports:
340
        level = 0
Sylvain Gugger's avatar
Sylvain Gugger committed
341
342
        while module.startswith("."):
            module = module[1:]
343
344
            level += 1

Sylvain Gugger's avatar
Sylvain Gugger committed
345
346
        if len(module) > 0:
            dep_parts = module_parts[: len(module_parts) - level] + module.split(".")
347
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
348
            dep_parts = module_parts[: len(module_parts) - level]
349
        imported_module = os.path.sep.join(dep_parts)
Sylvain Gugger's avatar
Sylvain Gugger committed
350
        imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")]))
351
352

    # Let's continue with direct imports
Sylvain Gugger's avatar
Sylvain Gugger committed
353
354
355
356
    direct_imports = _re_single_line_direct_imports.findall(content)
    direct_imports = [(mod, imp) for mod, imp in direct_imports if "# tests_ignore" not in imp and imp.strip() != "("]
    multiline_direct_imports = _re_multi_line_direct_imports.findall(content)
    direct_imports += [(mod, imp) for mod, imp in multiline_direct_imports if "# tests_ignore" not in imp]
357

Sylvain Gugger's avatar
Sylvain Gugger committed
358
359
360
361
362
    for module, imports in direct_imports:
        import_parts = module.split(".")[1:]  # ignore the first .
        dep_parts = ["src", "transformers"] + import_parts
        imported_module = os.path.sep.join(dep_parts)
        imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")]))
363

Sylvain Gugger's avatar
Sylvain Gugger committed
364
365
366
367
368
369
370
371
372
    result = []
    for module_file, imports in imported_modules:
        if (PATH_TO_REPO / f"{module_file}.py").is_file():
            module_file = f"{module_file}.py"
        elif (PATH_TO_REPO / module_file).is_dir() and (PATH_TO_REPO / module_file / "__init__.py").is_file():
            module_file = os.path.sep.join([module_file, "__init__.py"])
        imports = [imp for imp in imports if len(imp) > 0 and re.match("^[A-Za-z0-9_]*$", imp)]
        if len(imports) > 0:
            result.append((module_file, imports))
373

Sylvain Gugger's avatar
Sylvain Gugger committed
374
375
    if cache is not None:
        cache[module_fname] = result
Sylvain Gugger's avatar
Sylvain Gugger committed
376

Sylvain Gugger's avatar
Sylvain Gugger committed
377
    return result
378

379

Sylvain Gugger's avatar
Sylvain Gugger committed
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
def get_module_dependencies(module_fname, cache=None):
    """
    Get the dependencies of a module from the module filename as a list of module filenames. This will resolve any
    __init__ we pass: if we import from a submodule utils, the dependencies will be utils/foo.py and utils/bar.py (if
    the objects imported actually come from utils.foo and utils.bar) not utils/__init__.py.
    """
    dependencies = []
    imported_modules = extract_imports(module_fname, cache=cache)
    # The while loop is to recursively traverse all inits we may encounter.
    while len(imported_modules) > 0:
        new_modules = []
        for module, imports in imported_modules:
            # If we end up in an __init__ we are often not actually importing from this init (except in the case where
            # the object is fully defined in the __init__)
            if module.endswith("__init__.py"):
                # So we get the imports from that init then try to find where our objects come from.
                new_imported_modules = extract_imports(module, cache=cache)
                for new_module, new_imports in new_imported_modules:
                    if any([i in new_imports for i in imports]):
                        if new_module not in dependencies:
                            new_modules.append((new_module, [i for i in new_imports if i in imports]))
                        imports = [i for i in imports if i not in new_imports]
                if len(imports) > 0:
                    # If there are any objects lefts, they may be a submodule
                    path_to_module = PATH_TO_REPO / module.replace("__init__.py", "")
                    dependencies.extend(
                        [
                            os.path.join(module.replace("__init__.py", ""), f"{i}.py")
                            for i in imports
                            if (path_to_module / f"{i}.py").is_file()
                        ]
                    )
                    imports = [i for i in imports if not (path_to_module / f"{i}.py").is_file()]
                    if len(imports) > 0:
                        # Then if there are still objects left, they are fully defined in the init, so we keep it as a
                        # dependency.
                        dependencies.append(module)
            else:
                dependencies.append(module)
419

Sylvain Gugger's avatar
Sylvain Gugger committed
420
421
        imported_modules = new_modules
    return dependencies
Sylvain Gugger's avatar
Sylvain Gugger committed
422
423


424
425
426
427
def create_reverse_dependency_tree():
    """
    Create a list of all edges (a, b) which mean that modifying a impacts b with a going over all module and test files.
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
428
429
430
431
    cache = {}
    all_modules = list(PATH_TO_TRANFORMERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py"))
    all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules]
    edges = [(dep, mod) for mod in all_modules for dep in get_module_dependencies(mod, cache=cache)]
432

Sylvain Gugger's avatar
Sylvain Gugger committed
433
    return list(set(edges))
434
435
436
437
438
439
440
441


def get_tree_starting_at(module, edges):
    """
    Returns the tree starting at a given module following all edges in the following format: [module, [list of edges
    starting at module], [list of edges starting at the preceding level], ...]
    """
    vertices_seen = [module]
Sylvain Gugger's avatar
Sylvain Gugger committed
442
    new_edges = [edge for edge in edges if edge[0] == module and edge[1] != module and "__init__.py" not in edge[1]]
443
444
445
    tree = [module]
    while len(new_edges) > 0:
        tree.append(new_edges)
446
        final_vertices = list({edge[1] for edge in new_edges})
447
        vertices_seen.extend(final_vertices)
Sylvain Gugger's avatar
Sylvain Gugger committed
448
449
450
451
452
        new_edges = [
            edge
            for edge in edges
            if edge[0] in final_vertices and edge[1] not in vertices_seen and "__init__.py" not in edge[1]
        ]
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469

    return tree


def print_tree_deps_of(module, all_edges=None):
    """
    Prints the tree of modules depending on a given module.
    """
    if all_edges is None:
        all_edges = create_reverse_dependency_tree()
    tree = get_tree_starting_at(module, all_edges)

    # The list of lines is a list of tuples (line_to_be_printed, module)
    # Keeping the modules lets us know where to insert each new lines in the list.
    lines = [(tree[0], tree[0])]
    for index in range(1, len(tree)):
        edges = tree[index]
470
        start_edges = {edge[0] for edge in edges}
471
472

        for start in start_edges:
473
            end_edges = {edge[1] for edge in edges if edge[0] == start}
474
475
476
477
478
479
480
481
482
483
484
            # We will insert all those edges just after the line showing start.
            pos = 0
            while lines[pos][1] != start:
                pos += 1
            lines = lines[: pos + 1] + [(" " * (2 * index) + end, end) for end in end_edges] + lines[pos + 1 :]

    for line in lines:
        # We don't print the refs that where just here to help build lines.
        print(line[0])


485
486
def create_reverse_dependency_map():
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
487
488
    Create the dependency map from module/test filename to the list of modules/tests that depend on it (even
    recursively).
489
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
490
491
492
493
    cache = {}
    all_modules = list(PATH_TO_TRANFORMERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py"))
    all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules]
    direct_deps = {m: get_module_dependencies(m, cache=cache) for m in all_modules}
Sylvain Gugger's avatar
Sylvain Gugger committed
494

495
496
497
498
    # This recurses the dependencies
    something_changed = True
    while something_changed:
        something_changed = False
Sylvain Gugger's avatar
Sylvain Gugger committed
499
        for m in all_modules:
500
            for d in direct_deps[m]:
Sylvain Gugger's avatar
Sylvain Gugger committed
501
502
                if d.endswith("__init__.py"):
                    continue
503
504
                if d not in direct_deps:
                    raise ValueError(f"KeyError:{d}. From {m}")
Sylvain Gugger's avatar
Sylvain Gugger committed
505
506
507
508
                new_deps = set(direct_deps[d]) - set(direct_deps[m])
                if len(new_deps) > 0:
                    direct_deps[m].extend(list(new_deps))
                    something_changed = True
509
510
511

    # Finally we can build the reverse map.
    reverse_map = collections.defaultdict(list)
Sylvain Gugger's avatar
Sylvain Gugger committed
512
    for m in all_modules:
513
514
515
        for d in direct_deps[m]:
            reverse_map[d].append(m)

Sylvain Gugger's avatar
Sylvain Gugger committed
516
517
518
519
520
    for m in [f for f in all_modules if f.endswith("__init__.py")]:
        direct_deps = get_module_dependencies(m, cache=cache)
        deps = sum([reverse_map[d] for d in direct_deps if not d.endswith("__init__.py")], direct_deps)
        reverse_map[m] = list(set(deps) - {m})

521
522
523
    return reverse_map


Sylvain Gugger's avatar
Sylvain Gugger committed
524
525
526
527
528
529
530
def create_module_to_test_map(reverse_map=None, filter_models=False):
    """
    Extract the tests from the reverse_dependency_map and potentially filters the model tests.
    """
    if reverse_map is None:
        reverse_map = create_reverse_dependency_map()
    test_map = {module: [f for f in deps if f.startswith("tests")] for module, deps in reverse_map.items()}
531

Sylvain Gugger's avatar
Sylvain Gugger committed
532
533
    if not filter_models:
        return test_map
534

Sylvain Gugger's avatar
Sylvain Gugger committed
535
536
537
538
539
540
541
542
543
544
    num_model_tests = len(list(PATH_TO_TESTS.glob("models/*")))

    def has_many_models(tests):
        model_tests = {Path(t).parts[2] for t in tests if t.startswith("tests/models/")}
        return len(model_tests) > num_model_tests // 2

    def filter_tests(tests):
        return [t for t in tests if not t.startswith("tests/models/") or Path(t).parts[2] in IMPORTANT_MODELS]

    return {module: (filter_tests(tests) if has_many_models(tests) else tests) for module, tests in test_map.items()}
545
546


Sylvain Gugger's avatar
Sylvain Gugger committed
547
def check_imports_all_exist():
548
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
549
550
    Isn't used per se by the test fetcher but might be used later as a quality check. Putting this here for now so the
    code is not lost.
551
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
552
553
554
555
    cache = {}
    all_modules = list(PATH_TO_TRANFORMERS.glob("**/*.py")) + list(PATH_TO_TESTS.glob("**/*.py"))
    all_modules = [str(mod.relative_to(PATH_TO_REPO)) for mod in all_modules]
    direct_deps = {m: get_module_dependencies(m, cache=cache) for m in all_modules}
556

Sylvain Gugger's avatar
Sylvain Gugger committed
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
    for module, deps in direct_deps.items():
        for dep in deps:
            if not (PATH_TO_REPO / dep).is_file():
                print(f"{module} has dependency on {dep} which does not exist.")


def _print_list(l):
    return "\n".join([f"- {f}" for f in l])


def create_json_map(test_files_to_run, json_output_file):
    if json_output_file is None:
        return

    test_map = {}
    for test_file in test_files_to_run:
        # `test_file` is a path to a test folder/file, starting with `tests/`. For example,
        #   - `tests/models/bert/test_modeling_bert.py` or `tests/models/bert`
        #   - `tests/trainer/test_trainer.py` or `tests/trainer`
        #   - `tests/test_modeling_common.py`
        names = test_file.split(os.path.sep)
        if names[1] == "models":
            # take the part like `models/bert` for modeling tests
            key = os.path.sep.join(names[1:3])
        elif len(names) > 2 or not test_file.endswith(".py"):
            # test folders under `tests` or python files under them
            # take the part like tokenization, `pipeline`, etc. for other test categories
            key = os.path.sep.join(names[1:2])
585
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
586
587
            # common test files directly under `tests/`
            key = "common"
588

Sylvain Gugger's avatar
Sylvain Gugger committed
589
590
591
        if key not in test_map:
            test_map[key] = []
        test_map[key].append(test_file)
592

Sylvain Gugger's avatar
Sylvain Gugger committed
593
594
595
596
597
    # sort the keys & values
    keys = sorted(test_map.keys())
    test_map = {k: " ".join(sorted(test_map[k])) for k in keys}
    with open(json_output_file, "w", encoding="UTF-8") as fp:
        json.dump(test_map, fp, ensure_ascii=False)
598
599


Sylvain Gugger's avatar
Sylvain Gugger committed
600
601
602
def infer_tests_to_run(
    output_file, diff_with_last_commit=False, filters=None, filter_models=True, json_output_file=None
):
603
    modified_files = get_modified_python_files(diff_with_last_commit=diff_with_last_commit)
604
605
606
    print(f"\n### MODIFIED FILES ###\n{_print_list(modified_files)}")

    # Create the map that will give us all impacted modules.
Sylvain Gugger's avatar
Sylvain Gugger committed
607
    reverse_map = create_reverse_dependency_map()
608
609
    impacted_files = modified_files.copy()
    for f in modified_files:
Sylvain Gugger's avatar
Sylvain Gugger committed
610
611
        if f in reverse_map:
            impacted_files.extend(reverse_map[f])
612
613

    # Remove duplicates
614
    impacted_files = sorted(set(impacted_files))
615
616
617
    print(f"\n### IMPACTED FILES ###\n{_print_list(impacted_files)}")

    # Grab the corresponding test files:
Sylvain Gugger's avatar
Sylvain Gugger committed
618
    if "setup.py" in modified_files:
619
        test_files_to_run = ["tests"]
Sylvain Gugger's avatar
Sylvain Gugger committed
620
        repo_utils_launch = True
621
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
622
623
624
625
626
627
628
629
630
        # All modified tests need to be run.
        test_files_to_run = [
            f for f in modified_files if f.startswith("tests") and f.split(os.path.sep)[-1].startswith("test")
        ]
        # Then we grab the corresponding test files.
        test_map = create_module_to_test_map(reverse_map=reverse_map, filter_models=filter_models)
        for f in modified_files:
            if f in test_map:
                test_files_to_run.extend(test_map[f])
631
        test_files_to_run = sorted(set(test_files_to_run))
Sylvain Gugger's avatar
Sylvain Gugger committed
632
633
        # Remove SageMaker tests
        test_files_to_run = [f for f in test_files_to_run if not f.split(os.path.sep)[1] == "sagemaker"]
634
        # Make sure we did not end up with a test file that was removed
Sylvain Gugger's avatar
Sylvain Gugger committed
635
        test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()]
636
        if filters is not None:
637
            filtered_files = []
Sylvain Gugger's avatar
Sylvain Gugger committed
638
639
            for _filter in filters:
                filtered_files.extend([f for f in test_files_to_run if f.startswith(_filter)])
640
            test_files_to_run = filtered_files
Sylvain Gugger's avatar
Sylvain Gugger committed
641
642

        repo_utils_launch = any(f.split(os.path.sep)[1] == "repo_utils" for f in modified_files)
Sylvain Gugger's avatar
Sylvain Gugger committed
643
644
645
646
647

    if repo_utils_launch:
        repo_util_file = Path(output_file).parent / "test_repo_utils.txt"
        with open(repo_util_file, "w", encoding="utf-8") as f:
            f.write("tests/repo_utils")
648

649
650
651
652
653
    print(f"\n### TEST TO RUN ###\n{_print_list(test_files_to_run)}")
    if len(test_files_to_run) > 0:
        with open(output_file, "w", encoding="utf-8") as f:
            f.write(" ".join(test_files_to_run))

Yih-Dar's avatar
Yih-Dar committed
654
655
656
657
658
659
660
        # Create a map that maps test categories to test files, i.e. `models/bert` -> [...test_modeling_bert.py, ...]

        # Get all test directories (and some common test files) under `tests` and `tests/models` if `test_files_to_run`
        # contains `tests` (i.e. when `setup.py` is changed).
        if "tests" in test_files_to_run:
            test_files_to_run = get_all_tests()

Sylvain Gugger's avatar
Sylvain Gugger committed
661
        create_json_map(test_files_to_run, json_output_file)
Yih-Dar's avatar
Yih-Dar committed
662

663

Sylvain Gugger's avatar
Sylvain Gugger committed
664
665
666
667
668
669
670
671
def filter_tests(output_file, filters):
    """
    Reads the content of the output file and filters out all the tests in a list of given folders.

    Args:
        output_file (`str` or `os.PathLike`): The path to the output file of the tests fetcher.
        filters (`List[str]`): A list of folders to filter.
    """
672
673
674
675
676
677
    if not os.path.isfile(output_file):
        print("No test file found.")
        return
    with open(output_file, "r", encoding="utf-8") as f:
        test_files = f.read().split(" ")

Sylvain Gugger's avatar
Sylvain Gugger committed
678
    if len(test_files) == 0 or test_files == [""]:
679
680
        print("No tests to filter.")
        return
Sylvain Gugger's avatar
Sylvain Gugger committed
681

682
    if test_files == ["tests"]:
Sylvain Gugger's avatar
Sylvain Gugger committed
683
        test_files = [os.path.join("tests", f) for f in os.listdir("tests") if f not in ["__init__.py"] + filters]
684
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
685
        test_files = [f for f in test_files if f.split(os.path.sep)[1] not in filters]
686
687
688
689
690

    with open(output_file, "w", encoding="utf-8") as f:
        f.write(" ".join(test_files))


Sylvain Gugger's avatar
Sylvain Gugger committed
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
def parse_commit_message(commit_message):
    """
    Parses the commit message to detect if a command is there to skip, force all or part of the CI.

    Returns a dictionary of strings to bools with keys skip, test_all_models and test_all.
    """
    if commit_message is None:
        return {"skip": False, "no_filter": False, "test_all": False}

    command_search = re.search(r"\[([^\]]*)\]", commit_message)
    if command_search is not None:
        command = command_search.groups()[0]
        command = command.lower().replace("-", " ").replace("_", " ")
        skip = command in ["ci skip", "skip ci", "circleci skip", "skip circleci"]
        no_filter = set(command.split(" ")) == {"no", "filter"}
        test_all = set(command.split(" ")) == {"test", "all"}
        return {"skip": skip, "no_filter": no_filter, "test_all": test_all}
    else:
        return {"skip": False, "no_filter": False, "test_all": False}


712
713
714
715
716
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output_file", type=str, default="test_list.txt", help="Where to store the list of tests to run"
    )
Yih-Dar's avatar
Yih-Dar committed
717
718
719
720
721
722
    parser.add_argument(
        "--json_output_file",
        type=str,
        default="test_map.json",
        help="Where to store the tests to run in a dictionary format mapping test categories to test files",
    )
723
724
725
726
727
728
    parser.add_argument(
        "--diff_with_last_commit",
        action="store_true",
        help="To fetch the tests between the current commit and the last commit",
    )
    parser.add_argument(
729
730
731
732
733
        "--filters",
        type=str,
        nargs="*",
        default=["tests"],
        help="Only keep the test files matching one of those filters.",
734
    )
735
    parser.add_argument(
Sylvain Gugger's avatar
Sylvain Gugger committed
736
        "--filter_tests",
737
        action="store_true",
Sylvain Gugger's avatar
Sylvain Gugger committed
738
        help="Will filter the pipeline/repo utils tests outside of the generated list of tests.",
739
    )
740
741
742
743
744
745
    parser.add_argument(
        "--print_dependencies_of",
        type=str,
        help="Will only print the tree of modules depending on the file passed.",
        default=None,
    )
Sylvain Gugger's avatar
Sylvain Gugger committed
746
747
748
749
750
751
    parser.add_argument(
        "--commit_message",
        type=str,
        help="The commit message (which could contain a command to force all tests or skip the CI).",
        default=None,
    )
752
    args = parser.parse_args()
753
754
    if args.print_dependencies_of is not None:
        print_tree_deps_of(args.print_dependencies_of)
Sylvain Gugger's avatar
Sylvain Gugger committed
755
756
    elif args.filter_tests:
        filter_tests(args.output_file, ["pipelines", "repo_utils"])
757
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
758
759
760
761
762
763
764
765
766
767
        repo = Repo(PATH_TO_REPO)
        commit_message = repo.head.commit.message
        commit_flags = parse_commit_message(commit_message)
        if commit_flags["skip"]:
            print("Force-skipping the CI")
            quit()
        if commit_flags["no_filter"]:
            print("Running all tests fetched without filtering.")
        if commit_flags["test_all"]:
            print("Force-launching all tests")
768
769

        diff_with_last_commit = args.diff_with_last_commit
770
771
        if not diff_with_last_commit and not repo.head.is_detached and repo.head.ref == repo.refs.main:
            print("main branch detected, fetching tests against last commit.")
772
773
            diff_with_last_commit = True

Sylvain Gugger's avatar
Sylvain Gugger committed
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
        if not commit_flags["test_all"]:
            try:
                infer_tests_to_run(
                    args.output_file,
                    diff_with_last_commit=diff_with_last_commit,
                    filters=args.filters,
                    json_output_file=args.json_output_file,
                    filter_models=not commit_flags["no_filter"],
                )
                filter_tests(args.output_file, ["repo_utils"])
            except Exception as e:
                print(f"\nError when trying to grab the relevant tests: {e}\n\nRunning all tests.")
                commit_flags["test_all"] = True

        if commit_flags["test_all"]:
789
            with open(args.output_file, "w", encoding="utf-8") as f:
790
791
792
793
                if args.filters is None:
                    f.write("./tests/")
                else:
                    f.write(" ".join(args.filters))
Sylvain Gugger's avatar
Sylvain Gugger committed
794
795
796

            test_files_to_run = get_all_tests()
            create_json_map(test_files_to_run, args.json_output_file)