tests_fetcher.py 31.5 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
229
            if diff_obj.change_type in ["A", "M", "R"] and (
                diff_obj.b_path.endswith(".py") or diff_obj.b_path.endswith(".mdx")
230
            ):
231
                code_diff.append(diff_obj.b_path)
232
233
234
235
236
237
238
239
240
241
242
243
244
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

    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
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
# (:?^|\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):
301
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
302
303
    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.
304
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
305
306
307
308
    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:
309
310
        content = f.read()

Sylvain Gugger's avatar
Sylvain Gugger committed
311
312
313
314
315
    # 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)
316
317
318
    imported_modules = []

    # Let's start with relative imports
Sylvain Gugger's avatar
Sylvain Gugger committed
319
320
321
322
323
324
325
326
    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:
327
        level = 0
Sylvain Gugger's avatar
Sylvain Gugger committed
328
329
        while module.startswith("."):
            module = module[1:]
330
331
            level += 1

Sylvain Gugger's avatar
Sylvain Gugger committed
332
333
        if len(module) > 0:
            dep_parts = module_parts[: len(module_parts) - level] + module.split(".")
334
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
335
            dep_parts = module_parts[: len(module_parts) - level]
336
        imported_module = os.path.sep.join(dep_parts)
Sylvain Gugger's avatar
Sylvain Gugger committed
337
        imported_modules.append((imported_module, [imp.strip() for imp in imports.split(",")]))
338
339

    # Let's continue with direct imports
Sylvain Gugger's avatar
Sylvain Gugger committed
340
341
342
343
    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]
344

Sylvain Gugger's avatar
Sylvain Gugger committed
345
346
347
348
349
    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(",")]))
350

Sylvain Gugger's avatar
Sylvain Gugger committed
351
352
353
354
355
356
357
358
359
    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))
360

Sylvain Gugger's avatar
Sylvain Gugger committed
361
362
    if cache is not None:
        cache[module_fname] = result
Sylvain Gugger's avatar
Sylvain Gugger committed
363

Sylvain Gugger's avatar
Sylvain Gugger committed
364
    return result
365

366

Sylvain Gugger's avatar
Sylvain Gugger committed
367
368
369
370
371
372
373
374
375
376
377
378
379
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
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)
406

Sylvain Gugger's avatar
Sylvain Gugger committed
407
408
        imported_modules = new_modules
    return dependencies
Sylvain Gugger's avatar
Sylvain Gugger committed
409
410


411
412
413
414
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
415
416
417
418
    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)]
419

Sylvain Gugger's avatar
Sylvain Gugger committed
420
    return list(set(edges))
421
422
423
424
425
426
427
428


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
429
    new_edges = [edge for edge in edges if edge[0] == module and edge[1] != module and "__init__.py" not in edge[1]]
430
431
432
    tree = [module]
    while len(new_edges) > 0:
        tree.append(new_edges)
433
        final_vertices = list({edge[1] for edge in new_edges})
434
        vertices_seen.extend(final_vertices)
Sylvain Gugger's avatar
Sylvain Gugger committed
435
436
437
438
439
        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]
        ]
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456

    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]
457
        start_edges = {edge[0] for edge in edges}
458
459

        for start in start_edges:
460
            end_edges = {edge[1] for edge in edges if edge[0] == start}
461
462
463
464
465
466
467
468
469
470
471
            # 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])


472
473
def create_reverse_dependency_map():
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
474
475
    Create the dependency map from module/test filename to the list of modules/tests that depend on it (even
    recursively).
476
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
477
478
479
480
    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
481

482
483
484
485
    # This recurses the dependencies
    something_changed = True
    while something_changed:
        something_changed = False
Sylvain Gugger's avatar
Sylvain Gugger committed
486
        for m in all_modules:
487
            for d in direct_deps[m]:
Sylvain Gugger's avatar
Sylvain Gugger committed
488
489
                if d.endswith("__init__.py"):
                    continue
490
491
                if d not in direct_deps:
                    raise ValueError(f"KeyError:{d}. From {m}")
Sylvain Gugger's avatar
Sylvain Gugger committed
492
493
494
495
                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
496
497
498

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

Sylvain Gugger's avatar
Sylvain Gugger committed
503
504
505
506
507
    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})

508
509
510
    return reverse_map


Sylvain Gugger's avatar
Sylvain Gugger committed
511
512
513
514
515
516
517
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()}
518

Sylvain Gugger's avatar
Sylvain Gugger committed
519
520
    if not filter_models:
        return test_map
521

Sylvain Gugger's avatar
Sylvain Gugger committed
522
523
524
525
526
527
528
529
530
531
    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()}
532
533


Sylvain Gugger's avatar
Sylvain Gugger committed
534
def check_imports_all_exist():
535
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
536
537
    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.
538
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
539
540
541
542
    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}
543

Sylvain Gugger's avatar
Sylvain Gugger committed
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
    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])
572
        else:
Sylvain Gugger's avatar
Sylvain Gugger committed
573
574
            # common test files directly under `tests/`
            key = "common"
575

Sylvain Gugger's avatar
Sylvain Gugger committed
576
577
578
        if key not in test_map:
            test_map[key] = []
        test_map[key].append(test_file)
579

Sylvain Gugger's avatar
Sylvain Gugger committed
580
581
582
583
584
    # 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)
585
586


Sylvain Gugger's avatar
Sylvain Gugger committed
587
588
589
def infer_tests_to_run(
    output_file, diff_with_last_commit=False, filters=None, filter_models=True, json_output_file=None
):
590
    modified_files = get_modified_python_files(diff_with_last_commit=diff_with_last_commit)
591
592
593
    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
594
    reverse_map = create_reverse_dependency_map()
595
596
    impacted_files = modified_files.copy()
    for f in modified_files:
Sylvain Gugger's avatar
Sylvain Gugger committed
597
598
        if f in reverse_map:
            impacted_files.extend(reverse_map[f])
599
600

    # Remove duplicates
601
    impacted_files = sorted(set(impacted_files))
602
603
604
    print(f"\n### IMPACTED FILES ###\n{_print_list(impacted_files)}")

    # Grab the corresponding test files:
Sylvain Gugger's avatar
Sylvain Gugger committed
605
    if "setup.py" in modified_files:
606
        test_files_to_run = ["tests"]
Sylvain Gugger's avatar
Sylvain Gugger committed
607
        repo_utils_launch = True
608
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
609
610
611
612
613
614
615
616
617
        # 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])
618
        test_files_to_run = sorted(set(test_files_to_run))
Sylvain Gugger's avatar
Sylvain Gugger committed
619
620
        # Remove SageMaker tests
        test_files_to_run = [f for f in test_files_to_run if not f.split(os.path.sep)[1] == "sagemaker"]
621
        # Make sure we did not end up with a test file that was removed
Sylvain Gugger's avatar
Sylvain Gugger committed
622
        test_files_to_run = [f for f in test_files_to_run if (PATH_TO_REPO / f).exists()]
623
        if filters is not None:
624
            filtered_files = []
Sylvain Gugger's avatar
Sylvain Gugger committed
625
626
            for _filter in filters:
                filtered_files.extend([f for f in test_files_to_run if f.startswith(_filter)])
627
            test_files_to_run = filtered_files
Sylvain Gugger's avatar
Sylvain Gugger committed
628
629

        repo_utils_launch = any(f.split(os.path.sep)[1] == "repo_utils" for f in modified_files)
Sylvain Gugger's avatar
Sylvain Gugger committed
630
631
632
633
634

    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")
635

636
637
638
639
640
    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
641
642
643
644
645
646
647
        # 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
648
        create_json_map(test_files_to_run, json_output_file)
Yih-Dar's avatar
Yih-Dar committed
649

650

Sylvain Gugger's avatar
Sylvain Gugger committed
651
652
653
654
655
656
657
658
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.
    """
659
660
661
662
663
664
    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
665
    if len(test_files) == 0 or test_files == [""]:
666
667
        print("No tests to filter.")
        return
Sylvain Gugger's avatar
Sylvain Gugger committed
668

669
    if test_files == ["tests"]:
Sylvain Gugger's avatar
Sylvain Gugger committed
670
        test_files = [os.path.join("tests", f) for f in os.listdir("tests") if f not in ["__init__.py"] + filters]
671
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
672
        test_files = [f for f in test_files if f.split(os.path.sep)[1] not in filters]
673
674
675
676
677

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


Sylvain Gugger's avatar
Sylvain Gugger committed
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
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}


699
700
701
702
703
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
704
705
706
707
708
709
    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",
    )
710
711
712
713
714
715
    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(
716
717
718
719
720
        "--filters",
        type=str,
        nargs="*",
        default=["tests"],
        help="Only keep the test files matching one of those filters.",
721
    )
722
    parser.add_argument(
Sylvain Gugger's avatar
Sylvain Gugger committed
723
        "--filter_tests",
724
        action="store_true",
Sylvain Gugger's avatar
Sylvain Gugger committed
725
        help="Will filter the pipeline/repo utils tests outside of the generated list of tests.",
726
    )
727
728
729
730
731
732
    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
733
734
735
736
737
738
    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,
    )
739
    args = parser.parse_args()
740
741
    if args.print_dependencies_of is not None:
        print_tree_deps_of(args.print_dependencies_of)
Sylvain Gugger's avatar
Sylvain Gugger committed
742
743
    elif args.filter_tests:
        filter_tests(args.output_file, ["pipelines", "repo_utils"])
744
    else:
Sylvain Gugger's avatar
Sylvain Gugger committed
745
746
747
748
749
750
751
752
753
754
        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")
755
756

        diff_with_last_commit = args.diff_with_last_commit
757
758
        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.")
759
760
            diff_with_last_commit = True

Sylvain Gugger's avatar
Sylvain Gugger committed
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
        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"]:
776
            with open(args.output_file, "w", encoding="utf-8") as f:
777
778
779
780
                if args.filters is None:
                    f.write("./tests/")
                else:
                    f.write(" ".join(args.filters))
Sylvain Gugger's avatar
Sylvain Gugger committed
781
782
783

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