testing_utils.py 46 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.

15
import contextlib
16
import inspect
17
import logging
18
import os
19
import re
20
import shutil
21
import sys
22
import tempfile
Aymeric Augustin's avatar
Aymeric Augustin committed
23
import unittest
24
from distutils.util import strtobool
25
from io import StringIO
26
from pathlib import Path
27
from typing import Iterator, Union
28
from unittest import mock
29

30
31
from transformers import logging as transformers_logging

32
from .deepspeed import is_deepspeed_available
33
from .file_utils import (
34
    is_detectron2_available,
35
36
    is_faiss_available,
    is_flax_available,
37
    is_keras2onnx_available,
38
    is_librosa_available,
39
    is_onnx_available,
40
    is_pandas_available,
41
    is_pyctcdecode_available,
42
    is_pytesseract_available,
43
    is_pytorch_quantization_available,
yujun's avatar
yujun committed
44
    is_rjieba_available,
45
46
    is_scatter_available,
    is_sentencepiece_available,
Patrick von Platen's avatar
Patrick von Platen committed
47
    is_soundfile_availble,
Kamal Raj's avatar
Kamal Raj committed
48
    is_tensorflow_probability_available,
49
    is_tf_available,
NielsRogge's avatar
NielsRogge committed
50
    is_timm_available,
51
52
    is_tokenizers_available,
    is_torch_available,
53
    is_torch_bf16_available,
54
    is_torch_tf32_available,
55
    is_torch_tpu_available,
Suraj Patil's avatar
Suraj Patil committed
56
    is_torchaudio_available,
57
    is_vision_available,
58
)
59
from .integrations import is_optuna_available, is_ray_available, is_sigopt_available
60
61


Julien Chaumond's avatar
Julien Chaumond committed
62
SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy"
63
DUMMY_UNKNOWN_IDENTIFIER = "julien-c/dummy-unknown"
64
DUMMY_DIFF_TOKENIZER_IDENTIFIER = "julien-c/dummy-diff-tokenizer"
Julien Chaumond's avatar
Julien Chaumond committed
65
# Used to test Auto{Config, Model, Tokenizer} model_type detection.
Julien Chaumond's avatar
Julien Chaumond committed
66

Sylvain Gugger's avatar
Sylvain Gugger committed
67
68
69
70
71
# Used to test the hub
USER = "__DUMMY_TRANSFORMERS_USER__"
PASS = "__DUMMY_TRANSFORMERS_PASS__"
ENDPOINT_STAGING = "https://moon-staging.huggingface.co"

Julien Chaumond's avatar
Julien Chaumond committed
72

73
def parse_flag_from_env(key, default=False):
74
    try:
75
76
77
78
79
80
81
82
83
84
        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.
85
            raise ValueError(f"If set, {key} must be yes or no.")
86
87
    return _value

88

Julien Chaumond's avatar
Julien Chaumond committed
89
90
91
92
93
94
95
96
97
def parse_int_from_env(key, default=None):
    try:
        value = os.environ[key]
    except KeyError:
        _value = default
    else:
        try:
            _value = int(value)
        except ValueError:
98
            raise ValueError(f"If set, {key} must be a int.")
Julien Chaumond's avatar
Julien Chaumond committed
99
100
101
    return _value


102
_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)
103
_run_pt_tf_cross_tests = parse_flag_from_env("RUN_PT_TF_CROSS_TESTS", default=False)
104
_run_pt_flax_cross_tests = parse_flag_from_env("RUN_PT_FLAX_CROSS_TESTS", default=False)
105
_run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False)
Sylvain Gugger's avatar
Sylvain Gugger committed
106
_run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False)
107
_run_pipeline_tests = parse_flag_from_env("RUN_PIPELINE_TESTS", default=False)
108
_run_git_lfs_tests = parse_flag_from_env("RUN_GIT_LFS_TESTS", default=False)
Julien Chaumond's avatar
Julien Chaumond committed
109
_tf_gpu_memory_limit = parse_int_from_env("TF_GPU_MEMORY_LIMIT", default=None)
110
111


112
113
114
115
116
117
118
119
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.

    """
120
    if not _run_pt_tf_cross_tests or not is_torch_available() or not is_tf_available():
121
122
123
124
125
126
127
128
129
130
        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)


131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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)


150
151
152
153
def is_pipeline_test(test_case):
    """
    Decorator marking a test as a pipeline test.

154
155
    Pipeline tests are skipped by default and we can run only them by setting RUN_PIPELINE_TESTS environment variable
    to a truthy value and selecting the is_pipeline_test pytest mark.
156
157
158
159
160
161
162
163
164
165
166
167
168

    """
    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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
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)


186
187
188
189
def slow(test_case):
    """
    Decorator marking a test as slow.

Sylvain Gugger's avatar
Sylvain Gugger committed
190
    Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.
191
192
193

    """
    if not _run_slow_tests:
194
195
196
        return unittest.skip("test is slow")(test_case)
    else:
        return test_case
197
198


Lysandre Debut's avatar
Lysandre Debut committed
199
200
201
202
203
204
205
206
207
208
209
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)


210
211
212
213
def custom_tokenizers(test_case):
    """
    Decorator marking a test for a custom tokenizer.

Sylvain Gugger's avatar
Sylvain Gugger committed
214
215
    Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS
    environment variable to a truthy value to run them.
216
217
    """
    if not _run_custom_tokenizers:
218
219
220
        return unittest.skip("test of custom tokenizers")(test_case)
    else:
        return test_case
221
222


223
224
225
226
227
228
229
230
231
232
233
234
235
def require_git_lfs(test_case):
    """
    Decorator marking a test that requires git-lfs.

    git-lfs requires additional dependencies, and tests are skipped by default. Set the RUN_GIT_LFS_TESTS environment
    variable to a truthy value to run them.
    """
    if not _run_git_lfs_tests:
        return unittest.skip("test of git lfs workflow")(test_case)
    else:
        return test_case


yujun's avatar
yujun committed
236
237
238
239
240
241
242
243
244
245
def require_rjieba(test_case):
    """
    Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed.
    """
    if not is_rjieba_available():
        return unittest.skip("test requires rjieba")(test_case)
    else:
        return test_case


246
247
248
249
250
251
252
def require_keras2onnx(test_case):
    if not is_keras2onnx_available():
        return unittest.skip("test requires keras2onnx")(test_case)
    else:
        return test_case


253
254
255
256
257
258
259
def require_onnx(test_case):
    if not is_onnx_available():
        return unittest.skip("test requires ONNX")(test_case)
    else:
        return test_case


NielsRogge's avatar
NielsRogge committed
260
261
262
263
264
265
266
267
268
269
270
271
272
def require_timm(test_case):
    """
    Decorator marking a test that requires Timm.

    These tests are skipped when Timm isn't installed.

    """
    if not is_timm_available():
        return unittest.skip("test requires Timm")(test_case)
    else:
        return test_case


273
274
275
276
277
278
279
def require_torch(test_case):
    """
    Decorator marking a test that requires PyTorch.

    These tests are skipped when PyTorch isn't installed.

    """
280
    if not is_torch_available():
281
282
283
        return unittest.skip("test requires PyTorch")(test_case)
    else:
        return test_case
284
285


286
287
288
289
290
291
292
def require_torch_scatter(test_case):
    """
    Decorator marking a test that requires PyTorch scatter.

    These tests are skipped when PyTorch scatter isn't installed.

    """
293
    if not is_scatter_available():
294
295
        return unittest.skip("test requires PyTorch scatter")(test_case)
    else:
Suraj Patil's avatar
Suraj Patil committed
296
297
298
        return test_case


Kamal Raj's avatar
Kamal Raj committed
299
300
301
302
303
304
305
306
307
308
309
310
311
def require_tensorflow_probability(test_case):
    """
    Decorator marking a test that requires TensorFlow probability.

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

    """
    if not is_tensorflow_probability_available():
        return unittest.skip("test requires TensorFlow probability")(test_case)
    else:
        return test_case


Suraj Patil's avatar
Suraj Patil committed
312
313
def require_torchaudio(test_case):
    """
314
    Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed.
Suraj Patil's avatar
Suraj Patil committed
315
    """
316
    if not is_torchaudio_available():
Suraj Patil's avatar
Suraj Patil committed
317
318
        return unittest.skip("test requires torchaudio")(test_case)
    else:
319
320
321
        return test_case


322
323
def require_tf(test_case):
    """
324
    Decorator marking a test that requires TensorFlow. These tests are skipped when TensorFlow isn't installed.
325
    """
326
    if not is_tf_available():
327
328
329
        return unittest.skip("test requires TensorFlow")(test_case)
    else:
        return test_case
330
331


332
333
def require_flax(test_case):
    """
334
    Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed
335
    """
336
    if not is_flax_available():
337
338
339
340
        test_case = unittest.skip("test requires JAX & Flax")(test_case)
    return test_case


341
342
def require_sentencepiece(test_case):
    """
343
    Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.
344
    """
345
    if not is_sentencepiece_available():
346
347
348
349
350
351
352
        return unittest.skip("test requires SentencePiece")(test_case)
    else:
        return test_case


def require_tokenizers(test_case):
    """
353
    Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed.
354
    """
355
    if not is_tokenizers_available():
356
357
358
359
360
        return unittest.skip("test requires tokenizers")(test_case)
    else:
        return test_case


NielsRogge's avatar
NielsRogge committed
361
362
363
364
def require_pandas(test_case):
    """
    Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed.
    """
365
    if not is_pandas_available():
NielsRogge's avatar
NielsRogge committed
366
367
368
369
370
        return unittest.skip("test requires pandas")(test_case)
    else:
        return test_case


371
372
373
374
375
376
377
378
379
380
def require_pytesseract(test_case):
    """
    Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed.
    """
    if not is_pytesseract_available():
        return unittest.skip("test requires PyTesseract")(test_case)
    else:
        return test_case


NielsRogge's avatar
NielsRogge committed
381
382
383
384
385
def require_scatter(test_case):
    """
    Decorator marking a test that requires PyTorch Scatter. These tests are skipped when PyTorch Scatter isn't
    installed.
    """
386
    if not is_scatter_available():
NielsRogge's avatar
NielsRogge committed
387
388
389
390
391
        return unittest.skip("test requires PyTorch Scatter")(test_case)
    else:
        return test_case


392
393
394
395
396
397
398
399
400
401
402
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.
    """
    if not is_pytorch_quantization_available():
        return unittest.skip("test requires PyTorch Quantization Toolkit")(test_case)
    else:
        return test_case


403
def require_vision(test_case):
404
    """
405
406
407
408
409
410
411
412
    Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't
    installed.
    """
    if not is_vision_available():
        return unittest.skip("test requires vision")(test_case)
    else:
        return test_case

413

414
415
416
417
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.
418

419
    To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu"
420
    """
421
    if not is_torch_available():
422
423
424
425
426
427
        return unittest.skip("test requires PyTorch")(test_case)

    import torch

    if torch.cuda.device_count() < 2:
        return unittest.skip("test requires multiple GPUs")(test_case)
428
429
    else:
        return test_case
430
431


432
def require_torch_non_multi_gpu(test_case):
433
434
435
    """
    Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch).
    """
436
    if not is_torch_available():
437
438
439
440
441
442
        return unittest.skip("test requires PyTorch")(test_case)

    import torch

    if torch.cuda.device_count() > 1:
        return unittest.skip("test requires 0 or 1 GPU")(test_case)
443
444
    else:
        return test_case
445
446


447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
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

    if torch.cuda.device_count() > 2:
        return unittest.skip("test requires 0 or 1 or 2 GPUs")(test_case)
    else:
        return test_case


Lysandre Debut's avatar
Lysandre Debut committed
462
463
464
465
def require_torch_tpu(test_case):
    """
    Decorator marking a test that requires a TPU (in PyTorch).
    """
466
    if not is_torch_tpu_available():
Lysandre Debut's avatar
Lysandre Debut committed
467
        return unittest.skip("test requires PyTorch TPU")
468
469
    else:
        return test_case
Lysandre Debut's avatar
Lysandre Debut committed
470
471


472
if is_torch_available():
Stas Bekman's avatar
Stas Bekman committed
473
474
475
476
    # Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode
    import torch

    torch_device = "cuda" if torch.cuda.is_available() else "cpu"
477
478
else:
    torch_device = None
479

480
481
482
if is_tf_available():
    import tensorflow as tf

483
484
485
486
487
488
489
if is_flax_available():
    import jax

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

490

491
def require_torch_gpu(test_case):
Patrick von Platen's avatar
Patrick von Platen committed
492
    """Decorator marking a test that requires CUDA and PyTorch."""
493
    if torch_device != "cuda":
494
        return unittest.skip("test requires CUDA")(test_case)
495
496
    else:
        return test_case
497
498


499
def require_torch_bf16(test_case):
500
    """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.10."""
501
    if not is_torch_bf16_available():
502
503
504
505
506
507
508
509
510
        return unittest.skip("test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.10")(test_case)
    else:
        return test_case


def require_torch_tf32(test_case):
    """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7."""
    if not is_torch_tf32_available():
        return unittest.skip("test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7")(test_case)
511
512
513
514
    else:
        return test_case


515
516
517
518
519
520
521
522
def require_detectron2(test_case):
    """Decorator marking a test that requires detectron2."""
    if not is_detectron2_available():
        return unittest.skip("test requires `detectron2`")(test_case)
    else:
        return test_case


Ola Piktus's avatar
Ola Piktus committed
523
524
def require_faiss(test_case):
    """Decorator marking a test that requires faiss."""
525
    if not is_faiss_available():
526
527
528
        return unittest.skip("test requires `faiss`")(test_case)
    else:
        return test_case
Ola Piktus's avatar
Ola Piktus committed
529
530


531
532
533
534
535
536
537
def require_optuna(test_case):
    """
    Decorator marking a test that requires optuna.

    These tests are skipped when optuna isn't installed.

    """
538
    if not is_optuna_available():
539
540
541
542
543
544
545
546
547
548
549
550
        return unittest.skip("test requires optuna")(test_case)
    else:
        return test_case


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

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

    """
551
    if not is_ray_available():
552
553
554
555
556
        return unittest.skip("test requires Ray/tune")(test_case)
    else:
        return test_case


557
558
559
560
561
562
563
564
565
566
567
568
569
def require_sigopt(test_case):
    """
    Decorator marking a test that requires SigOpt.

    These tests are skipped when SigOpt isn't installed.

    """
    if not is_sigopt_available():
        return unittest.skip("test requires SigOpt")(test_case)
    else:
        return test_case


Patrick von Platen's avatar
Patrick von Platen committed
570
571
572
573
574
575
576
577
578
579
580
581
582
def require_soundfile(test_case):
    """
    Decorator marking a test that requires soundfile

    These tests are skipped when soundfile isn't installed.

    """
    if not is_soundfile_availble():
        return unittest.skip("test requires soundfile")(test_case)
    else:
        return test_case


583
584
585
586
587
588
589
590
591
592
def require_deepspeed(test_case):
    """
    Decorator marking a test that requires deepspeed
    """
    if not is_deepspeed_available():
        return unittest.skip("test requires deepspeed")(test_case)
    else:
        return test_case


593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
def require_pyctcdecode(test_case):
    """
    Decorator marking a test that requires pyctcdecode
    """
    if not is_pyctcdecode_available():
        return unittest.skip("test requires pyctcdecode")(test_case)
    else:
        return test_case


def require_librosa(test_case):
    """
    Decorator marking a test that requires librosa
    """
    if not is_librosa_available():
        return unittest.skip("test requires librosa")(test_case)
    else:
        return test_case


613
614
def get_gpu_count():
    """
Suraj Patil's avatar
Suraj Patil committed
615
    Return the number of available gpus (regardless of whether torch, tf or jax is used)
616
    """
617
    if is_torch_available():
618
619
620
        import torch

        return torch.cuda.device_count()
621
    elif is_tf_available():
622
623
624
        import tensorflow as tf

        return len(tf.config.list_physical_devices("GPU"))
Suraj Patil's avatar
Suraj Patil committed
625
626
627
628
    elif is_flax_available():
        import jax

        return jax.device_count()
629
630
631
632
    else:
        return 0


633
def get_tests_dir(append_path=None):
634
    """
635
636
637
638
    Args:
        append_path: optional path to append to the tests dir path

    Return:
Sylvain Gugger's avatar
Sylvain Gugger committed
639
640
        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.
641

642
643
644
    """
    # this function caller's __file__
    caller__file__ = inspect.stack()[1][1]
645
646
647
648
649
    tests_dir = os.path.abspath(os.path.dirname(caller__file__))
    if append_path:
        return os.path.join(tests_dir, append_path)
    else:
        return tests_dir
650
651


652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
#
# Helper functions for dealing with testing text outputs
# The original code came from:
# https://github.com/fastai/fastai/blob/master/tests/utils/text.py

# 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
677
678
    """
    Context manager to capture:
679

680
681
        - stdout: replay it, clean it up and make it available via ``obj.out``
        - stderr: replay it and make it available via ``obj.err``
682
683
684

        init arguments:

685
686
687
688
689
        - out - capture stdout:`` True``/``False``, default ``True``
        - err - capture stdout: ``True``/``False``, default ``True``
        - replay - whether to replay or not: ``True``/``False``, default ``True``. By default each
        captured stream gets replayed back on context's exit, so that one can see what the test was doing. If this is a
        not wanted behavior and the captured data shouldn't be replayed, pass ``replay=False`` to disable this feature.
690

Sylvain Gugger's avatar
Sylvain Gugger committed
691
        Examples::
692

693
            # to capture stdout only with auto-replay
Sylvain Gugger's avatar
Sylvain Gugger committed
694
695
            with CaptureStdout() as cs:
                print("Secret message")
696
            assert "message" in cs.out
697

698
            # to capture stderr only with auto-replay
Sylvain Gugger's avatar
Sylvain Gugger committed
699
700
701
            import sys
            with CaptureStderr() as cs:
                print("Warning: ", file=sys.stderr)
702
            assert "Warning" in cs.err
703

704
705
706
707
708
709
710
711
            # 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
Sylvain Gugger's avatar
Sylvain Gugger committed
712
713
            with CaptureStd(err=False) as cs:
                print("Secret message")
714
            assert "message" in cs.out
Sylvain Gugger's avatar
Sylvain Gugger committed
715
            # but best use the stream-specific subclasses
716

717
718
719
720
721
            # to capture without auto-replay
            with CaptureStd(replay=False) as cs:
                print("Secret message")
            assert "message" in cs.out

722
723
    """

724
725
726
727
    def __init__(self, out=True, err=True, replay=True):

        self.replay = replay

728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
        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
756
757
758
759
            captured = self.out_buf.getvalue()
            if self.replay:
                sys.stdout.write(captured)
            self.out = apply_print_resets(captured)
760
761
762

        if self.err_buf:
            sys.stderr = self.err_old
763
764
765
766
            captured = self.err_buf.getvalue()
            if self.replay:
                sys.stderr.write(captured)
            self.err = captured
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783

    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
784
    """Same as CaptureStd but captures only stdout"""
785

786
787
    def __init__(self, replay=True):
        super().__init__(err=False, replay=replay)
788
789
790


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

793
794
    def __init__(self, replay=True):
        super().__init__(out=False, replay=replay)
795
796


797
class CaptureLogger:
Sylvain Gugger's avatar
Sylvain Gugger committed
798
799
    """
    Context manager to capture `logging` streams
800
801

    Args:
802

803
804
805
806
807
    - logger: 'logging` logger object

    Results:
        The captured output is available via `self.out`

808
    Example::
809

810
811
        >>> from transformers import logging
        >>> from transformers.testing_utils import CaptureLogger
812

813
814
        >>> msg = "Testing 1, 2, 3"
        >>> logging.set_verbosity_info()
815
        >>> logger = logging.get_logger("transformers.models.bart.tokenization_bart")
816
817
818
        >>> with CaptureLogger(logger) as cl:
        ...     logger.info(msg)
        >>> assert cl.out, msg+"\n"
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
    """

    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"


839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
@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.

    For example ::

        with LoggingLevel(logging.INFO):
            AutoModel.from_pretrained("gpt2") # calls logger.info() several times

    """
    orig_level = transformers_logging.get_verbosity()
    try:
        transformers_logging.set_verbosity(level)
        yield
    finally:
        transformers_logging.set_verbosity(orig_level)


859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
@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`.

    Usage ::

       with ExtendSysPath('/path/to/dir'):
           mymodule = importlib.import_module('mymodule')

    """

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


880
class TestCasePlus(unittest.TestCase):
Sylvain Gugger's avatar
Sylvain Gugger committed
881
882
    """
    This class extends `unittest.TestCase` with additional features.
883

884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
    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:

    * ``pathlib`` objects (all fully resolved):

       - ``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)

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

       - ``test_file_path_str``
       - ``test_file_dir_str``
       - ``tests_dir_str``
       - ``examples_dir_str``
       - ``repo_root_dir_str``
       - ``src_dir_str``

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

910
    1. Create a unique temporary dir:
911
912
913
914
915

    ::

        def test_whatever(self):
            tmp_dir = self.get_auto_remove_tmp_dir()
916

917
918
919
920
921
922
    ``tmp_dir`` will contain the path to the created temporary dir. It will be automatically removed at the end of the
    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.
923
924
925
926

    ::

        def test_whatever(self):
927
            tmp_dir = self.get_auto_remove_tmp_dir("./xxx")
928

929
930
    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.
931

932
933
    3. You can override the first two options by directly overriding the ``before`` and ``after`` args, leading to the
       following behavior:
934

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

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

939
940
941
    ``after=True``: the temporary dir will always be deleted at the end of the test.

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

943
944
945
    Note 1: In order to run the equivalent of ``rm -r`` safely, only subdirs of the project repository checkout are
    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 ``./``
946

947
948
    Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested
    otherwise.
949

950
951
952
953
954
955
956
957
    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.


    ::
        def test_whatever(self):
            env = self.get_env()

958
959
960
    """

    def setUp(self):
961
        # get_auto_remove_tmp_dir feature:
962
963
        self.teardown_tmp_dirs = []

964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
        # 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):
        """
        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.

        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).

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

1049
    def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None):
1050
1051
        """
        Args:
1052
            tmp_dir (:obj:`string`, `optional`):
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
                if :obj:`None`:

                   - a unique temporary path will be created
                   - sets ``before=True`` if ``before`` is :obj:`None`
                   - sets ``after=True`` if ``after`` is :obj:`None`
                else:

                   - :obj:`tmp_dir` will be created
                   - sets ``before=True`` if ``before`` is :obj:`None`
                   - sets ``after=False`` if ``after`` is :obj:`None`
            before (:obj:`bool`, `optional`):
                If :obj:`True` and the :obj:`tmp_dir` already exists, make sure to empty it right away if :obj:`False`
                and the :obj:`tmp_dir` already exists, any existing files will remain there.
            after (:obj:`bool`, `optional`):
                If :obj:`True`, delete the :obj:`tmp_dir` at the end of the test if :obj:`False`, leave the
                :obj:`tmp_dir` and its contents intact at the end of the test.
1069
1070

        Returns:
1071
            tmp_dir(:obj:`string`): either the same value as passed via `tmp_dir` or the path to the auto-selected tmp
Sylvain Gugger's avatar
Sylvain Gugger committed
1072
            dir
1073
1074
        """
        if tmp_dir is not None:
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084

            # 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

1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
            # 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:
1101
1102
1103
1104
1105
1106
1107
1108
1109
            # 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

1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
            # 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

    def tearDown(self):
1120
1121

        # get_auto_remove_tmp_dir feature: remove registered temp dirs
1122
1123
1124
        for path in self.teardown_tmp_dirs:
            shutil.rmtree(path, ignore_errors=True)
        self.teardown_tmp_dirs = []
1125
1126
1127


def mockenv(**kwargs):
Sylvain Gugger's avatar
Sylvain Gugger committed
1128
    """
1129
1130
1131
1132
1133
1134
    this is a convenience wrapper, that allows this ::

    @mockenv(RUN_SLOW=True, USE_TF=False)
    def test_something():
        run_slow = os.getenv("RUN_SLOW", False)
        use_tf = os.getenv("USE_TF", False)
1135
1136

    """
1137
    return mock.patch.dict(os.environ, kwargs)
1138
1139


1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
# from https://stackoverflow.com/a/34333710/9201239
@contextlib.contextmanager
def mockenv_context(*remove, **update):
    """
    Temporarily updates the ``os.environ`` dictionary in-place. Similar to mockenv

    The ``os.environ`` dictionary is updated in-place so that the modification is sure to work in all situations.

    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]


1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
# --- 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


1197
1198
def pytest_terminal_summary_main(tr, id):
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
1199
1200
    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.
1201
1202
1203

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

Sylvain Gugger's avatar
Sylvain Gugger committed
1204
1205
    This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined
    there.
1206
1207

    Args:
1208

1209
    - tr: `terminalreporter` passed from `conftest.py`
1210
1211
    - 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.
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229

    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.

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

    dir = "reports"
    Path(dir).mkdir(parents=True, exist_ok=True)
Stas Bekman's avatar
Stas Bekman committed
1230
    report_files = {
1231
        k: f"{dir}/{id}_{k}.txt"
Stas Bekman's avatar
Stas Bekman committed
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
        for k in [
            "durations",
            "errors",
            "failures_long",
            "failures_short",
            "failures_line",
            "passes",
            "stats",
            "summary_short",
            "warnings",
        ]
    }
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263

    # 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
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
    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

1278
1279
1280
1281
    # 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
1282
1283
1284
1285

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

Stas Bekman's avatar
Stas Bekman committed
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
    # 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:
1300
1301
1302
        tr._tw = create_terminal_writer(config, f)
        tr.summary_errors()

Stas Bekman's avatar
Stas Bekman committed
1303
    with open(report_files["warnings"], "w") as f:
1304
1305
1306
1307
        tr._tw = create_terminal_writer(config, f)
        tr.summary_warnings()  # normal warnings
        tr.summary_warnings()  # final warnings

Stas Bekman's avatar
Stas Bekman committed
1308
1309
    tr.reportchars = "wPpsxXEf"  # emulate -rA (used in summary_passes() and short_test_summary())
    with open(report_files["passes"], "w") as f:
1310
1311
1312
        tr._tw = create_terminal_writer(config, f)
        tr.summary_passes()

Stas Bekman's avatar
Stas Bekman committed
1313
    with open(report_files["summary_short"], "w") as f:
1314
1315
1316
        tr._tw = create_terminal_writer(config, f)
        tr.short_test_summary()

Stas Bekman's avatar
Stas Bekman committed
1317
    with open(report_files["stats"], "w") as f:
1318
1319
1320
1321
1322
1323
1324
        tr._tw = create_terminal_writer(config, f)
        tr.summary_stats()

    # restore:
    tr._tw = orig_writer
    tr.reportchars = orig_reportchars
    config.option.tbstyle = orig_tbstyle
1325
1326


1327
# --- distributed testing functions --- #
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381

# 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
1382
            _read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:")),
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
            _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:
1399
        stderr = "\n".join(result.stderr)
1400
        raise RuntimeError(
1401
1402
            f"'{cmd_str}' failed with returncode {result.returncode}\n\n"
            f"The combined stderr from workers follows:\n{stderr}"
1403
        )
Stas Bekman's avatar
Stas Bekman committed
1404
1405
1406
1407

    # 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:
1408
1409
1410
        raise RuntimeError(f"'{cmd_str}' produced no output.")

    return result
1411
1412


1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
def pytest_xdist_worker_id():
    """
    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.
    """
    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():
    """
    Returns a port number that can be fed to ``torch.distributed.launch``'s ``--master_port`` argument.

    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.
    """
    port = 29500
    uniq_delta = pytest_xdist_worker_id()
    return port + uniq_delta


1435
1436
1437
1438
1439
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.
    """
1440
1441
    import numpy as np

1442
1443
1444
1445
    from transformers.tokenization_utils import BatchEncoding

    if isinstance(obj, list):
        return [nested_simplify(item, decimals) for item in obj]
1446
1447
    elif isinstance(obj, np.ndarray):
        return nested_simplify(obj.tolist())
1448
1449
    elif isinstance(obj, (dict, BatchEncoding)):
        return {nested_simplify(k, decimals): nested_simplify(v, decimals) for k, v in obj.items()}
1450
    elif isinstance(obj, (str, int, np.int64)):
1451
        return obj
1452
1453
    elif obj is None:
        return obj
1454
    elif is_torch_available() and isinstance(obj, torch.Tensor):
1455
        return nested_simplify(obj.tolist(), decimals)
1456
1457
1458
1459
    elif is_tf_available() and tf.is_tensor(obj):
        return nested_simplify(obj.numpy().tolist())
    elif isinstance(obj, float):
        return round(obj, decimals)
1460
    elif isinstance(obj, (np.int32, np.float32)):
1461
        return nested_simplify(obj.item(), decimals)
1462
1463
    else:
        raise Exception(f"Not supported: {type(obj)}")