testing_utils.py 81.4 KB
Newer Older
Sylvain Gugger's avatar
Sylvain Gugger committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# 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.

NielsRogge's avatar
NielsRogge committed
15
import collections
16
import contextlib
17
import doctest
18
import functools
19
import importlib
20
import inspect
21
import logging
22
import multiprocessing
23
import os
24
import re
25
import shlex
26
import shutil
Zachary Mueller's avatar
Zachary Mueller committed
27
import subprocess
28
import sys
29
import tempfile
30
import time
Aymeric Augustin's avatar
Aymeric Augustin committed
31
import unittest
32
from collections import defaultdict
33
from collections.abc import Mapping
34
from functools import wraps
35
from io import StringIO
36
from pathlib import Path
37
from typing import Callable, Dict, Iterable, Iterator, List, Optional, Union
38
from unittest import mock
39
from unittest.mock import patch
40

41
import urllib3
42

43
44
from transformers import logging as transformers_logging

45
from .integrations import (
46
    is_clearml_available,
47
48
49
    is_optuna_available,
    is_ray_available,
    is_sigopt_available,
50
    is_tensorboard_available,
51
52
    is_wandb_available,
)
53
from .integrations.deepspeed import is_deepspeed_available
54
from .utils import (
55
    is_accelerate_available,
56
    is_apex_available,
57
    is_aqlm_available,
58
    is_auto_awq_available,
Marc Sun's avatar
Marc Sun committed
59
    is_auto_gptq_available,
60
    is_bitsandbytes_available,
NielsRogge's avatar
NielsRogge committed
61
    is_bs4_available,
NielsRogge's avatar
NielsRogge committed
62
    is_cv2_available,
63
    is_cython_available,
64
    is_decord_available,
65
    is_detectron2_available,
Susnato Dhar's avatar
Susnato Dhar committed
66
    is_essentia_available,
67
    is_faiss_available,
68
    is_flash_attn_2_available,
69
    is_flax_available,
70
    is_fsdp_available,
71
    is_ftfy_available,
72
    is_g2p_en_available,
73
    is_ipex_available,
74
    is_jieba_available,
75
    is_jinja_available,
76
    is_jumanpp_available,
Matt's avatar
Matt committed
77
    is_keras_nlp_available,
NielsRogge's avatar
NielsRogge committed
78
    is_levenshtein_available,
79
    is_librosa_available,
80
    is_natten_available,
NielsRogge's avatar
NielsRogge committed
81
    is_nltk_available,
82
    is_onnx_available,
83
    is_optimum_available,
84
    is_pandas_available,
85
    is_peft_available,
86
    is_phonemizer_available,
Susnato Dhar's avatar
Susnato Dhar committed
87
    is_pretty_midi_available,
88
    is_pyctcdecode_available,
89
    is_pytesseract_available,
90
    is_pytest_available,
91
    is_pytorch_quantization_available,
92
    is_quanto_available,
yujun's avatar
yujun committed
93
    is_rjieba_available,
94
    is_sacremoses_available,
95
    is_safetensors_available,
96
    is_scipy_available,
97
    is_sentencepiece_available,
98
    is_seqio_available,
Patrick von Platen's avatar
Patrick von Platen committed
99
    is_soundfile_availble,
100
    is_spacy_available,
101
    is_sudachi_available,
102
    is_sudachi_projection_available,
Kamal Raj's avatar
Kamal Raj committed
103
    is_tensorflow_probability_available,
104
    is_tensorflow_text_available,
105
    is_tf2onnx_available,
106
    is_tf_available,
NielsRogge's avatar
NielsRogge committed
107
    is_timm_available,
108
109
    is_tokenizers_available,
    is_torch_available,
110
    is_torch_bf16_available_on_device,
111
112
    is_torch_bf16_cpu_available,
    is_torch_bf16_gpu_available,
113
    is_torch_fp16_available_on_device,
114
    is_torch_neuroncore_available,
115
    is_torch_npu_available,
116
    is_torch_sdpa_available,
117
    is_torch_tensorrt_fx_available,
118
    is_torch_tf32_available,
119
    is_torch_xla_available,
120
    is_torch_xpu_available,
Suraj Patil's avatar
Suraj Patil committed
121
    is_torchaudio_available,
122
    is_torchdynamo_available,
NielsRogge's avatar
NielsRogge committed
123
    is_torchvision_available,
124
    is_vision_available,
125
    strtobool,
126
)
127
128


129
130
131
132
if is_accelerate_available():
    from accelerate.state import AcceleratorState, PartialState


133
134
135
136
137
138
139
140
141
142
143
if is_pytest_available():
    from _pytest.doctest import (
        Module,
        _get_checker,
        _get_continue_on_failure,
        _get_runner,
        _is_mocked,
        _patch_unwrap_mock_aware,
        get_optionflags,
    )
    from _pytest.outcomes import skip
144
    from _pytest.pathlib import import_path
145
146
147
148
149
150
    from pytest import DoctestItem
else:
    Module = object
    DoctestItem = object


Julien Chaumond's avatar
Julien Chaumond committed
151
SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy"
152
DUMMY_UNKNOWN_IDENTIFIER = "julien-c/dummy-unknown"
153
DUMMY_DIFF_TOKENIZER_IDENTIFIER = "julien-c/dummy-diff-tokenizer"
Julien Chaumond's avatar
Julien Chaumond committed
154
# Used to test Auto{Config, Model, Tokenizer} model_type detection.
Julien Chaumond's avatar
Julien Chaumond committed
155

Sylvain Gugger's avatar
Sylvain Gugger committed
156
157
# Used to test the hub
USER = "__DUMMY_TRANSFORMERS_USER__"
158
159
160
161
ENDPOINT_STAGING = "https://hub-ci.huggingface.co"

# Not critical, only usable on the sandboxed CI instance.
TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL"
Sylvain Gugger's avatar
Sylvain Gugger committed
162

Julien Chaumond's avatar
Julien Chaumond committed
163

164
def parse_flag_from_env(key, default=False):
165
    try:
166
167
168
169
170
171
172
173
174
175
        value = os.environ[key]
    except KeyError:
        # KEY isn't set, default to `default`.
        _value = default
    else:
        # KEY is set, convert it to True or False.
        try:
            _value = strtobool(value)
        except ValueError:
            # More values are supported, but let's keep the message simple.
176
            raise ValueError(f"If set, {key} must be yes or no.")
177
178
    return _value

179

Julien Chaumond's avatar
Julien Chaumond committed
180
181
182
183
184
185
186
187
188
def parse_int_from_env(key, default=None):
    try:
        value = os.environ[key]
    except KeyError:
        _value = default
    else:
        try:
            _value = int(value)
        except ValueError:
189
            raise ValueError(f"If set, {key} must be a int.")
Julien Chaumond's avatar
Julien Chaumond committed
190
191
192
    return _value


193
_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
194
195
_run_pt_tf_cross_tests = parse_flag_from_env("RUN_PT_TF_CROSS_TESTS", default=True)
_run_pt_flax_cross_tests = parse_flag_from_env("RUN_PT_FLAX_CROSS_TESTS", default=True)
196
_run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False)
Sylvain Gugger's avatar
Sylvain Gugger committed
197
_run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False)
Julien Chaumond's avatar
Julien Chaumond committed
198
_tf_gpu_memory_limit = parse_int_from_env("TF_GPU_MEMORY_LIMIT", default=None)
199
_run_pipeline_tests = parse_flag_from_env("RUN_PIPELINE_TESTS", default=True)
Sylvain Gugger's avatar
Sylvain Gugger committed
200
_run_tool_tests = parse_flag_from_env("RUN_TOOL_TESTS", default=False)
201
_run_third_party_device_tests = parse_flag_from_env("RUN_THIRD_PARTY_DEVICE_TESTS", default=False)
202
203


204
205
206
207
208
209
210
211
def is_pt_tf_cross_test(test_case):
    """
    Decorator marking a test as a test that control interactions between PyTorch and TensorFlow.

    PT+TF tests are skipped by default and we can run only them by setting RUN_PT_TF_CROSS_TESTS environment variable
    to a truthy value and selecting the is_pt_tf_cross_test pytest mark.

    """
212
    if not _run_pt_tf_cross_tests or not is_torch_available() or not is_tf_available():
213
214
215
216
217
218
219
220
221
222
        return unittest.skip("test is PT+TF test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:
            return pytest.mark.is_pt_tf_cross_test()(test_case)


223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
def is_pt_flax_cross_test(test_case):
    """
    Decorator marking a test as a test that control interactions between PyTorch and Flax

    PT+FLAX tests are skipped by default and we can run only them by setting RUN_PT_FLAX_CROSS_TESTS environment
    variable to a truthy value and selecting the is_pt_flax_cross_test pytest mark.

    """
    if not _run_pt_flax_cross_tests or not is_torch_available() or not is_flax_available():
        return unittest.skip("test is PT+FLAX test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:
            return pytest.mark.is_pt_flax_cross_test()(test_case)


Sylvain Gugger's avatar
Sylvain Gugger committed
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
def is_staging_test(test_case):
    """
    Decorator marking a test as a staging test.

    Those tests will run using the staging environment of huggingface.co instead of the real model hub.
    """
    if not _run_staging:
        return unittest.skip("test is staging test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:
            return pytest.mark.is_staging_test()(test_case)


259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def is_pipeline_test(test_case):
    """
    Decorator marking a test as a pipeline test. If RUN_PIPELINE_TESTS is set to a falsy value, those tests will be
    skipped.
    """
    if not _run_pipeline_tests:
        return unittest.skip("test is pipeline test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:
            return pytest.mark.is_pipeline_test()(test_case)


Sylvain Gugger's avatar
Sylvain Gugger committed
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
def is_tool_test(test_case):
    """
    Decorator marking a test as a tool test. If RUN_TOOL_TESTS is set to a falsy value, those tests will be skipped.
    """
    if not _run_tool_tests:
        return unittest.skip("test is a tool test")(test_case)
    else:
        try:
            import pytest  # We don't need a hard dependency on pytest in the main library
        except ImportError:
            return test_case
        else:
            return pytest.mark.is_tool_test()(test_case)


290
291
292
293
def slow(test_case):
    """
    Decorator marking a test as slow.

Sylvain Gugger's avatar
Sylvain Gugger committed
294
    Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.
295
296

    """
297
    return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)
298
299


Lysandre Debut's avatar
Lysandre Debut committed
300
301
302
303
304
305
306
307
308
309
310
def tooslow(test_case):
    """
    Decorator marking a test as too slow.

    Slow tests are skipped while they're in the process of being fixed. No test should stay tagged as "tooslow" as
    these will not be tested by the CI.

    """
    return unittest.skip("test is too slow")(test_case)


311
312
313
314
def custom_tokenizers(test_case):
    """
    Decorator marking a test for a custom tokenizer.

Sylvain Gugger's avatar
Sylvain Gugger committed
315
316
    Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS
    environment variable to a truthy value to run them.
317
    """
318
    return unittest.skipUnless(_run_custom_tokenizers, "test of custom tokenizers")(test_case)
319
320


NielsRogge's avatar
NielsRogge committed
321
322
323
324
325
326
327
def require_bs4(test_case):
    """
    Decorator marking a test that requires BeautifulSoup4. These tests are skipped when BeautifulSoup4 isn't installed.
    """
    return unittest.skipUnless(is_bs4_available(), "test requires BeautifulSoup4")(test_case)


NielsRogge's avatar
NielsRogge committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
def require_cv2(test_case):
    """
    Decorator marking a test that requires OpenCV.

    These tests are skipped when OpenCV isn't installed.

    """
    return unittest.skipUnless(is_cv2_available(), "test requires OpenCV")(test_case)


def require_levenshtein(test_case):
    """
    Decorator marking a test that requires Levenshtein.

    These tests are skipped when Levenshtein isn't installed.

    """
    return unittest.skipUnless(is_levenshtein_available(), "test requires Levenshtein")(test_case)


def require_nltk(test_case):
    """
    Decorator marking a test that requires NLTK.

    These tests are skipped when NLTK isn't installed.

    """
    return unittest.skipUnless(is_nltk_available(), "test requires NLTK")(test_case)


358
359
360
361
362
363
364
def require_accelerate(test_case):
    """
    Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed.
    """
    return unittest.skipUnless(is_accelerate_available(), "test requires accelerate")(test_case)


365
366
367
368
369
370
371
372
373
def require_fsdp(test_case, min_version: str = "1.12.0"):
    """
    Decorator marking a test that requires fsdp. These tests are skipped when fsdp isn't installed.
    """
    return unittest.skipUnless(is_fsdp_available(min_version), f"test requires torch version >= {min_version}")(
        test_case
    )


374
375
376
377
378
379
380
def require_g2p_en(test_case):
    """
    Decorator marking a test that requires g2p_en. These tests are skipped when SentencePiece isn't installed.
    """
    return unittest.skipUnless(is_g2p_en_available(), "test requires g2p_en")(test_case)


381
382
383
384
385
386
387
def require_safetensors(test_case):
    """
    Decorator marking a test that requires safetensors. These tests are skipped when safetensors isn't installed.
    """
    return unittest.skipUnless(is_safetensors_available(), "test requires safetensors")(test_case)


yujun's avatar
yujun committed
388
389
390
391
def require_rjieba(test_case):
    """
    Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed.
    """
392
    return unittest.skipUnless(is_rjieba_available(), "test requires rjieba")(test_case)
yujun's avatar
yujun committed
393
394


395
396
397
398
399
400
401
def require_jieba(test_case):
    """
    Decorator marking a test that requires jieba. These tests are skipped when jieba isn't installed.
    """
    return unittest.skipUnless(is_jieba_available(), "test requires jieba")(test_case)


402
403
404
405
406
407
408
def require_jinja(test_case):
    """
    Decorator marking a test that requires jinja. These tests are skipped when jinja isn't installed.
    """
    return unittest.skipUnless(is_jinja_available(), "test requires jinja")(test_case)


409
def require_tf2onnx(test_case):
410
    return unittest.skipUnless(is_tf2onnx_available(), "test requires tf2onnx")(test_case)
411
412


413
def require_onnx(test_case):
414
    return unittest.skipUnless(is_onnx_available(), "test requires ONNX")(test_case)
415
416


NielsRogge's avatar
NielsRogge committed
417
418
419
420
421
422
423
def require_timm(test_case):
    """
    Decorator marking a test that requires Timm.

    These tests are skipped when Timm isn't installed.

    """
424
    return unittest.skipUnless(is_timm_available(), "test requires Timm")(test_case)
NielsRogge's avatar
NielsRogge committed
425
426


427
428
429
430
431
432
433
434
435
436
def require_natten(test_case):
    """
    Decorator marking a test that requires NATTEN.

    These tests are skipped when NATTEN isn't installed.

    """
    return unittest.skipUnless(is_natten_available(), "test requires natten")(test_case)


437
438
439
440
441
442
443
def require_torch(test_case):
    """
    Decorator marking a test that requires PyTorch.

    These tests are skipped when PyTorch isn't installed.

    """
444
    return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case)
445
446


447
448
449
450
451
452
453
def require_flash_attn(test_case):
    """
    Decorator marking a test that requires Flash Attention.

    These tests are skipped when Flash Attention isn't installed.

    """
454
    return unittest.skipUnless(is_flash_attn_2_available(), "test requires Flash Attention")(test_case)
455
456


457
458
459
460
461
462
463
464
465
def require_torch_sdpa(test_case):
    """
    Decorator marking a test that requires PyTorch's SDPA.

    These tests are skipped when requirements are not met (torch version).
    """
    return unittest.skipUnless(is_torch_sdpa_available(), "test requires PyTorch SDPA")(test_case)


466
467
468
469
def require_read_token(fn):
    """
    A decorator that loads the HF token for tests that require to load gated models.
    """
470
    token = os.getenv("HF_HUB_READ_TOKEN")
471
472
473

    @wraps(fn)
    def _inner(*args, **kwargs):
474
        with patch("huggingface_hub.utils._headers.get_token", return_value=token):
475
476
477
478
479
            return fn(*args, **kwargs)

    return _inner


480
481
482
483
484
485
486
487
def require_peft(test_case):
    """
    Decorator marking a test that requires PEFT.

    These tests are skipped when PEFT isn't installed.

    """
    return unittest.skipUnless(is_peft_available(), "test requires PEFT")(test_case)
488
489


NielsRogge's avatar
NielsRogge committed
490
491
492
493
494
495
496
497
498
499
def require_torchvision(test_case):
    """
    Decorator marking a test that requires Torchvision.

    These tests are skipped when Torchvision isn't installed.

    """
    return unittest.skipUnless(is_torchvision_available(), "test requires Torchvision")(test_case)


500
501
502
503
504
505
506
507
508
509
510
511
def require_torch_or_tf(test_case):
    """
    Decorator marking a test that requires PyTorch or TensorFlow.

    These tests are skipped when neither PyTorch not TensorFlow is installed.

    """
    return unittest.skipUnless(is_torch_available() or is_tf_available(), "test requires PyTorch or TensorFlow")(
        test_case
    )


512
513
514
515
def require_intel_extension_for_pytorch(test_case):
    """
    Decorator marking a test that requires Intel Extension for PyTorch.

516
517
    These tests are skipped when Intel Extension for PyTorch isn't installed or it does not match current PyTorch
    version.
518
519

    """
520
521
522
523
524
    return unittest.skipUnless(
        is_ipex_available(),
        "test requires Intel Extension for PyTorch to be installed and match current PyTorch version, see"
        " https://github.com/intel/intel-extension-for-pytorch",
    )(test_case)
525
526


Kamal Raj's avatar
Kamal Raj committed
527
528
529
530
531
532
533
def require_tensorflow_probability(test_case):
    """
    Decorator marking a test that requires TensorFlow probability.

    These tests are skipped when TensorFlow probability isn't installed.

    """
534
535
536
    return unittest.skipUnless(is_tensorflow_probability_available(), "test requires TensorFlow probability")(
        test_case
    )
Kamal Raj's avatar
Kamal Raj committed
537
538


Suraj Patil's avatar
Suraj Patil committed
539
540
def require_torchaudio(test_case):
    """
541
    Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed.
Suraj Patil's avatar
Suraj Patil committed
542
    """
543
    return unittest.skipUnless(is_torchaudio_available(), "test requires torchaudio")(test_case)
544
545


546
547
def require_tf(test_case):
    """
548
    Decorator marking a test that requires TensorFlow. These tests are skipped when TensorFlow isn't installed.
549
    """
550
    return unittest.skipUnless(is_tf_available(), "test requires TensorFlow")(test_case)
551
552


553
554
def require_flax(test_case):
    """
555
    Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed
556
    """
557
    return unittest.skipUnless(is_flax_available(), "test requires JAX & Flax")(test_case)
558
559


560
561
def require_sentencepiece(test_case):
    """
562
    Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
563
    """
564
    return unittest.skipUnless(is_sentencepiece_available(), "test requires SentencePiece")(test_case)
565
566


567
568
569
570
571
572
573
def require_sacremoses(test_case):
    """
    Decorator marking a test that requires Sacremoses. These tests are skipped when Sacremoses isn't installed.
    """
    return unittest.skipUnless(is_sacremoses_available(), "test requires Sacremoses")(test_case)


574
575
576
577
578
579
580
def require_seqio(test_case):
    """
    Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
    """
    return unittest.skipUnless(is_seqio_available(), "test requires Seqio")(test_case)


581
582
583
584
def require_scipy(test_case):
    """
    Decorator marking a test that requires Scipy. These tests are skipped when SentencePiece isn't installed.
    """
585
    return unittest.skipUnless(is_scipy_available(), "test requires Scipy")(test_case)
586
587


588
589
def require_tokenizers(test_case):
    """
590
    Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed.
591
    """
592
    return unittest.skipUnless(is_tokenizers_available(), "test requires tokenizers")(test_case)
593
594


595
596
597
598
599
600
601
602
def require_tensorflow_text(test_case):
    """
    Decorator marking a test that requires tensorflow_text. These tests are skipped when tensroflow_text isn't
    installed.
    """
    return unittest.skipUnless(is_tensorflow_text_available(), "test requires tensorflow_text")(test_case)


Matt's avatar
Matt committed
603
604
605
606
607
608
609
def require_keras_nlp(test_case):
    """
    Decorator marking a test that requires keras_nlp. These tests are skipped when keras_nlp isn't installed.
    """
    return unittest.skipUnless(is_keras_nlp_available(), "test requires keras_nlp")(test_case)


NielsRogge's avatar
NielsRogge committed
610
611
612
613
def require_pandas(test_case):
    """
    Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed.
    """
614
    return unittest.skipUnless(is_pandas_available(), "test requires pandas")(test_case)
NielsRogge's avatar
NielsRogge committed
615
616


617
618
619
620
def require_pytesseract(test_case):
    """
    Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed.
    """
621
    return unittest.skipUnless(is_pytesseract_available(), "test requires PyTesseract")(test_case)
622
623


624
625
626
627
628
def require_pytorch_quantization(test_case):
    """
    Decorator marking a test that requires PyTorch Quantization Toolkit. These tests are skipped when PyTorch
    Quantization Toolkit isn't installed.
    """
629
630
631
    return unittest.skipUnless(is_pytorch_quantization_available(), "test requires PyTorch Quantization Toolkit")(
        test_case
    )
632
633


634
def require_vision(test_case):
635
    """
636
637
638
    Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't
    installed.
    """
639
    return unittest.skipUnless(is_vision_available(), "test requires vision")(test_case)
640

641

642
643
644
645
def require_ftfy(test_case):
    """
    Decorator marking a test that requires ftfy. These tests are skipped when ftfy isn't installed.
    """
646
    return unittest.skipUnless(is_ftfy_available(), "test requires ftfy")(test_case)
647
648
649
650
651
652


def require_spacy(test_case):
    """
    Decorator marking a test that requires SpaCy. These tests are skipped when SpaCy isn't installed.
    """
653
    return unittest.skipUnless(is_spacy_available(), "test requires spacy")(test_case)
654
655


656
657
658
659
660
661
662
def require_decord(test_case):
    """
    Decorator marking a test that requires decord. These tests are skipped when decord isn't installed.
    """
    return unittest.skipUnless(is_decord_available(), "test requires decord")(test_case)


663
664
665
666
def require_torch_multi_gpu(test_case):
    """
    Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without
    multiple GPUs.
667

668
    To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu"
669
    """
670
    if not is_torch_available():
671
672
673
674
        return unittest.skip("test requires PyTorch")(test_case)

    import torch

675
    return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)
676
677


678
679
680
681
682
683
684
685
686
687
688
689
690
691
def require_torch_multi_accelerator(test_case):
    """
    Decorator marking a test that requires a multi-accelerator (in PyTorch). These tests are skipped on a machine
    without multiple accelerators. To run *only* the multi_accelerator tests, assuming all test names contain
    multi_accelerator: $ pytest -sv ./tests -k "multi_accelerator"
    """
    if not is_torch_available():
        return unittest.skip("test requires PyTorch")(test_case)

    return unittest.skipUnless(backend_device_count(torch_device) > 1, "test requires multiple accelerators")(
        test_case
    )


692
def require_torch_non_multi_gpu(test_case):
693
694
695
    """
    Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch).
    """
696
    if not is_torch_available():
697
698
699
700
        return unittest.skip("test requires PyTorch")(test_case)

    import torch

701
    return unittest.skipUnless(torch.cuda.device_count() < 2, "test requires 0 or 1 GPU")(test_case)
702
703


704
705
706
707
708
709
710
711
712
713
def require_torch_non_multi_accelerator(test_case):
    """
    Decorator marking a test that requires 0 or 1 accelerator setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip("test requires PyTorch")(test_case)

    return unittest.skipUnless(backend_device_count(torch_device) < 2, "test requires 0 or 1 accelerator")(test_case)


714
715
716
717
718
719
720
721
722
def require_torch_up_to_2_gpus(test_case):
    """
    Decorator marking a test that requires 0 or 1 or 2 GPU setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip("test requires PyTorch")(test_case)

    import torch

723
    return unittest.skipUnless(torch.cuda.device_count() < 3, "test requires 0 or 1 or 2 GPUs")(test_case)
724
725


726
727
728
729
730
731
732
733
734
735
736
def require_torch_up_to_2_accelerators(test_case):
    """
    Decorator marking a test that requires 0 or 1 or 2 accelerator setup (in PyTorch).
    """
    if not is_torch_available():
        return unittest.skip("test requires PyTorch")(test_case)

    return unittest.skipUnless(backend_device_count(torch_device) < 3, "test requires 0 or 1 or 2 accelerators")
    (test_case)


737
def require_torch_xla(test_case):
Lysandre Debut's avatar
Lysandre Debut committed
738
    """
739
    Decorator marking a test that requires TorchXLA (in PyTorch).
Lysandre Debut's avatar
Lysandre Debut committed
740
    """
741
    return unittest.skipUnless(is_torch_xla_available(), "test requires TorchXLA")(test_case)
Lysandre Debut's avatar
Lysandre Debut committed
742
743


744
745
746
747
748
749
750
751
752
def require_torch_neuroncore(test_case):
    """
    Decorator marking a test that requires NeuronCore (in PyTorch).
    """
    return unittest.skipUnless(is_torch_neuroncore_available(check_device=False), "test requires PyTorch NeuronCore")(
        test_case
    )


753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
def require_torch_npu(test_case):
    """
    Decorator marking a test that requires NPU (in PyTorch).
    """
    return unittest.skipUnless(is_torch_npu_available(), "test requires PyTorch NPU")(test_case)


def require_torch_multi_npu(test_case):
    """
    Decorator marking a test that requires a multi-NPU setup (in PyTorch). These tests are skipped on a machine without
    multiple NPUs.

    To run *only* the multi_npu tests, assuming all test names contain multi_npu: $ pytest -sv ./tests -k "multi_npu"
    """
    if not is_torch_npu_available():
        return unittest.skip("test requires PyTorch NPU")(test_case)

    return unittest.skipUnless(torch.npu.device_count() > 1, "test requires multiple NPUs")(test_case)


773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
def require_torch_xpu(test_case):
    """
    Decorator marking a test that requires XPU and IPEX.

    These tests are skipped when Intel Extension for PyTorch isn't installed or it does not match current PyTorch
    version.
    """
    return unittest.skipUnless(is_torch_xpu_available(), "test requires IPEX and an XPU device")(test_case)


def require_torch_multi_xpu(test_case):
    """
    Decorator marking a test that requires a multi-XPU setup with IPEX and atleast one XPU device. These tests are
    skipped on a machine without IPEX or multiple XPUs.

    To run *only* the multi_xpu tests, assuming all test names contain multi_xpu: $ pytest -sv ./tests -k "multi_xpu"
    """
    if not is_torch_xpu_available():
        return unittest.skip("test requires IPEX and atleast one XPU device")(test_case)

    return unittest.skipUnless(torch.xpu.device_count() > 1, "test requires multiple XPUs")(test_case)


796
if is_torch_available():
Stas Bekman's avatar
Stas Bekman committed
797
798
799
    # Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode
    import torch

800
801
802
803
804
805
806
807
808
809
    if "TRANSFORMERS_TEST_BACKEND" in os.environ:
        backend = os.environ["TRANSFORMERS_TEST_BACKEND"]
        try:
            _ = importlib.import_module(backend)
        except ModuleNotFoundError as e:
            raise ModuleNotFoundError(
                f"Failed to import `TRANSFORMERS_TEST_BACKEND` '{backend}'! This should be the name of an installed module. The original error (look up to see its"
                f" traceback):\n{e}"
            ) from e

810
811
    if "TRANSFORMERS_TEST_DEVICE" in os.environ:
        torch_device = os.environ["TRANSFORMERS_TEST_DEVICE"]
812
813
814
815
816
817
818
819
820
821
822
823
824
        if torch_device == "cuda" and not torch.cuda.is_available():
            raise ValueError(
                f"TRANSFORMERS_TEST_DEVICE={torch_device}, but CUDA is unavailable. Please double-check your testing environment."
            )
        if torch_device == "xpu" and not is_torch_xpu_available():
            raise ValueError(
                f"TRANSFORMERS_TEST_DEVICE={torch_device}, but XPU is unavailable. Please double-check your testing environment."
            )
        if torch_device == "npu" and not is_torch_npu_available():
            raise ValueError(
                f"TRANSFORMERS_TEST_DEVICE={torch_device}, but NPU is unavailable. Please double-check your testing environment."
            )

825
826
827
828
829
830
831
832
        try:
            # try creating device to see if provided device is valid
            _ = torch.device(torch_device)
        except RuntimeError as e:
            raise RuntimeError(
                f"Unknown testing device specified by environment variable `TRANSFORMERS_TEST_DEVICE`: {torch_device}"
            ) from e
    elif torch.cuda.is_available():
833
834
835
        torch_device = "cuda"
    elif _run_third_party_device_tests and is_torch_npu_available():
        torch_device = "npu"
836
837
    elif _run_third_party_device_tests and is_torch_xpu_available():
        torch_device = "xpu"
838
839
    else:
        torch_device = "cpu"
840
841
else:
    torch_device = None
842

843
844
845
if is_tf_available():
    import tensorflow as tf

846
847
848
849
850
851
852
if is_flax_available():
    import jax

    jax_device = jax.default_backend()
else:
    jax_device = None

853

854
855
856
857
858
def require_torchdynamo(test_case):
    """Decorator marking a test that requires TorchDynamo"""
    return unittest.skipUnless(is_torchdynamo_available(), "test requires TorchDynamo")(test_case)


859
860
861
862
863
def require_torch_tensorrt_fx(test_case):
    """Decorator marking a test that requires Torch-TensorRT FX"""
    return unittest.skipUnless(is_torch_tensorrt_fx_available(), "test requires Torch-TensorRT FX")(test_case)


864
def require_torch_gpu(test_case):
Patrick von Platen's avatar
Patrick von Platen committed
865
    """Decorator marking a test that requires CUDA and PyTorch."""
866
    return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case)
867
868


869
870
def require_torch_accelerator(test_case):
    """Decorator marking a test that requires an accessible accelerator and PyTorch."""
871
872
873
    return unittest.skipUnless(torch_device is not None and torch_device != "cpu", "test requires accelerator")(
        test_case
    )
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889


def require_torch_fp16(test_case):
    """Decorator marking a test that requires a device that supports fp16"""
    return unittest.skipUnless(
        is_torch_fp16_available_on_device(torch_device), "test requires device with fp16 support"
    )(test_case)


def require_torch_bf16(test_case):
    """Decorator marking a test that requires a device that supports bf16"""
    return unittest.skipUnless(
        is_torch_bf16_available_on_device(torch_device), "test requires device with bf16 support"
    )(test_case)


890
891
def require_torch_bf16_gpu(test_case):
    """Decorator marking a test that requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0"""
892
    return unittest.skipUnless(
893
894
895
896
897
898
899
900
901
902
        is_torch_bf16_gpu_available(),
        "test requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0",
    )(test_case)


def require_torch_bf16_cpu(test_case):
    """Decorator marking a test that requires torch>=1.10, using CPU."""
    return unittest.skipUnless(
        is_torch_bf16_cpu_available(),
        "test requires torch>=1.10, using CPU",
903
    )(test_case)
904
905
906
907


def require_torch_tf32(test_case):
    """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7."""
908
909
910
    return unittest.skipUnless(
        is_torch_tf32_available(), "test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7"
    )(test_case)
911
912


913
914
def require_detectron2(test_case):
    """Decorator marking a test that requires detectron2."""
915
    return unittest.skipUnless(is_detectron2_available(), "test requires `detectron2`")(test_case)
916
917


Ola Piktus's avatar
Ola Piktus committed
918
919
def require_faiss(test_case):
    """Decorator marking a test that requires faiss."""
920
    return unittest.skipUnless(is_faiss_available(), "test requires `faiss`")(test_case)
Ola Piktus's avatar
Ola Piktus committed
921
922


923
924
925
926
927
928
929
def require_optuna(test_case):
    """
    Decorator marking a test that requires optuna.

    These tests are skipped when optuna isn't installed.

    """
930
    return unittest.skipUnless(is_optuna_available(), "test requires optuna")(test_case)
931
932
933
934
935
936
937
938
939


def require_ray(test_case):
    """
    Decorator marking a test that requires Ray/tune.

    These tests are skipped when Ray/tune isn't installed.

    """
940
    return unittest.skipUnless(is_ray_available(), "test requires Ray/tune")(test_case)
941
942


943
944
945
946
947
948
949
def require_sigopt(test_case):
    """
    Decorator marking a test that requires SigOpt.

    These tests are skipped when SigOpt isn't installed.

    """
950
    return unittest.skipUnless(is_sigopt_available(), "test requires SigOpt")(test_case)
951
952


953
954
955
956
957
958
959
def require_wandb(test_case):
    """
    Decorator marking a test that requires wandb.

    These tests are skipped when wandb isn't installed.

    """
960
    return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case)
961
962


963
964
965
966
967
968
969
970
971
972
def require_clearml(test_case):
    """
    Decorator marking a test requires clearml.

    These tests are skipped when clearml isn't installed.

    """
    return unittest.skipUnless(is_clearml_available(), "test requires clearml")(test_case)


Patrick von Platen's avatar
Patrick von Platen committed
973
974
975
976
977
978
979
def require_soundfile(test_case):
    """
    Decorator marking a test that requires soundfile

    These tests are skipped when soundfile isn't installed.

    """
980
    return unittest.skipUnless(is_soundfile_availble(), "test requires soundfile")(test_case)
Patrick von Platen's avatar
Patrick von Platen committed
981
982


983
984
985
986
def require_deepspeed(test_case):
    """
    Decorator marking a test that requires deepspeed
    """
987
    return unittest.skipUnless(is_deepspeed_available(), "test requires deepspeed")(test_case)
988
989


990
991
992
993
def require_apex(test_case):
    """
    Decorator marking a test that requires apex
    """
994
    return unittest.skipUnless(is_apex_available(), "test requires apex")(test_case)
995
996


997
998
999
1000
1001
1002
1003
def require_aqlm(test_case):
    """
    Decorator marking a test that requires aqlm
    """
    return unittest.skipUnless(is_aqlm_available(), "test requires aqlm")(test_case)


1004
1005
def require_bitsandbytes(test_case):
    """
1006
    Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed.
1007
    """
1008
1009
1010
1011
1012
1013
1014
1015
1016
    if is_bitsandbytes_available() and is_torch_available():
        try:
            import pytest

            return pytest.mark.bitsandbytes(test_case)
        except ImportError:
            return test_case
    else:
        return unittest.skip("test requires bitsandbytes and torch")(test_case)
1017
1018


1019
1020
1021
1022
1023
1024
1025
def require_optimum(test_case):
    """
    Decorator for optimum dependency
    """
    return unittest.skipUnless(is_optimum_available(), "test requires optimum")(test_case)


1026
1027
1028
1029
1030
1031
1032
def require_tensorboard(test_case):
    """
    Decorator for `tensorboard` dependency
    """
    return unittest.skipUnless(is_tensorboard_available(), "test requires tensorboard")


Marc Sun's avatar
Marc Sun committed
1033
1034
1035
1036
1037
1038
1039
def require_auto_gptq(test_case):
    """
    Decorator for auto_gptq dependency
    """
    return unittest.skipUnless(is_auto_gptq_available(), "test requires auto-gptq")(test_case)


1040
1041
1042
1043
1044
1045
1046
def require_auto_awq(test_case):
    """
    Decorator for auto_awq dependency
    """
    return unittest.skipUnless(is_auto_awq_available(), "test requires autoawq")(test_case)


1047
1048
1049
1050
1051
1052
1053
def require_quanto(test_case):
    """
    Decorator for quanto dependency
    """
    return unittest.skipUnless(is_quanto_available(), "test requires quanto")(test_case)


1054
1055
1056
1057
def require_phonemizer(test_case):
    """
    Decorator marking a test that requires phonemizer
    """
1058
    return unittest.skipUnless(is_phonemizer_available(), "test requires phonemizer")(test_case)
1059
1060


1061
1062
1063
1064
def require_pyctcdecode(test_case):
    """
    Decorator marking a test that requires pyctcdecode
    """
1065
    return unittest.skipUnless(is_pyctcdecode_available(), "test requires pyctcdecode")(test_case)
1066
1067
1068
1069
1070
1071


def require_librosa(test_case):
    """
    Decorator marking a test that requires librosa
    """
1072
    return unittest.skipUnless(is_librosa_available(), "test requires librosa")(test_case)
1073
1074


Susnato Dhar's avatar
Susnato Dhar committed
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
def require_essentia(test_case):
    """
    Decorator marking a test that requires essentia
    """
    return unittest.skipUnless(is_essentia_available(), "test requires essentia")(test_case)


def require_pretty_midi(test_case):
    """
    Decorator marking a test that requires pretty_midi
    """
    return unittest.skipUnless(is_pretty_midi_available(), "test requires pretty_midi")(test_case)


1089
1090
1091
1092
1093
1094
1095
1096
def cmd_exists(cmd):
    return shutil.which(cmd) is not None


def require_usr_bin_time(test_case):
    """
    Decorator marking a test that requires `/usr/bin/time`
    """
1097
    return unittest.skipUnless(cmd_exists("/usr/bin/time"), "test requires /usr/bin/time")(test_case)
1098
1099
1100
1101
1102
1103
1104
1105
1106


def require_sudachi(test_case):
    """
    Decorator marking a test that requires sudachi
    """
    return unittest.skipUnless(is_sudachi_available(), "test requires sudachi")(test_case)


1107
1108
1109
1110
1111
1112
1113
1114
1115
def require_sudachi_projection(test_case):
    """
    Decorator marking a test that requires sudachi_projection
    """
    return unittest.skipUnless(is_sudachi_projection_available(), "test requires sudachi which supports projection")(
        test_case
    )


1116
1117
1118
1119
1120
def require_jumanpp(test_case):
    """
    Decorator marking a test that requires jumanpp
    """
    return unittest.skipUnless(is_jumanpp_available(), "test requires jumanpp")(test_case)
1121
1122


1123
1124
1125
1126
1127
1128
1129
def require_cython(test_case):
    """
    Decorator marking a test that requires jumanpp
    """
    return unittest.skipUnless(is_cython_available(), "test requires cython")(test_case)


1130
1131
def get_gpu_count():
    """
Suraj Patil's avatar
Suraj Patil committed
1132
    Return the number of available gpus (regardless of whether torch, tf or jax is used)
1133
    """
1134
    if is_torch_available():
1135
1136
1137
        import torch

        return torch.cuda.device_count()
1138
    elif is_tf_available():
1139
1140
1141
        import tensorflow as tf

        return len(tf.config.list_physical_devices("GPU"))
Suraj Patil's avatar
Suraj Patil committed
1142
1143
1144
1145
    elif is_flax_available():
        import jax

        return jax.device_count()
1146
1147
1148
1149
    else:
        return 0


1150
def get_tests_dir(append_path=None):
1151
    """
1152
1153
1154
1155
    Args:
        append_path: optional path to append to the tests dir path

    Return:
Sylvain Gugger's avatar
Sylvain Gugger committed
1156
1157
        The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is
        joined after the `tests` dir the former is provided.
1158

1159
1160
1161
    """
    # this function caller's __file__
    caller__file__ = inspect.stack()[1][1]
1162
    tests_dir = os.path.abspath(os.path.dirname(caller__file__))
1163
1164
1165
1166

    while not tests_dir.endswith("tests"):
        tests_dir = os.path.dirname(tests_dir)

1167
1168
1169
1170
    if append_path:
        return os.path.join(tests_dir, append_path)
    else:
        return tests_dir
1171
1172


1173
1174
1175
1176
1177
#
# Helper functions for dealing with testing text outputs
# The original code came from:
# https://github.com/fastai/fastai/blob/master/tests/utils/text.py

1178

1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
# When any function contains print() calls that get overwritten, like progress bars,
# a special care needs to be applied, since under pytest -s captured output (capsys
# or contextlib.redirect_stdout) contains any temporary printed strings, followed by
# \r's. This helper function ensures that the buffer will contain the same output
# with and without -s in pytest, by turning:
# foo bar\r tar mar\r final message
# into:
# final message
# it can handle a single string or a multiline buffer
def apply_print_resets(buf):
    return re.sub(r"^.*\r", "", buf, 0, re.M)


def assert_screenout(out, what):
    out_pr = apply_print_resets(out).lower()
    match_str = out_pr.find(what.lower())
    assert match_str != -1, f"expecting to find {what} in output: f{out_pr}"


class CaptureStd:
Sylvain Gugger's avatar
Sylvain Gugger committed
1199
1200
    """
    Context manager to capture:
1201

1202
1203
        - stdout: replay it, clean it up and make it available via `obj.out`
        - stderr: replay it and make it available via `obj.err`
1204

1205
1206
1207
1208
1209
    Args:
        out (`bool`, *optional*, defaults to `True`): Whether to capture stdout or not.
        err (`bool`, *optional*, defaults to `True`): Whether to capture stderr or not.
        replay (`bool`, *optional*, defaults to `True`): Whether to replay or not.
            By default each captured stream gets replayed back on context's exit, so that one can see what the test was
Sylvain Gugger's avatar
Sylvain Gugger committed
1210
1211
            doing. If this is a not wanted behavior and the captured data shouldn't be replayed, pass `replay=False` to
            disable this feature.
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222

    Examples:

    ```python
    # to capture stdout only with auto-replay
    with CaptureStdout() as cs:
        print("Secret message")
    assert "message" in cs.out

    # to capture stderr only with auto-replay
    import sys
Sylvain Gugger's avatar
Sylvain Gugger committed
1223

1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
    with CaptureStderr() as cs:
        print("Warning: ", file=sys.stderr)
    assert "Warning" in cs.err

    # to capture both streams with auto-replay
    with CaptureStd() as cs:
        print("Secret message")
        print("Warning: ", file=sys.stderr)
    assert "message" in cs.out
    assert "Warning" in cs.err

    # to capture just one of the streams, and not the other, with auto-replay
    with CaptureStd(err=False) as cs:
        print("Secret message")
    assert "message" in cs.out
    # but best use the stream-specific subclasses

    # to capture without auto-replay
    with CaptureStd(replay=False) as cs:
        print("Secret message")
    assert "message" in cs.out
    ```"""
1246

1247
1248
1249
    def __init__(self, out=True, err=True, replay=True):
        self.replay = replay

1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
        if out:
            self.out_buf = StringIO()
            self.out = "error: CaptureStd context is unfinished yet, called too early"
        else:
            self.out_buf = None
            self.out = "not capturing stdout"

        if err:
            self.err_buf = StringIO()
            self.err = "error: CaptureStd context is unfinished yet, called too early"
        else:
            self.err_buf = None
            self.err = "not capturing stderr"

    def __enter__(self):
        if self.out_buf:
            self.out_old = sys.stdout
            sys.stdout = self.out_buf

        if self.err_buf:
            self.err_old = sys.stderr
            sys.stderr = self.err_buf

        return self

    def __exit__(self, *exc):
        if self.out_buf:
            sys.stdout = self.out_old
1278
1279
1280
1281
            captured = self.out_buf.getvalue()
            if self.replay:
                sys.stdout.write(captured)
            self.out = apply_print_resets(captured)
1282
1283
1284

        if self.err_buf:
            sys.stderr = self.err_old
1285
1286
1287
1288
            captured = self.err_buf.getvalue()
            if self.replay:
                sys.stderr.write(captured)
            self.err = captured
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305

    def __repr__(self):
        msg = ""
        if self.out_buf:
            msg += f"stdout: {self.out}\n"
        if self.err_buf:
            msg += f"stderr: {self.err}\n"
        return msg


# in tests it's the best to capture only the stream that's wanted, otherwise
# it's easy to miss things, so unless you need to capture both streams, use the
# subclasses below (less typing). Or alternatively, configure `CaptureStd` to
# disable the stream you don't need to test.


class CaptureStdout(CaptureStd):
Patrick von Platen's avatar
Patrick von Platen committed
1306
    """Same as CaptureStd but captures only stdout"""
1307

1308
1309
    def __init__(self, replay=True):
        super().__init__(err=False, replay=replay)
1310
1311
1312


class CaptureStderr(CaptureStd):
Patrick von Platen's avatar
Patrick von Platen committed
1313
    """Same as CaptureStd but captures only stderr"""
1314

1315
1316
    def __init__(self, replay=True):
        super().__init__(out=False, replay=replay)
1317
1318


1319
class CaptureLogger:
Sylvain Gugger's avatar
Sylvain Gugger committed
1320
1321
    """
    Context manager to capture `logging` streams
1322
1323

    Args:
1324
        logger: 'logging` logger object
1325

1326
    Returns:
1327
1328
        The captured output is available via `self.out`

1329
    Example:
1330

1331
1332
1333
    ```python
    >>> from transformers import logging
    >>> from transformers.testing_utils import CaptureLogger
1334

1335
1336
1337
1338
1339
    >>> msg = "Testing 1, 2, 3"
    >>> logging.set_verbosity_info()
    >>> logger = logging.get_logger("transformers.models.bart.tokenization_bart")
    >>> with CaptureLogger(logger) as cl:
    ...     logger.info(msg)
Sylvain Gugger's avatar
Sylvain Gugger committed
1340
    >>> assert cl.out, msg + "\n"
1341
    ```
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
    """

    def __init__(self, logger):
        self.logger = logger
        self.io = StringIO()
        self.sh = logging.StreamHandler(self.io)
        self.out = ""

    def __enter__(self):
        self.logger.addHandler(self.sh)
        return self

    def __exit__(self, *exc):
        self.logger.removeHandler(self.sh)
        self.out = self.io.getvalue()

    def __repr__(self):
        return f"captured: {self.out}\n"


1362
1363
1364
1365
1366
1367
@contextlib.contextmanager
def LoggingLevel(level):
    """
    This is a context manager to temporarily change transformers modules logging level to the desired value and have it
    restored to the original setting at the end of the scope.

1368
    Example:
1369

1370
1371
    ```python
    with LoggingLevel(logging.INFO):
1372
        AutoModel.from_pretrained("openai-community/gpt2")  # calls logger.info() several times
1373
    ```
1374
1375
1376
1377
1378
1379
1380
1381
1382
    """
    orig_level = transformers_logging.get_verbosity()
    try:
        transformers_logging.set_verbosity(level)
        yield
    finally:
        transformers_logging.set_verbosity(orig_level)


1383
1384
1385
1386
1387
1388
@contextlib.contextmanager
# adapted from https://stackoverflow.com/a/64789046/9201239
def ExtendSysPath(path: Union[str, os.PathLike]) -> Iterator[None]:
    """
    Temporary add given path to `sys.path`.

1389
    Usage :
1390

1391
    ```python
Sylvain Gugger's avatar
Sylvain Gugger committed
1392
1393
    with ExtendSysPath("/path/to/dir"):
        mymodule = importlib.import_module("mymodule")
1394
    ```
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
    """

    path = os.fspath(path)
    try:
        sys.path.insert(0, path)
        yield
    finally:
        sys.path.remove(path)


1405
class TestCasePlus(unittest.TestCase):
Sylvain Gugger's avatar
Sylvain Gugger committed
1406
    """
1407
    This class extends *unittest.TestCase* with additional features.
1408

1409
1410
1411
1412
1413
1414
    Feature 1: A set of fully resolved important file and dir path accessors.

    In tests often we need to know where things are relative to the current test file, and it's not trivial since the
    test could be invoked from more than one directory or could reside in sub-directories with different depths. This
    class solves this problem by sorting out all the basic paths and provides easy accessors to them:

1415
    - `pathlib` objects (all fully resolved):
1416

1417
1418
1419
1420
1421
1422
       - `test_file_path` - the current test file path (=`__file__`)
       - `test_file_dir` - the directory containing the current test file
       - `tests_dir` - the directory of the `tests` test suite
       - `examples_dir` - the directory of the `examples` test suite
       - `repo_root_dir` - the directory of the repository
       - `src_dir` - the directory of `src` (i.e. where the `transformers` sub-dir resides)
1423

1424
    - stringified paths---same as above but these return paths as strings, rather than `pathlib` objects:
1425

1426
1427
1428
1429
1430
1431
       - `test_file_path_str`
       - `test_file_dir_str`
       - `tests_dir_str`
       - `examples_dir_str`
       - `repo_root_dir_str`
       - `src_dir_str`
1432

1433
    Feature 2: Flexible auto-removable temporary dirs which are guaranteed to get removed at the end of test.
1434

1435
    1. Create a unique temporary dir:
1436

1437
1438
1439
1440
    ```python
    def test_whatever(self):
        tmp_dir = self.get_auto_remove_tmp_dir()
    ```
1441

1442
    `tmp_dir` will contain the path to the created temporary dir. It will be automatically removed at the end of the
1443
1444
1445
1446
1447
    test.


    2. Create a temporary dir of my choice, ensure it's empty before the test starts and don't
    empty it after the test.
1448

1449
1450
1451
1452
    ```python
    def test_whatever(self):
        tmp_dir = self.get_auto_remove_tmp_dir("./xxx")
    ```
1453

1454
1455
    This is useful for debug when you want to monitor a specific directory and want to make sure the previous tests
    didn't leave any data in there.
1456

1457
1458
    3. You can override the first two options by directly overriding the `before` and `after` args, leading to the
        following behavior:
1459

1460
    `before=True`: the temporary dir will always be cleared at the beginning of the test.
1461

1462
    `before=False`: if the temporary dir already existed, any existing files will remain there.
1463

1464
    `after=True`: the temporary dir will always be deleted at the end of the test.
1465

1466
    `after=False`: the temporary dir will always be left intact at the end of the test.
1467

1468
    Note 1: In order to run the equivalent of `rm -r` safely, only subdirs of the project repository checkout are
Sylvain Gugger's avatar
Sylvain Gugger committed
1469
1470
    allowed if an explicit `tmp_dir` is used, so that by mistake no `/tmp` or similar important part of the filesystem
    will get nuked. i.e. please always pass paths that start with `./`
1471

1472
1473
    Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested
    otherwise.
1474

Sylvain Gugger's avatar
Sylvain Gugger committed
1475
1476
    Feature 3: Get a copy of the `os.environ` object that sets up `PYTHONPATH` specific to the current test suite. This
    is useful for invoking external programs from the test suite - e.g. distributed training.
1477
1478


1479
1480
1481
1482
    ```python
    def test_whatever(self):
        env = self.get_env()
    ```"""
1483
1484

    def setUp(self):
1485
        # get_auto_remove_tmp_dir feature:
1486
1487
        self.teardown_tmp_dirs = []

1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
        # figure out the resolved paths for repo_root, tests, examples, etc.
        self._test_file_path = inspect.getfile(self.__class__)
        path = Path(self._test_file_path).resolve()
        self._test_file_dir = path.parents[0]
        for up in [1, 2, 3]:
            tmp_dir = path.parents[up]
            if (tmp_dir / "src").is_dir() and (tmp_dir / "tests").is_dir():
                break
        if tmp_dir:
            self._repo_root_dir = tmp_dir
        else:
            raise ValueError(f"can't figure out the root of the repo from {self._test_file_path}")
        self._tests_dir = self._repo_root_dir / "tests"
        self._examples_dir = self._repo_root_dir / "examples"
        self._src_dir = self._repo_root_dir / "src"

    @property
    def test_file_path(self):
        return self._test_file_path

    @property
    def test_file_path_str(self):
        return str(self._test_file_path)

    @property
    def test_file_dir(self):
        return self._test_file_dir

    @property
    def test_file_dir_str(self):
        return str(self._test_file_dir)

    @property
    def tests_dir(self):
        return self._tests_dir

    @property
    def tests_dir_str(self):
        return str(self._tests_dir)

    @property
    def examples_dir(self):
        return self._examples_dir

    @property
    def examples_dir_str(self):
        return str(self._examples_dir)

    @property
    def repo_root_dir(self):
        return self._repo_root_dir

    @property
    def repo_root_dir_str(self):
        return str(self._repo_root_dir)

    @property
    def src_dir(self):
        return self._src_dir

    @property
    def src_dir_str(self):
        return str(self._src_dir)

    def get_env(self):
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
1554
1555
        Return a copy of the `os.environ` object that sets up `PYTHONPATH` correctly, depending on the test suite it's
        invoked from. This is useful for invoking external programs from the test suite - e.g. distributed training.
1556

Sylvain Gugger's avatar
Sylvain Gugger committed
1557
1558
        It always inserts `./src` first, then `./tests` or `./examples` depending on the test suite type and finally
        the preset `PYTHONPATH` if any (all full resolved paths).
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571

        """
        env = os.environ.copy()
        paths = [self.src_dir_str]
        if "/examples" in self.test_file_dir_str:
            paths.append(self.examples_dir_str)
        else:
            paths.append(self.tests_dir_str)
        paths.append(env.get("PYTHONPATH", ""))

        env["PYTHONPATH"] = ":".join(paths)
        return env

1572
    def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None):
1573
1574
        """
        Args:
1575
1576
            tmp_dir (`string`, *optional*):
                if `None`:
1577
1578

                   - a unique temporary path will be created
1579
1580
                   - sets `before=True` if `before` is `None`
                   - sets `after=True` if `after` is `None`
1581
1582
                else:

1583
1584
1585
1586
                   - `tmp_dir` will be created
                   - sets `before=True` if `before` is `None`
                   - sets `after=False` if `after` is `None`
            before (`bool`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1587
1588
                If `True` and the `tmp_dir` already exists, make sure to empty it right away if `False` and the
                `tmp_dir` already exists, any existing files will remain there.
1589
            after (`bool`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1590
1591
                If `True`, delete the `tmp_dir` at the end of the test if `False`, leave the `tmp_dir` and its contents
                intact at the end of the test.
1592
1593

        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
1594
            tmp_dir(`string`): either the same value as passed via *tmp_dir* or the path to the auto-selected tmp dir
1595
1596
        """
        if tmp_dir is not None:
1597
1598
1599
1600
1601
1602
1603
1604
1605
            # defining the most likely desired behavior for when a custom path is provided.
            # this most likely indicates the debug mode where we want an easily locatable dir that:
            # 1. gets cleared out before the test (if it already exists)
            # 2. is left intact after the test
            if before is None:
                before = True
            if after is None:
                after = False

1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
            # using provided path
            path = Path(tmp_dir).resolve()

            # to avoid nuking parts of the filesystem, only relative paths are allowed
            if not tmp_dir.startswith("./"):
                raise ValueError(
                    f"`tmp_dir` can only be a relative path, i.e. `./some/path`, but received `{tmp_dir}`"
                )

            # ensure the dir is empty to start with
            if before is True and path.exists():
                shutil.rmtree(tmp_dir, ignore_errors=True)

            path.mkdir(parents=True, exist_ok=True)

        else:
1622
1623
1624
1625
1626
1627
1628
1629
1630
            # defining the most likely desired behavior for when a unique tmp path is auto generated
            # (not a debug mode), here we require a unique tmp dir that:
            # 1. is empty before the test (it will be empty in this situation anyway)
            # 2. gets fully removed after the test
            if before is None:
                before = True
            if after is None:
                after = True

1631
1632
1633
1634
1635
1636
1637
1638
1639
            # using unique tmp dir (always empty, regardless of `before`)
            tmp_dir = tempfile.mkdtemp()

        if after is True:
            # register for deletion
            self.teardown_tmp_dirs.append(tmp_dir)

        return tmp_dir

1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
    def python_one_liner_max_rss(self, one_liner_str):
        """
        Runs the passed python one liner (just the code) and returns how much max cpu memory was used to run the
        program.

        Args:
            one_liner_str (`string`):
                a python one liner code that gets passed to `python -c`

        Returns:
            max cpu memory bytes used to run the program. This value is likely to vary slightly from run to run.

        Requirements:
            this helper needs `/usr/bin/time` to be installed (`apt install time`)

        Example:

        ```
1658
        one_liner_str = 'from transformers import AutoModel; AutoModel.from_pretrained("google-t5/t5-large")'
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
        max_rss = self.python_one_liner_max_rss(one_liner_str)
        ```
        """

        if not cmd_exists("/usr/bin/time"):
            raise ValueError("/usr/bin/time is required, install with `apt install time`")

        cmd = shlex.split(f"/usr/bin/time -f %M python -c '{one_liner_str}'")
        with CaptureStd() as cs:
            execute_subprocess_async(cmd, env=self.get_env())
        # returned data is in KB so convert to bytes
        max_rss = int(cs.err.split("\n")[-2].replace("stderr: ", "")) * 1024
        return max_rss

1673
    def tearDown(self):
1674
        # get_auto_remove_tmp_dir feature: remove registered temp dirs
1675
1676
1677
        for path in self.teardown_tmp_dirs:
            shutil.rmtree(path, ignore_errors=True)
        self.teardown_tmp_dirs = []
1678
1679
1680
        if is_accelerate_available():
            AcceleratorState._reset_state()
            PartialState._reset_state()
1681

1682
1683
1684
1685
1686
            # delete all the env variables having `ACCELERATE` in them
            for k in list(os.environ.keys()):
                if "ACCELERATE" in k:
                    del os.environ[k]

1687
1688

def mockenv(**kwargs):
Sylvain Gugger's avatar
Sylvain Gugger committed
1689
    """
1690
1691
    this is a convenience wrapper, that allows this ::

Sylvain Gugger's avatar
Sylvain Gugger committed
1692
1693
    @mockenv(RUN_SLOW=True, USE_TF=False) def test_something():
        run_slow = os.getenv("RUN_SLOW", False) use_tf = os.getenv("USE_TF", False)
1694
1695

    """
1696
    return mock.patch.dict(os.environ, kwargs)
1697
1698


1699
1700
1701
1702
# from https://stackoverflow.com/a/34333710/9201239
@contextlib.contextmanager
def mockenv_context(*remove, **update):
    """
1703
    Temporarily updates the `os.environ` dictionary in-place. Similar to mockenv
1704

1705
    The `os.environ` dictionary is updated in-place so that the modification is sure to work in all situations.
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
1729
1730

    Args:
      remove: Environment variables to remove.
      update: Dictionary of environment variables and values to add/update.
    """
    env = os.environ
    update = update or {}
    remove = remove or []

    # List of environment variables being updated or removed.
    stomped = (set(update.keys()) | set(remove)) & set(env.keys())
    # Environment variables and values to restore on exit.
    update_after = {k: env[k] for k in stomped}
    # Environment variables and values to remove on exit.
    remove_after = frozenset(k for k in update if k not in env)

    try:
        env.update(update)
        [env.pop(k, None) for k in remove]
        yield
    finally:
        env.update(update_after)
        [env.pop(k) for k in remove_after]


1731
1732
1733
1734
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
# --- pytest conf functions --- #

# to avoid multiple invocation from tests/conftest.py and examples/conftest.py - make sure it's called only once
pytest_opt_registered = {}


def pytest_addoption_shared(parser):
    """
    This function is to be called from `conftest.py` via `pytest_addoption` wrapper that has to be defined there.

    It allows loading both `conftest.py` files at once without causing a failure due to adding the same `pytest`
    option.

    """
    option = "--make-reports"
    if option not in pytest_opt_registered:
        parser.addoption(
            option,
            action="store",
            default=False,
            help="generate report files. The value of this option is used as a prefix to report names",
        )
        pytest_opt_registered[option] = 1


1756
1757
def pytest_terminal_summary_main(tr, id):
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
1758
1759
    Generate multiple reports at the end of test suite run - each report goes into a dedicated file in the current
    directory. The report files are prefixed with the test suite name.
1760
1761
1762

    This function emulates --duration and -rA pytest arguments.

Sylvain Gugger's avatar
Sylvain Gugger committed
1763
1764
    This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined
    there.
1765
1766
1767

    Args:
    - tr: `terminalreporter` passed from `conftest.py`
1768
1769
    - id: unique id like `tests` or `examples` that will be incorporated into the final reports filenames - this is
      needed as some jobs have multiple runs of pytest, so we can't have them overwrite each other.
1770

Sylvain Gugger's avatar
Sylvain Gugger committed
1771
1772
1773
    NB: this functions taps into a private _pytest API and while unlikely, it could break should pytest do internal
    changes - also it calls default internal methods of terminalreporter which can be hijacked by various `pytest-`
    plugins and interfere.
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785

    """
    from _pytest.config import create_terminal_writer

    if not len(id):
        id = "tests"

    config = tr.config
    orig_writer = config.get_terminal_writer()
    orig_tbstyle = config.option.tbstyle
    orig_reportchars = tr.reportchars

1786
    dir = f"reports/{id}"
1787
    Path(dir).mkdir(parents=True, exist_ok=True)
Stas Bekman's avatar
Stas Bekman committed
1788
    report_files = {
1789
        k: f"{dir}/{k}.txt"
Stas Bekman's avatar
Stas Bekman committed
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
        for k in [
            "durations",
            "errors",
            "failures_long",
            "failures_short",
            "failures_line",
            "passes",
            "stats",
            "summary_short",
            "warnings",
        ]
    }
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
1821

    # custom durations report
    # note: there is no need to call pytest --durations=XX to get this separate report
    # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/runner.py#L66
    dlist = []
    for replist in tr.stats.values():
        for rep in replist:
            if hasattr(rep, "duration"):
                dlist.append(rep)
    if dlist:
        dlist.sort(key=lambda x: x.duration, reverse=True)
        with open(report_files["durations"], "w") as f:
            durations_min = 0.05  # sec
            f.write("slowest durations\n")
            for i, rep in enumerate(dlist):
                if rep.duration < durations_min:
                    f.write(f"{len(dlist)-i} durations < {durations_min} secs were omitted")
                    break
                f.write(f"{rep.duration:02.2f}s {rep.when:<8} {rep.nodeid}\n")

Stas Bekman's avatar
Stas Bekman committed
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
    def summary_failures_short(tr):
        # expecting that the reports were --tb=long (default) so we chop them off here to the last frame
        reports = tr.getreports("failed")
        if not reports:
            return
        tr.write_sep("=", "FAILURES SHORT STACK")
        for rep in reports:
            msg = tr._getfailureheadline(rep)
            tr.write_sep("_", msg, red=True, bold=True)
            # chop off the optional leading extra frames, leaving only the last one
            longrepr = re.sub(r".*_ _ _ (_ ){10,}_ _ ", "", rep.longreprtext, 0, re.M | re.S)
            tr._tw.line(longrepr)
            # note: not printing out any rep.sections to keep the report short

1836
1837
1838
1839
    # use ready-made report funcs, we are just hijacking the filehandle to log to a dedicated file each
    # adapted from https://github.com/pytest-dev/pytest/blob/897f151e/src/_pytest/terminal.py#L814
    # note: some pytest plugins may interfere by hijacking the default `terminalreporter` (e.g.
    # pytest-instafail does that)
Stas Bekman's avatar
Stas Bekman committed
1840
1841
1842
1843

    # report failures with line/short/long styles
    config.option.tbstyle = "auto"  # full tb
    with open(report_files["failures_long"], "w") as f:
1844
1845
1846
        tr._tw = create_terminal_writer(config, f)
        tr.summary_failures()

Stas Bekman's avatar
Stas Bekman committed
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
    # config.option.tbstyle = "short" # short tb
    with open(report_files["failures_short"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        summary_failures_short(tr)

    config.option.tbstyle = "line"  # one line per error
    with open(report_files["failures_line"], "w") as f:
        tr._tw = create_terminal_writer(config, f)
        tr.summary_failures()

    with open(report_files["errors"], "w") as f:
1858
1859
1860
        tr._tw = create_terminal_writer(config, f)
        tr.summary_errors()

Stas Bekman's avatar
Stas Bekman committed
1861
    with open(report_files["warnings"], "w") as f:
1862
1863
1864
1865
        tr._tw = create_terminal_writer(config, f)
        tr.summary_warnings()  # normal warnings
        tr.summary_warnings()  # final warnings

Stas Bekman's avatar
Stas Bekman committed
1866
    tr.reportchars = "wPpsxXEf"  # emulate -rA (used in summary_passes() and short_test_summary())
1867
1868
1869
1870
1871
1872
1873

    # Skip the `passes` report, as it starts to take more than 5 minutes, and sometimes it timeouts on CircleCI if it
    # takes > 10 minutes (as this part doesn't generate any output on the terminal).
    # (also, it seems there is no useful information in this report, and we rarely need to read it)
    # with open(report_files["passes"], "w") as f:
    #     tr._tw = create_terminal_writer(config, f)
    #     tr.summary_passes()
1874

Stas Bekman's avatar
Stas Bekman committed
1875
    with open(report_files["summary_short"], "w") as f:
1876
1877
1878
        tr._tw = create_terminal_writer(config, f)
        tr.short_test_summary()

Stas Bekman's avatar
Stas Bekman committed
1879
    with open(report_files["stats"], "w") as f:
1880
1881
1882
1883
1884
1885
1886
        tr._tw = create_terminal_writer(config, f)
        tr.summary_stats()

    # restore:
    tr._tw = orig_writer
    tr.reportchars = orig_reportchars
    config.option.tbstyle = orig_tbstyle
1887
1888


1889
# --- distributed testing functions --- #
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920
1921
1922
1923
1924
1925
1926
1927
1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
1939
1940
1941
1942
1943

# adapted from https://stackoverflow.com/a/59041913/9201239
import asyncio  # noqa


class _RunOutput:
    def __init__(self, returncode, stdout, stderr):
        self.returncode = returncode
        self.stdout = stdout
        self.stderr = stderr


async def _read_stream(stream, callback):
    while True:
        line = await stream.readline()
        if line:
            callback(line)
        else:
            break


async def _stream_subprocess(cmd, env=None, stdin=None, timeout=None, quiet=False, echo=False) -> _RunOutput:
    if echo:
        print("\nRunning: ", " ".join(cmd))

    p = await asyncio.create_subprocess_exec(
        cmd[0],
        *cmd[1:],
        stdin=stdin,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
        env=env,
    )

    # note: there is a warning for a possible deadlock when using `wait` with huge amounts of data in the pipe
    # https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.wait
    #
    # If it starts hanging, will need to switch to the following code. The problem is that no data
    # will be seen until it's done and if it hangs for example there will be no debug info.
    # out, err = await p.communicate()
    # return _RunOutput(p.returncode, out, err)

    out = []
    err = []

    def tee(line, sink, pipe, label=""):
        line = line.decode("utf-8").rstrip()
        sink.append(line)
        if not quiet:
            print(label, line, file=pipe)

    # XXX: the timeout doesn't seem to make any difference here
    await asyncio.wait(
        [
Stas Bekman's avatar
Stas Bekman committed
1944
            _read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:")),
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
            _read_stream(p.stderr, lambda l: tee(l, err, sys.stderr, label="stderr:")),
        ],
        timeout=timeout,
    )
    return _RunOutput(await p.wait(), out, err)


def execute_subprocess_async(cmd, env=None, stdin=None, timeout=180, quiet=False, echo=True) -> _RunOutput:
    loop = asyncio.get_event_loop()
    result = loop.run_until_complete(
        _stream_subprocess(cmd, env=env, stdin=stdin, timeout=timeout, quiet=quiet, echo=echo)
    )

    cmd_str = " ".join(cmd)
    if result.returncode > 0:
1960
        stderr = "\n".join(result.stderr)
1961
        raise RuntimeError(
1962
1963
            f"'{cmd_str}' failed with returncode {result.returncode}\n\n"
            f"The combined stderr from workers follows:\n{stderr}"
1964
        )
Stas Bekman's avatar
Stas Bekman committed
1965
1966
1967
1968

    # check that the subprocess actually did run and produced some output, should the test rely on
    # the remote side to do the testing
    if not result.stdout and not result.stderr:
1969
1970
1971
        raise RuntimeError(f"'{cmd_str}' produced no output.")

    return result
1972
1973


1974
1975
def pytest_xdist_worker_id():
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
1976
1977
    Returns an int value of worker's numerical id under `pytest-xdist`'s concurrent workers `pytest -n N` regime, or 0
    if `-n 1` or `pytest-xdist` isn't being used.
1978
1979
1980
1981
1982
1983
1984
1985
    """
    worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
    worker = re.sub(r"^gw", "", worker, 0, re.M)
    return int(worker)


def get_torch_dist_unique_port():
    """
1986
    Returns a port number that can be fed to `torch.distributed.launch`'s `--master_port` argument.
1987

Sylvain Gugger's avatar
Sylvain Gugger committed
1988
1989
    Under `pytest-xdist` it adds a delta number based on a worker id so that concurrent tests don't try to use the same
    port at once.
1990
1991
1992
1993
1994
1995
    """
    port = 29500
    uniq_delta = pytest_xdist_worker_id()
    return port + uniq_delta


1996
1997
1998
1999
2000
def nested_simplify(obj, decimals=3):
    """
    Simplifies an object by rounding float numbers, and downcasting tensors/numpy arrays to get simple equality test
    within tests.
    """
2001
2002
    import numpy as np

2003
2004
    if isinstance(obj, list):
        return [nested_simplify(item, decimals) for item in obj]
2005
2006
    if isinstance(obj, tuple):
        return tuple([nested_simplify(item, decimals) for item in obj])
2007
2008
    elif isinstance(obj, np.ndarray):
        return nested_simplify(obj.tolist())
2009
    elif isinstance(obj, Mapping):
2010
        return {nested_simplify(k, decimals): nested_simplify(v, decimals) for k, v in obj.items()}
2011
    elif isinstance(obj, (str, int, np.int64)):
2012
        return obj
2013
2014
    elif obj is None:
        return obj
2015
    elif is_torch_available() and isinstance(obj, torch.Tensor):
2016
        return nested_simplify(obj.tolist(), decimals)
2017
2018
2019
2020
    elif is_tf_available() and tf.is_tensor(obj):
        return nested_simplify(obj.numpy().tolist())
    elif isinstance(obj, float):
        return round(obj, decimals)
2021
    elif isinstance(obj, (np.int32, np.float32)):
2022
        return nested_simplify(obj.item(), decimals)
2023
2024
    else:
        raise Exception(f"Not supported: {type(obj)}")
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041


def check_json_file_has_correct_format(file_path):
    with open(file_path, "r") as f:
        lines = f.readlines()
        if len(lines) == 1:
            # length can only be 1 if dict is empty
            assert lines[0] == "{}"
        else:
            # otherwise make sure json has correct format (at least 3 lines)
            assert len(lines) >= 3
            # each key one line, ident should be 2, min length is 3
            assert lines[0].strip() == "{"
            for line in lines[1:-1]:
                left_indent = len(lines[1]) - len(lines[1].lstrip())
                assert left_indent == 2
            assert lines[-1].strip() == "}"
NielsRogge's avatar
NielsRogge committed
2042
2043
2044
2045
2046
2047


def to_2tuple(x):
    if isinstance(x, collections.abc.Iterable):
        return x
    return (x, x)
Zachary Mueller's avatar
Zachary Mueller committed
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069


# These utils relate to ensuring the right error message is received when running scripts
class SubprocessCallException(Exception):
    pass


def run_command(command: List[str], return_stdout=False):
    """
    Runs `command` with `subprocess.check_output` and will potentially return the `stdout`. Will also properly capture
    if an error occured while running `command`
    """
    try:
        output = subprocess.check_output(command, stderr=subprocess.STDOUT)
        if return_stdout:
            if hasattr(output, "decode"):
                output = output.decode("utf-8")
            return output
    except subprocess.CalledProcessError as e:
        raise SubprocessCallException(
            f"Command `{' '.join(command)}` failed with the following error:\n\n{e.output.decode()}"
        ) from e
2070
2071
2072
2073
2074


class RequestCounter:
    """
    Helper class that will count all requests made online.
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085

    Might not be robust if urllib3 changes its logging format but should be good enough for us.

    Usage:
    ```py
    with RequestCounter() as counter:
        _ = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-bert")
    assert counter["GET"] == 0
    assert counter["HEAD"] == 1
    assert counter.total_calls == 1
    ```
2086
2087
2088
    """

    def __enter__(self):
2089
2090
2091
        self._counter = defaultdict(int)
        self.patcher = patch.object(urllib3.connectionpool.log, "debug", wraps=urllib3.connectionpool.log.debug)
        self.mock = self.patcher.start()
2092
2093
        return self

2094
2095
2096
2097
2098
2099
2100
2101
    def __exit__(self, *args, **kwargs) -> None:
        for call in self.mock.call_args_list:
            log = call.args[0] % call.args[1:]
            for method in ("HEAD", "GET", "POST", "PUT", "DELETE", "CONNECT", "OPTIONS", "TRACE", "PATCH"):
                if method in log:
                    self._counter[method] += 1
                    break
        self.patcher.stop()
2102

2103
2104
    def __getitem__(self, key: str) -> int:
        return self._counter[key]
2105

2106
2107
2108
    @property
    def total_calls(self) -> int:
        return sum(self._counter.values())
2109
2110


2111
def is_flaky(max_attempts: int = 5, wait_before_retry: Optional[float] = None, description: Optional[str] = None):
2112
2113
2114
2115
2116
2117
2118
2119
    """
    To decorate flaky tests. They will be retried on failures.

    Args:
        max_attempts (`int`, *optional*, defaults to 5):
            The maximum number of attempts to retry the flaky test.
        wait_before_retry (`float`, *optional*):
            If provided, will wait that number of seconds before retrying the test.
2120
2121
2122
        description (`str`, *optional*):
            A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors,
            etc.)
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
    """

    def decorator(test_func_ref):
        @functools.wraps(test_func_ref)
        def wrapper(*args, **kwargs):
            retry_count = 1

            while retry_count < max_attempts:
                try:
                    return test_func_ref(*args, **kwargs)

                except Exception as err:
                    print(f"Test failed with {err} at try {retry_count}/{max_attempts}.", file=sys.stderr)
                    if wait_before_retry is not None:
                        time.sleep(wait_before_retry)
                    retry_count += 1

            return test_func_ref(*args, **kwargs)

        return wrapper

    return decorator
2145
2146


2147
def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None):
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
    """
    To run a test in a subprocess. In particular, this can avoid (GPU) memory issue.

    Args:
        test_case (`unittest.TestCase`):
            The test that will run `target_func`.
        target_func (`Callable`):
            The function implementing the actual testing logic.
        inputs (`dict`, *optional*, defaults to `None`):
            The inputs that will be passed to `target_func` through an (input) queue.
2158
2159
2160
        timeout (`int`, *optional*, defaults to `None`):
            The timeout (in seconds) that will be passed to the input and output queues. If not specified, the env.
            variable `PYTEST_TIMEOUT` will be checked. If still `None`, its value will be set to `600`.
2161
    """
2162
2163
    if timeout is None:
        timeout = int(os.environ.get("PYTEST_TIMEOUT", 600))
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187

    start_methohd = "spawn"
    ctx = multiprocessing.get_context(start_methohd)

    input_queue = ctx.Queue(1)
    output_queue = ctx.JoinableQueue(1)

    # We can't send `unittest.TestCase` to the child, otherwise we get issues regarding pickle.
    input_queue.put(inputs, timeout=timeout)

    process = ctx.Process(target=target_func, args=(input_queue, output_queue, timeout))
    process.start()
    # Kill the child process if we can't get outputs from it in time: otherwise, the hanging subprocess prevents
    # the test to exit properly.
    try:
        results = output_queue.get(timeout=timeout)
        output_queue.task_done()
    except Exception as e:
        process.terminate()
        test_case.fail(e)
    process.join(timeout=timeout)

    if results["error"] is not None:
        test_case.fail(f'{results["error"]}')
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200


"""
The following contains utils to run the documentation tests without having to overwrite any files.

The `preprocess_string` function adds `# doctest: +IGNORE_RESULT` markers on the fly anywhere a `load_dataset` call is
made as a print would otherwise fail the corresonding line.

To skip cuda tests, make sure to call `SKIP_CUDA_DOCTEST=1 pytest --doctest-modules <path_to_files_to_test>
"""


def preprocess_string(string, skip_cuda_tests):
2201
    """Prepare a docstring or a `.md` file to be run by doctest.
2202

2203
    The argument `string` would be the whole file content if it is a `.md` file. For a python file, it would be one of
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
    its docstring. In each case, it may contain multiple python code examples. If `skip_cuda_tests` is `True` and a
    cuda stuff is detective (with a heuristic), this method will return an empty string so no doctest will be run for
    `string`.
    """
    codeblock_pattern = r"(```(?:python|py)\s*\n\s*>>> )((?:.*?\n)*?.*?```)"
    codeblocks = re.split(re.compile(codeblock_pattern, flags=re.MULTILINE | re.DOTALL), string)
    is_cuda_found = False
    for i, codeblock in enumerate(codeblocks):
        if "load_dataset(" in codeblock and "# doctest: +IGNORE_RESULT" not in codeblock:
            codeblocks[i] = re.sub(r"(>>> .*load_dataset\(.*)", r"\1 # doctest: +IGNORE_RESULT", codeblock)
        if (
            (">>>" in codeblock or "..." in codeblock)
            and re.search(r"cuda|to\(0\)|device=0", codeblock)
            and skip_cuda_tests
        ):
            is_cuda_found = True
            break
Yih-Dar's avatar
Yih-Dar committed
2221

2222
2223
2224
    modified_string = ""
    if not is_cuda_found:
        modified_string = "".join(codeblocks)
Yih-Dar's avatar
Yih-Dar committed
2225

2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
    return modified_string


class HfDocTestParser(doctest.DocTestParser):
    """
    Overwrites the DocTestParser from doctest to properly parse the codeblocks that are formatted with black. This
    means that there are no extra lines at the end of our snippets. The `# doctest: +IGNORE_RESULT` marker is also
    added anywhere a `load_dataset` call is made as a print would otherwise fail the corresponding line.

    Tests involving cuda are skipped base on a naive pattern that should be updated if it is not enough.
    """

    # This regular expression is used to find doctest examples in a
    # string.  It defines three groups: `source` is the source code
    # (including leading indentation and prompts); `indent` is the
    # indentation of the first (PS1) line of the source code; and
    # `want` is the expected output (including leading indentation).
    # fmt: off
    _EXAMPLE_RE = re.compile(r'''
        # Source consists of a PS1 line followed by zero or more PS2 lines.
        (?P<source>
            (?:^(?P<indent> [ ]*) >>>    .*)    # PS1 line
            (?:\n           [ ]*  \.\.\. .*)*)  # PS2 lines
        \n?
        # Want consists of any non-blank lines that do not start with PS1.
        (?P<want> (?:(?![ ]*$)    # Not a blank line
             (?![ ]*>>>)          # Not a line starting with PS1
             # !!!!!!!!!!! HF Specific !!!!!!!!!!!
             (?:(?!```).)*        # Match any character except '`' until a '```' is found (this is specific to HF because black removes the last line)
             # !!!!!!!!!!! HF Specific !!!!!!!!!!!
             (?:\n|$)  # Match a new line or end of string
          )*)
        ''', re.MULTILINE | re.VERBOSE
    )
    # fmt: on

    # !!!!!!!!!!! HF Specific !!!!!!!!!!!
    skip_cuda_tests: bool = bool(os.environ.get("SKIP_CUDA_DOCTEST", False))
    # !!!!!!!!!!! HF Specific !!!!!!!!!!!

    def parse(self, string, name="<string>"):
        """
        Overwrites the `parse` method to incorporate a skip for CUDA tests, and remove logs and dataset prints before
        calling `super().parse`
        """
        string = preprocess_string(string, self.skip_cuda_tests)
        return super().parse(string, name)


class HfDoctestModule(Module):
    """
    Overwrites the `DoctestModule` of the pytest package to make sure the HFDocTestParser is used when discovering
    tests.
    """

    def collect(self) -> Iterable[DoctestItem]:
        class MockAwareDocTestFinder(doctest.DocTestFinder):
            """A hackish doctest finder that overrides stdlib internals to fix a stdlib bug.

            https://github.com/pytest-dev/pytest/issues/3456 https://bugs.python.org/issue25532
            """

            def _find_lineno(self, obj, source_lines):
                """Doctest code does not take into account `@property`, this
                is a hackish way to fix it. https://bugs.python.org/issue17446

                Wrapped Doctests will need to be unwrapped so the correct line number is returned. This will be
                reported upstream. #8796
                """
                if isinstance(obj, property):
                    obj = getattr(obj, "fget", obj)

                if hasattr(obj, "__wrapped__"):
                    # Get the main obj in case of it being wrapped
                    obj = inspect.unwrap(obj)

                # Type ignored because this is a private function.
                return super()._find_lineno(  # type:ignore[misc]
                    obj,
                    source_lines,
                )

            def _find(self, tests, obj, name, module, source_lines, globs, seen) -> None:
                if _is_mocked(obj):
                    return
                with _patch_unwrap_mock_aware():
                    # Type ignored because this is a private function.
                    super()._find(  # type:ignore[misc]
                        tests, obj, name, module, source_lines, globs, seen
                    )

        if self.path.name == "conftest.py":
            module = self.config.pluginmanager._importconftest(
                self.path,
                self.config.getoption("importmode"),
                rootpath=self.config.rootpath,
            )
        else:
            try:
                module = import_path(
                    self.path,
                    root=self.config.rootpath,
                    mode=self.config.getoption("importmode"),
                )
            except ImportError:
                if self.config.getvalue("doctest_ignore_import_errors"):
                    skip("unable to import module %r" % self.path)
                else:
                    raise

        # !!!!!!!!!!! HF Specific !!!!!!!!!!!
        finder = MockAwareDocTestFinder(parser=HfDocTestParser())
        # !!!!!!!!!!! HF Specific !!!!!!!!!!!
        optionflags = get_optionflags(self)
        runner = _get_runner(
            verbose=False,
            optionflags=optionflags,
            checker=_get_checker(),
            continue_on_failure=_get_continue_on_failure(self.config),
        )
        for test in finder.find(module, module.__name__):
            if test.examples:  # skip empty doctests and cuda
                yield DoctestItem.from_parent(self, name=test.name, runner=runner, dtest=test)
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431


def _device_agnostic_dispatch(device: str, dispatch_table: Dict[str, Callable], *args, **kwargs):
    if device not in dispatch_table:
        return dispatch_table["default"](*args, **kwargs)

    fn = dispatch_table[device]

    # Some device agnostic functions return values. Need to guard against `None`
    # instead at user level.
    if fn is None:
        return None
    return fn(*args, **kwargs)


if is_torch_available():
    # Mappings from device names to callable functions to support device agnostic
    # testing.
    BACKEND_MANUAL_SEED = {"cuda": torch.cuda.manual_seed, "cpu": torch.manual_seed, "default": torch.manual_seed}
    BACKEND_EMPTY_CACHE = {"cuda": torch.cuda.empty_cache, "cpu": None, "default": None}
    BACKEND_DEVICE_COUNT = {"cuda": torch.cuda.device_count, "cpu": lambda: 0, "default": lambda: 1}


def backend_manual_seed(device: str, seed: int):
    return _device_agnostic_dispatch(device, BACKEND_MANUAL_SEED, seed)


def backend_empty_cache(device: str):
    return _device_agnostic_dispatch(device, BACKEND_EMPTY_CACHE)


def backend_device_count(device: str):
    return _device_agnostic_dispatch(device, BACKEND_DEVICE_COUNT)


if is_torch_available():
    # If `TRANSFORMERS_TEST_DEVICE_SPEC` is enabled we need to import extra entries
    # into device to function mappings.
    if "TRANSFORMERS_TEST_DEVICE_SPEC" in os.environ:
        device_spec_path = os.environ["TRANSFORMERS_TEST_DEVICE_SPEC"]
        if not Path(device_spec_path).is_file():
            raise ValueError(
                f"Specified path to device spec file is not a file or not found. Received '{device_spec_path}"
            )

        # Try to strip extension for later import – also verifies we are importing a
        # python file.
        try:
            import_name = device_spec_path[: device_spec_path.index(".py")]
        except ValueError as e:
            raise ValueError(f"Provided device spec file was not a Python file! Received '{device_spec_path}") from e

        device_spec_module = importlib.import_module(import_name)

        # Imported file must contain `DEVICE_NAME`. If it doesn't, terminate early.
        try:
            device_name = device_spec_module.DEVICE_NAME
        except AttributeError as e:
            raise AttributeError("Device spec file did not contain `DEVICE_NAME`") from e

        if "TRANSFORMERS_TEST_DEVICE" in os.environ and torch_device != device_name:
            msg = f"Mismatch between environment variable `TRANSFORMERS_TEST_DEVICE` '{torch_device}' and device found in spec '{device_name}'\n"
            msg += "Either unset `TRANSFORMERS_TEST_DEVICE` or ensure it matches device spec name."
            raise ValueError(msg)

        torch_device = device_name

        def update_mapping_from_spec(device_fn_dict: Dict[str, Callable], attribute_name: str):
            try:
                # Try to import the function directly
                spec_fn = getattr(device_spec_module, attribute_name)
                device_fn_dict[torch_device] = spec_fn
            except AttributeError as e:
                # If the function doesn't exist, and there is no default, throw an error
                if "default" not in device_fn_dict:
                    raise AttributeError(
                        f"`{attribute_name}` not found in '{device_spec_path}' and no default fallback function found."
                    ) from e

        # Add one entry here for each `BACKEND_*` dictionary.
        update_mapping_from_spec(BACKEND_MANUAL_SEED, "MANUAL_SEED_FN")
        update_mapping_from_spec(BACKEND_EMPTY_CACHE, "EMPTY_CACHE_FN")
        update_mapping_from_spec(BACKEND_DEVICE_COUNT, "DEVICE_COUNT_FN")