diff --git "a/third_party/transformers/src/transformers/testing_utils.py" "b/third_party/transformers/src/transformers/testing_utils.py" new file mode 100644--- /dev/null +++ "b/third_party/transformers/src/transformers/testing_utils.py" @@ -0,0 +1,4401 @@ +# 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. + +import ast +import collections +import contextlib +import copy +import doctest +import functools +import gc +import importlib +import inspect +import json +import logging +import multiprocessing +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +import threading +import time +import traceback +import types +import unittest +from collections import UserDict, defaultdict +from collections.abc import Callable, Generator, Iterable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import MISSING, fields +from functools import cache, wraps +from io import StringIO +from pathlib import Path +from typing import TYPE_CHECKING, Any +from unittest import mock +from unittest.mock import patch + +import httpx +from huggingface_hub import create_repo, delete_repo +from packaging import version + +from transformers import logging as transformers_logging + + +if TYPE_CHECKING: + from .trainer import Trainer +else: + Trainer = Any # type: ignore + +from .integrations import ( + is_clearml_available, + is_optuna_available, + is_ray_available, + is_swanlab_available, + is_tensorboard_available, + is_trackio_available, + is_wandb_available, +) +from .integrations.deepspeed import is_deepspeed_available +from .utils import ( + ACCELERATE_MIN_VERSION, + GGUF_MIN_VERSION, + SAFE_WEIGHTS_INDEX_NAME, + TRITON_MIN_VERSION, + WEIGHTS_INDEX_NAME, + is_accelerate_available, + is_apex_available, + is_apollo_torch_available, + is_aqlm_available, + is_auto_round_available, + is_av_available, + is_bitsandbytes_available, + is_bs4_available, + is_compressed_tensors_available, + is_cv2_available, + is_cython_available, + is_decord_available, + is_detectron2_available, + is_essentia_available, + is_faiss_available, + is_fbgemm_gpu_available, + is_flash_attn_2_available, + is_flash_attn_3_available, + is_flash_attn_4_available, + is_flute_available, + is_fouroversix_available, + is_fp_quant_available, + is_fsdp_available, + is_g2p_en_available, + is_galore_torch_available, + is_gguf_available, + is_gptqmodel_available, + is_grokadamw_available, + is_hadamard_available, + is_hqq_available, + is_huggingface_hub_greater_or_equal, + is_jinja_available, + is_jmespath_available, + is_jumanpp_available, + is_kernels_available, + is_levenshtein_available, + is_librosa_available, + is_liger_kernel_available, + is_lomo_available, + is_mistral_common_available, + is_multipart_available, + is_natten_available, + is_nltk_available, + is_numba_available, + is_onnx_available, + is_openai_available, + is_optimum_available, + is_optimum_quanto_available, + is_pandas_available, + is_peft_available, + is_phonemizer_available, + is_pretty_midi_available, + is_psutil_available, + is_pyctcdecode_available, + is_pytesseract_available, + is_pytest_available, + is_pytest_order_available, + is_pytorch_quantization_available, + is_quark_available, + is_qutlass_available, + is_rjieba_available, + is_sacremoses_available, + is_schedulefree_available, + is_scipy_available, + is_sentencepiece_available, + is_seqio_available, + is_serve_available, + is_spacy_available, + is_speech_available, + is_spqr_available, + is_sudachi_available, + is_sudachi_projection_available, + is_tiktoken_available, + is_timm_available, + is_tokenizers_available, + is_torch_available, + is_torch_bf16_available_on_device, + is_torch_fp16_available_on_device, + is_torch_greater_or_equal, + is_torch_hpu_available, + is_torch_mlu_available, + is_torch_neuroncore_available, + is_torch_npu_available, + is_torch_optimi_available, + is_torch_tensorrt_fx_available, + is_torch_tf32_available, + is_torch_xla_available, + is_torch_xpu_available, + is_torchao_available, + is_torchaudio_available, + is_torchcodec_available, + is_torchvision_available, + is_triton_available, + is_vision_available, + is_vptq_available, + strtobool, +) + + +if is_accelerate_available(): + from accelerate.state import AcceleratorState, PartialState + from accelerate.utils.imports import is_fp8_available + + +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 + from _pytest.pathlib import import_path + from pytest import DoctestItem +else: + Module = object + DoctestItem = object + + +SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy" +DUMMY_UNKNOWN_IDENTIFIER = "julien-c/dummy-unknown" +DUMMY_DIFF_TOKENIZER_IDENTIFIER = "julien-c/dummy-diff-tokenizer" +# Used to test Auto{Config, Model, Tokenizer} model_type detection. + +# Used to test the hub +USER = "__DUMMY_TRANSFORMERS_USER__" +ENDPOINT_STAGING = "https://hub-ci.huggingface.co" + +# Not critical, only usable on the sandboxed CI instance. +TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL" + + +# Used in CausalLMModelTester (and related classes/methods) to infer the common model classes from the base model class +_COMMON_MODEL_NAMES_MAP = { + "config_class": "Config", + "causal_lm_class": "ForCausalLM", + "question_answering_class": "ForQuestionAnswering", + "sequence_classification_class": "ForSequenceClassification", + "token_classification_class": "ForTokenClassification", +} + +# Used in VLMModelTester (and related classes/methods) to infer the common model classes from the base model class +_VLM_COMMON_MODEL_NAMES_MAP = { + "config_class": "Config", + "text_config_class": "TextConfig", + "vision_config_class": "VisionConfig", + "conditional_generation_class": "ForConditionalGeneration", +} + + +if is_torch_available(): + import torch + from safetensors.torch import load_file + + from .modeling_utils import FLASH_ATTN_KERNEL_FALLBACK, PreTrainedModel + + IS_ROCM_SYSTEM = torch.version.hip is not None + IS_CUDA_SYSTEM = torch.version.cuda is not None + IS_XPU_SYSTEM = getattr(torch.version, "xpu", None) is not None + IS_NPU_SYSTEM = getattr(torch, "npu", None) is not None +else: + IS_ROCM_SYSTEM = False + IS_CUDA_SYSTEM = False + IS_XPU_SYSTEM = False + IS_NPU_SYSTEM = False + +logger = transformers_logging.get_logger(__name__) + + +def parse_flag_from_env(key, default=False): + try: + 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. + raise ValueError(f"If set, {key} must be yes or no.") + return _value + + +def parse_int_from_env(key, default=None): + try: + value = os.environ[key] + except KeyError: + _value = default + else: + try: + _value = int(value) + except ValueError: + raise ValueError(f"If set, {key} must be a int.") + return _value + + +_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False) +_run_flaky_tests = parse_flag_from_env("RUN_FLAKY", default=True) +_run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False) +_run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False) +_run_pipeline_tests = parse_flag_from_env("RUN_PIPELINE_TESTS", default=True) +_run_agent_tests = parse_flag_from_env("RUN_AGENT_TESTS", default=False) +_run_training_tests = parse_flag_from_env("RUN_TRAINING_TESTS", default=True) +_run_tensor_parallel_tests = parse_flag_from_env("RUN_TENSOR_PARALLEL_TESTS", default=True) + + +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(reason="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) + + +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(reason="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) + + +def is_agent_test(test_case): + """ + Decorator marking a test as an agent test. If RUN_TOOL_TESTS is set to a falsy value, those tests will be skipped. + """ + if not _run_agent_tests: + return unittest.skip(reason="test is an agent 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_agent_test()(test_case) + + +def is_training_test(test_case): + """ + Decorator marking a test as a training test. If RUN_TRAINING_TESTS is set to a falsy value, those tests will be + skipped. + """ + if not _run_training_tests: + return unittest.skip(reason="test is training 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_training_test()(test_case) + + +def is_tensor_parallel_test(test_case): + """ + Decorator marking a test as a tensor parallel test. If RUN_TENSOR_PARALLEL_TESTS is set to a falsy value, those + tests will be skipped. + """ + if not _run_tensor_parallel_tests: + return unittest.skip(reason="test is tensor parallel 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_tensor_parallel_test()(test_case) + + +def slow(test_case): + """ + Decorator marking a test as slow. + + Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them. + + """ + return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case) + + +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(reason="test is too slow")(test_case) + + +def skip_if_not_implemented(test_func): + @functools.wraps(test_func) + def wrapper(*args, **kwargs): + try: + return test_func(*args, **kwargs) + except NotImplementedError as e: + raise unittest.SkipTest(f"Test skipped due to NotImplementedError: {e}") + + return wrapper + + +def apply_skip_if_not_implemented(cls): + """ + Class decorator to apply @skip_if_not_implemented to all test methods. + """ + for attr_name in dir(cls): + if attr_name.startswith("test_"): + attr = getattr(cls, attr_name) + if callable(attr): + setattr(cls, attr_name, skip_if_not_implemented(attr)) + return cls + + +def custom_tokenizers(test_case): + """ + Decorator marking a test for a custom tokenizer. + + Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS + environment variable to a truthy value to run them. + """ + return unittest.skipUnless(_run_custom_tokenizers, "test of custom tokenizers")(test_case) + + +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) + + +def require_galore_torch(test_case): + """ + Decorator marking a test that requires GaLore. These tests are skipped when GaLore isn't installed. + https://github.com/jiaweizzhao/GaLore + """ + return unittest.skipUnless(is_galore_torch_available(), "test requires GaLore")(test_case) + + +def require_apollo_torch(test_case): + """ + Decorator marking a test that requires GaLore. These tests are skipped when APOLLO isn't installed. + https://github.com/zhuhanqing/APOLLO + """ + return unittest.skipUnless(is_apollo_torch_available(), "test requires APOLLO")(test_case) + + +def require_torch_optimi(test_case): + """ + Decorator marking a test that requires torch-optimi. These tests are skipped when torch-optimi isn't installed. + https://github.com/jxnl/torch-optimi + """ + return unittest.skipUnless(is_torch_optimi_available(), "test requires torch-optimi")(test_case) + + +def require_lomo(test_case): + """ + Decorator marking a test that requires LOMO. These tests are skipped when LOMO-optim isn't installed. + https://github.com/OpenLMLab/LOMO + """ + return unittest.skipUnless(is_lomo_available(), "test requires LOMO")(test_case) + + +def require_grokadamw(test_case): + """ + Decorator marking a test that requires GrokAdamW. These tests are skipped when GrokAdamW isn't installed. + """ + return unittest.skipUnless(is_grokadamw_available(), "test requires GrokAdamW")(test_case) + + +def require_schedulefree(test_case): + """ + Decorator marking a test that requires schedulefree. These tests are skipped when schedulefree isn't installed. + https://github.com/facebookresearch/schedule_free + """ + return unittest.skipUnless(is_schedulefree_available(), "test requires schedulefree")(test_case) + + +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) + + +def require_accelerate(test_case, min_version: str = ACCELERATE_MIN_VERSION): + """ + Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed. + """ + return unittest.skipUnless( + is_accelerate_available(min_version), f"test requires accelerate version >= {min_version}" + )(test_case) + + +def require_triton(min_version: str = TRITON_MIN_VERSION): + """ + Decorator marking a test that requires triton. These tests are skipped when triton isn't installed. + """ + + def decorator(test_case): + return unittest.skipUnless(is_triton_available(min_version), f"test requires triton version >= {min_version}")( + test_case + ) + + return decorator + + +def require_gguf(test_case, min_version: str = GGUF_MIN_VERSION): + """ + Decorator marking a test that requires ggguf. These tests are skipped when gguf isn't installed. + """ + return unittest.skipUnless(is_gguf_available(min_version), f"test requires gguf version >= {min_version}")( + test_case + ) + + +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 + ) + + +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) + + +def require_rjieba(test_case): + """ + Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed. + """ + return unittest.skipUnless(is_rjieba_available(), "test requires rjieba")(test_case) + + +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) + + +def require_jmespath(test_case): + """ + Decorator marking a test that requires jmespath. These tests are skipped when jmespath isn't installed. + """ + return unittest.skipUnless(is_jmespath_available(), "test requires jmespath")(test_case) + + +def require_onnx(test_case): + return unittest.skipUnless(is_onnx_available(), "test requires ONNX")(test_case) + + +def require_timm(test_case): + """ + Decorator marking a test that requires Timm. + + These tests are skipped when Timm isn't installed. + + """ + return unittest.skipUnless(is_timm_available(), "test requires Timm")(test_case) + + +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) + + +def require_torch(test_case): + """ + Decorator marking a test that requires PyTorch. + + These tests are skipped when PyTorch isn't installed. + + """ + return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case) + + +def require_torch_greater_or_equal(version: str): + """ + Decorator marking a test that requires PyTorch version >= `version`. + + These tests are skipped when PyTorch version is less than `version`. + """ + + def decorator(test_case): + return unittest.skipUnless(is_torch_greater_or_equal(version), f"test requires PyTorch version >= {version}")( + test_case + ) + + return decorator + + +def require_huggingface_hub_greater_or_equal(version: str): + """ + Decorator marking a test that requires huggingface_hub version >= `version`. + + These tests are skipped when huggingface_hub version is less than `version`. + """ + + def decorator(test_case): + return unittest.skipUnless( + is_huggingface_hub_greater_or_equal(version), f"test requires huggingface_hub version >= {version}" + )(test_case) + + return decorator + + +def require_flash_attn(test_case): + """ + Decorator marking a test that requires Flash Attention. + + These tests are skipped when Flash Attention isn't installed. + + """ + flash_attn_available = is_flash_attn_2_available() + kernels_available = is_kernels_available() + try: + from kernels import get_kernel + + get_kernel(FLASH_ATTN_KERNEL_FALLBACK["flash_attention_2"]) + except Exception as _: + kernels_available = False + + return unittest.skipUnless(kernels_available | flash_attn_available, "test requires Flash Attention")(test_case) + + +def require_kernels(test_case): + """ + Decorator marking a test that requires the kernels library. + + These tests are skipped when the kernels library isn't installed. + + """ + return unittest.skipUnless(is_kernels_available(), "test requires the kernels library")(test_case) + + +def require_flash_attn_3(test_case): + """ + Decorator marking a test that requires Flash Attention 3. + + These tests are skipped when Flash Attention 3 isn't installed. + """ + return unittest.skipUnless(is_flash_attn_3_available(), "test requires Flash Attention 3")(test_case) + + +def require_flash_attn_4(test_case): + """ + Decorator marking a test that requires Flash Attention 4. + + These tests are skipped when Flash Attention 4 isn't installed. + """ + return unittest.skipUnless(is_flash_attn_4_available(), "test requires Flash Attention 4")(test_case) + + +def require_all_flash_attn(test_case): + flash_attn_available = is_flash_attn_2_available() + kernels_available = is_kernels_available() + try: + from kernels import get_kernel + + get_kernel(FLASH_ATTN_KERNEL_FALLBACK["flash_attention_2"]) + except Exception as _: + kernels_available = False + + return unittest.skipUnless( + all( + ( + flash_attn_available | kernels_available, + is_flash_attn_3_available(), + is_flash_attn_4_available(), + ) + ), + "test requires all mainline Flash Attention packages", + )(test_case) + + +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) + + +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) + + +def require_torchcodec(test_case): + """ + Decorator marking a test that requires Torchcodec. + + These tests are skipped when Torchcodec isn't installed. + + """ + return unittest.skipUnless(is_torchcodec_available(), "test requires Torchcodec")(test_case) + + +def require_torchaudio(test_case): + """ + Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed. + """ + return unittest.skipUnless(is_torchaudio_available(), "test requires torchaudio")(test_case) + + +def require_sentencepiece(test_case): + """ + Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed. + """ + return unittest.skipUnless(is_sentencepiece_available(), "test requires SentencePiece")(test_case) + + +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) + + +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) + + +def require_scipy(test_case): + """ + Decorator marking a test that requires Scipy. These tests are skipped when SentencePiece isn't installed. + """ + return unittest.skipUnless(is_scipy_available(), "test requires Scipy")(test_case) + + +def require_tokenizers(test_case): + """ + Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed. + """ + return unittest.skipUnless(is_tokenizers_available(), "test requires tokenizers")(test_case) + + +def require_pandas(test_case): + """ + Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed. + """ + return unittest.skipUnless(is_pandas_available(), "test requires pandas")(test_case) + + +def require_pytesseract(test_case): + """ + Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed. + """ + return unittest.skipUnless(is_pytesseract_available(), "test requires PyTesseract")(test_case) + + +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. + """ + return unittest.skipUnless(is_pytorch_quantization_available(), "test requires PyTorch Quantization Toolkit")( + test_case + ) + + +def require_vision(test_case): + """ + Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't + installed. + """ + return unittest.skipUnless(is_vision_available(), "test requires vision")(test_case) + + +def require_spacy(test_case): + """ + Decorator marking a test that requires SpaCy. These tests are skipped when SpaCy isn't installed. + """ + return unittest.skipUnless(is_spacy_available(), "test requires spacy")(test_case) + + +def require_torch_multi_gpu(test_case): + """ + Decorator marking a test that requires a multi-GPU CUDA setup (in PyTorch). These tests are skipped on a machine without + multiple CUDA GPUs. + + To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu" + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + import torch + + return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple CUDA GPUs")(test_case) + + +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(reason="test requires PyTorch")(test_case) + + return unittest.skipUnless(backend_device_count(torch_device) > 1, "test requires multiple accelerators")( + test_case + ) + + +def require_torch_non_multi_gpu(test_case): + """ + Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch). + """ + if not is_torch_available(): + return unittest.skip(reason="test requires PyTorch")(test_case) + + import torch + + return unittest.skipUnless(torch.cuda.device_count() < 2, "test requires 0 or 1 GPU")(test_case) + + +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(reason="test requires PyTorch")(test_case) + + return unittest.skipUnless(backend_device_count(torch_device) < 2, "test requires 0 or 1 accelerator")(test_case) + + +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(reason="test requires PyTorch")(test_case) + + import torch + + return unittest.skipUnless(torch.cuda.device_count() < 3, "test requires 0 or 1 or 2 GPUs")(test_case) + + +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(reason="test requires PyTorch")(test_case) + + return unittest.skipUnless(backend_device_count(torch_device) < 3, "test requires 0 or 1 or 2 accelerators")( + test_case + ) + + +def require_torch_xla(test_case): + """ + Decorator marking a test that requires TorchXLA (in PyTorch). + """ + return unittest.skipUnless(is_torch_xla_available(), "test requires TorchXLA")(test_case) + + +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 + ) + + +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(reason="test requires PyTorch NPU")(test_case) + + return unittest.skipUnless(torch.npu.device_count() > 1, "test requires multiple NPUs")(test_case) + + +def require_non_hpu(test_case): + """ + Decorator marking a test that should be skipped for HPU. + """ + return unittest.skipUnless(torch_device != "hpu", "test requires a non-HPU")(test_case) + + +def require_torch_xpu(test_case): + """ + Decorator marking a test that requires XPU (in PyTorch). + + These tests are skipped when XPU backend is not available. + """ + return unittest.skipUnless(is_torch_xpu_available(), "test requires XPU device")(test_case) + + +def require_non_xpu(test_case): + """ + Decorator marking a test that should be skipped for XPU. + """ + return unittest.skipUnless(torch_device != "xpu", "test requires a non-XPU")(test_case) + + +def require_torch_multi_xpu(test_case): + """ + Decorator marking a test that requires a multi-XPU setup (in PyTorch). These tests are skipped on a machine without + 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(reason="test requires PyTorch XPU")(test_case) + + return unittest.skipUnless(torch.xpu.device_count() > 1, "test requires multiple XPUs")(test_case) + + +def require_torch_multi_hpu(test_case): + """ + Decorator marking a test that requires a multi-HPU setup (in PyTorch). These tests are skipped on a machine without + multiple HPUs. + + To run *only* the multi_hpu tests, assuming all test names contain multi_hpu: $ pytest -sv ./tests -k "multi_hpu" + """ + if not is_torch_hpu_available(): + return unittest.skip(reason="test requires PyTorch HPU")(test_case) + + return unittest.skipUnless(torch.hpu.device_count() > 1, "test requires multiple HPUs")(test_case) + + +if is_torch_available(): + # Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode + import torch + + 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 + + if "TRANSFORMERS_TEST_DEVICE" in os.environ: + torch_device = os.environ["TRANSFORMERS_TEST_DEVICE"] + 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." + ) + if torch_device == "mlu" and not is_torch_mlu_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but MLU is unavailable. Please double-check your testing environment." + ) + if torch_device == "hpu" and not is_torch_hpu_available(): + raise ValueError( + f"TRANSFORMERS_TEST_DEVICE={torch_device}, but HPU is unavailable. Please double-check your testing environment." + ) + + 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(): + torch_device = "cuda" + elif is_torch_npu_available(): + torch_device = "npu" + elif is_torch_mlu_available(): + torch_device = "mlu" + elif is_torch_hpu_available(): + torch_device = "hpu" + elif is_torch_xpu_available(): + torch_device = "xpu" + else: + torch_device = "cpu" +else: + torch_device = None + + +def require_torchao(test_case): + """Decorator marking a test that requires torchao""" + return unittest.skipUnless(is_torchao_available(), "test requires torchao")(test_case) + + +def require_torchao_version_greater_or_equal(torchao_version): + def decorator(test_case): + correct_torchao_version = is_torchao_available() and version.parse( + version.parse(importlib.metadata.version("torchao")).base_version + ) >= version.parse(torchao_version) + return unittest.skipUnless( + correct_torchao_version, f"Test requires torchao with the version greater than {torchao_version}." + )(test_case) + + return decorator + + +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) + + +def require_torch_gpu(test_case): + """Decorator marking a test that requires CUDA and PyTorch.""" + return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case) + + +def require_torch_mps(test_case): + """Decorator marking a test that requires CUDA and PyTorch.""" + return unittest.skipUnless(torch_device == "mps", "test requires MPS")(test_case) + + +def require_large_cpu_ram(test_case, memory: float = 80): + """Decorator marking a test that requires a CPU RAM with more than `memory` GiB of memory.""" + if not is_psutil_available(): + return test_case + + import psutil + + return unittest.skipUnless( + psutil.virtual_memory().total / 1024**3 > memory, + f"test requires a machine with more than {memory} GiB of CPU RAM memory", + )(test_case) + + +def require_torch_large_gpu(test_case, memory: float = 20): + """Decorator marking a test that requires a CUDA GPU with more than `memory` GiB of memory.""" + if torch_device != "cuda": + return unittest.skip(reason=f"test requires a CUDA GPU with more than {memory} GiB of memory")(test_case) + + return unittest.skipUnless( + torch.cuda.get_device_properties(0).total_memory / 1024**3 > memory, + f"test requires a GPU with more than {memory} GiB of memory", + )(test_case) + + +def require_torch_large_accelerator(test_case=None, *, memory: float = 20): + """Decorator marking a test that requires an accelerator with more than `memory` GiB of memory.""" + + def memory_decorator(tc): + if torch_device not in ("cuda", "xpu"): + return unittest.skip(f"test requires a GPU or XPU with more than {memory} GiB of memory")(tc) + + torch_accel = getattr(torch, torch_device) + return unittest.skipUnless( + torch_accel.get_device_properties(0).total_memory / 1024**3 > memory, + f"test requires a GPU or XPU with more than {memory} GiB of memory", + )(tc) + + return memory_decorator if test_case is None else memory_decorator(test_case) + + +def require_torch_accelerator(test_case): + """Decorator marking a test that requires an accessible accelerator and PyTorch.""" + return unittest.skipUnless(torch_device is not None and torch_device != "cpu", "test requires accelerator")( + test_case + ) + + +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_fp8(test_case): + """Decorator marking a test that requires supports for fp8""" + return unittest.skipUnless(is_accelerate_available() and is_fp8_available(), "test requires fp8 support")( + test_case + ) + + +def require_cuda_capability_at_least(major, minor): + """Decorator to skip tests when CUDA capability is below the given version.""" + import torch + + if not torch.cuda.is_available(): + return unittest.skip("CUDA not available") + capability = torch.cuda.get_device_capability() + return unittest.skipIf(capability < (major, minor), f"Requires CUDA capability >= {major}.{minor}") + + +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) + + +def require_deterministic_for_xpu(test_case): + @wraps(test_case) + def wrapper(*args, **kwargs): + if is_torch_xpu_available(): + original_state = torch.are_deterministic_algorithms_enabled() + try: + torch.use_deterministic_algorithms(True) + return test_case(*args, **kwargs) + finally: + torch.use_deterministic_algorithms(original_state) + else: + return test_case(*args, **kwargs) + + return wrapper + + +def require_torch_tf32(test_case): + """Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7.""" + return unittest.skipUnless( + is_torch_tf32_available(), "test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7" + )(test_case) + + +def require_detectron2(test_case): + """Decorator marking a test that requires detectron2.""" + return unittest.skipUnless(is_detectron2_available(), "test requires `detectron2`")(test_case) + + +def require_faiss(test_case): + """Decorator marking a test that requires faiss.""" + return unittest.skipUnless(is_faiss_available(), "test requires `faiss`")(test_case) + + +def require_optuna(test_case): + """ + Decorator marking a test that requires optuna. + + These tests are skipped when optuna isn't installed. + + """ + return unittest.skipUnless(is_optuna_available(), "test requires optuna")(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. + + """ + return unittest.skipUnless(is_ray_available(), "test requires Ray/tune")(test_case) + + +def require_swanlab(test_case): + """ + Decorator marking a test that requires swanlab. + + These tests are skipped when swanlab isn't installed. + + """ + return unittest.skipUnless(is_swanlab_available(), "test requires swanlab")(test_case) + + +def require_trackio(test_case): + """ + Decorator marking a test that requires trackio. + + These tests are skipped when trackio isn't installed. + + """ + return unittest.skipUnless(is_trackio_available(), "test requires trackio")(test_case) + + +def require_wandb(test_case): + """ + Decorator marking a test that requires wandb. + + These tests are skipped when wandb isn't installed. + + """ + return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case) + + +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) + + +def require_deepspeed(test_case): + """ + Decorator marking a test that requires deepspeed + """ + return unittest.skipUnless(is_deepspeed_available(), "test requires deepspeed")(test_case) + + +def require_apex(test_case): + """ + Decorator marking a test that requires apex + """ + return unittest.skipUnless(is_apex_available(), "test requires apex")(test_case) + + +def require_aqlm(test_case): + """ + Decorator marking a test that requires aqlm + """ + return unittest.skipUnless(is_aqlm_available(), "test requires aqlm")(test_case) + + +def require_vptq(test_case): + """ + Decorator marking a test that requires vptq + """ + return unittest.skipUnless(is_vptq_available(), "test requires vptq")(test_case) + + +def require_spqr(test_case): + """ + Decorator marking a test that requires spqr + """ + return unittest.skipUnless(is_spqr_available(), "test requires spqr")(test_case) + + +def require_av(test_case): + """ + Decorator marking a test that requires av + """ + return unittest.skipUnless(is_av_available(), "test requires av")(test_case) + + +def require_decord(test_case): + """ + Decorator marking a test that requires decord + """ + return unittest.skipUnless(is_decord_available(), "test requires decord")(test_case) + + +def require_bitsandbytes(test_case): + """ + Decorator marking a test that requires the bitsandbytes library. Will be skipped when the library or its hard dependency torch is not installed. + """ + return unittest.skipUnless(is_bitsandbytes_available(), "test requires bitsandbytes")(test_case) + + +def require_optimum(test_case): + """ + Decorator for optimum dependency + """ + return unittest.skipUnless(is_optimum_available(), "test requires optimum")(test_case) + + +def require_tensorboard(test_case): + """ + Decorator for `tensorboard` dependency + """ + return unittest.skipUnless(is_tensorboard_available(), "test requires tensorboard") + + +def require_gptqmodel(test_case): + """ + Decorator for gptqmodel dependency + """ + return unittest.skipUnless(is_gptqmodel_available(), "test requires gptqmodel")(test_case) + + +def require_hqq(test_case): + """ + Decorator for hqq dependency + """ + return unittest.skipUnless(is_hqq_available(), "test requires hqq")(test_case) + + +def require_auto_round(test_case): + """ + Decorator for auto_round dependency + """ + return unittest.skipUnless(is_auto_round_available(), "test requires autoround")(test_case) + + +def require_optimum_quanto(test_case): + """ + Decorator for quanto dependency + """ + return unittest.skipUnless(is_optimum_quanto_available(), "test requires optimum-quanto")(test_case) + + +def require_compressed_tensors(test_case): + """ + Decorator for compressed_tensors dependency + """ + return unittest.skipUnless(is_compressed_tensors_available(), "test requires compressed_tensors")(test_case) + + +def require_fbgemm_gpu(test_case): + """ + Decorator for fbgemm_gpu dependency + """ + return unittest.skipUnless(is_fbgemm_gpu_available(), "test requires fbgemm-gpu")(test_case) + + +def require_quark(test_case): + """ + Decorator for quark dependency + """ + return unittest.skipUnless(is_quark_available(), "test requires quark")(test_case) + + +def require_flute_hadamard(test_case): + """ + Decorator marking a test that requires higgs and hadamard + """ + return unittest.skipUnless( + is_flute_available() and is_hadamard_available(), "test requires flute and fast_hadamard_transform" + )(test_case) + + +def require_fouroversix(test_case): + """ + Decorator marking a test that requires fouroversix + """ + return unittest.skipUnless(is_fouroversix_available(), "test requires fouroversix")(test_case) + + +def require_fp_quant(test_case): + """ + Decorator marking a test that requires fp_quant and qutlass + """ + return unittest.skipUnless(is_fp_quant_available(), "test requires fp_quant")(test_case) + + +def require_qutlass(test_case): + """ + Decorator marking a test that requires qutlass + """ + return unittest.skipUnless(is_qutlass_available(), "test requires qutlass")(test_case) + + +def require_phonemizer(test_case): + """ + Decorator marking a test that requires phonemizer + """ + return unittest.skipUnless(is_phonemizer_available(), "test requires phonemizer")(test_case) + + +def require_pyctcdecode(test_case): + """ + Decorator marking a test that requires pyctcdecode + """ + return unittest.skipUnless(is_pyctcdecode_available(), "test requires pyctcdecode")(test_case) + + +def require_numba(test_case): + """ + Decorator marking a test that requires numba + """ + return unittest.skipUnless(is_numba_available(), "test requires numba")(test_case) + + +def require_librosa(test_case): + """ + Decorator marking a test that requires librosa + """ + return unittest.skipUnless(is_librosa_available(), "test requires librosa")(test_case) + + +def require_multipart(test_case): + """ + Decorator marking a test that requires python-multipart + """ + return unittest.skipUnless(is_multipart_available(), "test requires python-multipart")(test_case) + + +def require_liger_kernel(test_case): + """ + Decorator marking a test that requires liger_kernel + """ + return unittest.skipUnless(is_liger_kernel_available(), "test requires liger_kernel")(test_case) + + +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) + + +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` + """ + return unittest.skipUnless(cmd_exists("/usr/bin/time"), "test requires /usr/bin/time")(test_case) + + +def require_sudachi(test_case): + """ + Decorator marking a test that requires sudachi + """ + return unittest.skipUnless(is_sudachi_available(), "test requires sudachi")(test_case) + + +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 + ) + + +def require_jumanpp(test_case): + """ + Decorator marking a test that requires jumanpp + """ + return unittest.skipUnless(is_jumanpp_available(), "test requires jumanpp")(test_case) + + +def require_cython(test_case): + """ + Decorator marking a test that requires jumanpp + """ + return unittest.skipUnless(is_cython_available(), "test requires cython")(test_case) + + +def require_tiktoken(test_case): + """ + Decorator marking a test that requires TikToken. These tests are skipped when TikToken isn't installed. + """ + return unittest.skipUnless(is_tiktoken_available(), "test requires TikToken")(test_case) + + +def require_speech(test_case): + """ + Decorator marking a test that requires speech. These tests are skipped when speech isn't available. + """ + return unittest.skipUnless(is_speech_available(), "test requires torchaudio")(test_case) + + +def require_openai(test_case): + """ + Decorator marking a test that requires openai + """ + return unittest.skipUnless(is_openai_available(), "test requires openai")(test_case) + + +def require_serve(test_case): + """ + Decorator marking a test that requires the serving dependencies (fastapi, uvicorn, pydantic, openai). + """ + return unittest.skipUnless(is_serve_available(), "test requires serving dependencies")(test_case) + + +def require_mistral_common(test_case): + """ + Decorator marking a test that requires mistral-common. These tests are skipped when mistral-common isn't available. + """ + return unittest.skipUnless(is_mistral_common_available(), "test requires mistral-common")(test_case) + + +def get_gpu_count(): + """ + Return the number of available gpus + """ + if is_torch_available(): + import torch + + return torch.cuda.device_count() + else: + return 0 + + +def get_tests_dir(append_path=None): + """ + Args: + append_path: optional path to append to the tests dir path + + Return: + 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. + + """ + # this function caller's __file__ + caller__file__ = inspect.stack()[1][1] + tests_dir = os.path.abspath(os.path.dirname(caller__file__)) + + while not tests_dir.endswith("tests"): + tests_dir = os.path.dirname(tests_dir) + + if append_path: + return os.path.join(tests_dir, append_path) + else: + return tests_dir + + +def get_steps_per_epoch(trainer: Trainer) -> int: + training_args = trainer.args + train_dataloader = trainer.get_train_dataloader() + + initial_training_values = trainer.set_initial_training_values(args=training_args, dataloader=train_dataloader) + steps_per_epoch = initial_training_values[5] + + return steps_per_epoch + + +def evaluate_side_effect_factory( + side_effect_values: list[dict[str, float]], +) -> Generator[dict[str, float], None, None]: + """ + Function that returns side effects for the _evaluate method. + Used when we're unsure of exactly how many times _evaluate will be called. + """ + yield from side_effect_values + + while True: + yield side_effect_values[-1] + + +# +# 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.MULTILINE) + + +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}" + + +def set_config_for_less_flaky_test(config): + target_attrs = [ + "rms_norm_eps", + "layer_norm_eps", + "norm_eps", + "norm_epsilon", + "layer_norm_epsilon", + "batch_norm_eps", + ] + for target_attr in target_attrs: + setattr(config, target_attr, 1.0) + + # norm layers (layer/group norm, etc.) could cause flaky tests when the tensors have very small variance. + # (We don't need the original epsilon values to check eager/sdpa matches) + attrs = ["text_config", "vision_config", "audio_config", "text_encoder", "audio_encoder", "decoder"] + for attr in attrs: + if hasattr(config, attr) and getattr(config, attr) is not None: + for target_attr in target_attrs: + setattr(getattr(config, attr), target_attr, 1.0) + + +def set_model_for_less_flaky_test(model): + # Another way to make sure norm layers have desired epsilon. (Some models don't set it from its config.) + target_names = ( + "LayerNorm", + "GroupNorm", + "BatchNorm", + "RMSNorm", + "BatchNorm2d", + "BatchNorm1d", + "BitGroupNormActivation", + "WeightStandardizedConv2d", + ) + target_attrs = ["eps", "epsilon", "variance_epsilon"] + if is_torch_available() and isinstance(model, torch.nn.Module): + for module in model.modules(): + if type(module).__name__.endswith(target_names): + for attr in target_attrs: + if hasattr(module, attr): + setattr(module, attr, 1.0) + + +class CaptureStd: + """ + Context manager to capture: + + - stdout: replay it, clean it up and make it available via `obj.out` + - stderr: replay it and make it available via `obj.err` + + 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 + doing. If this is a not wanted behavior and the captured data shouldn't be replayed, pass `replay=False` to + disable this feature. + + 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 + + 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 + ```""" + + def __init__(self, out=True, err=True, replay=True): + self.replay = replay + + 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 + captured = self.out_buf.getvalue() + if self.replay: + sys.stdout.write(captured) + self.out = apply_print_resets(captured) + + if self.err_buf: + sys.stderr = self.err_old + captured = self.err_buf.getvalue() + if self.replay: + sys.stderr.write(captured) + self.err = captured + + 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): + """Same as CaptureStd but captures only stdout""" + + def __init__(self, replay=True): + super().__init__(err=False, replay=replay) + + +class CaptureStderr(CaptureStd): + """Same as CaptureStd but captures only stderr""" + + def __init__(self, replay=True): + super().__init__(out=False, replay=replay) + + +class CaptureLogger: + """ + Context manager to capture `logging` streams + + Args: + logger: 'logging` logger object + + Returns: + The captured output is available via `self.out` + + Example: + + ```python + >>> from transformers import logging + >>> from transformers.testing_utils import CaptureLogger + + >>> 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) + >>> assert cl.out, msg + "\n" + ``` + """ + + 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" + + +@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. + + Example: + + ```python + with LoggingLevel(logging.INFO): + AutoModel.from_pretrained("openai-community/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) + + +class TemporaryHubRepo: + """Create a temporary Hub repository and return its `RepoUrl` object. This is similar to + `tempfile.TemporaryDirectory` and can be used as a context manager. For example: + + with TemporaryHubRepo(token=self._token) as temp_repo: + ... + + Upon exiting the context, the repository and everything contained in it are removed. + + Example: + + ```python + with TemporaryHubRepo(token=self._token) as temp_repo: + model.push_to_hub(tmp_repo.repo_id, token=self._token) + ``` + """ + + def __init__(self, namespace: str | None = None, token: str | None = None) -> None: + self.token = token + with tempfile.TemporaryDirectory() as tmp_dir: + repo_id = Path(tmp_dir).name + if namespace is not None: + repo_id = f"{namespace}/{repo_id}" + self.repo_url = create_repo(repo_id, token=self.token) + + def __enter__(self): + return self.repo_url + + def __exit__(self, exc, value, tb): + delete_repo(repo_id=self.repo_url.repo_id, token=self.token, missing_ok=True) + + +@contextlib.contextmanager +# adapted from https://stackoverflow.com/a/64789046/9201239 +def ExtendSysPath(path: str | os.PathLike) -> Iterator[None]: + """ + Temporary add given path to `sys.path`. + + Usage : + + ```python + 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) + + +class TestCasePlus(unittest.TestCase): + """ + This class extends *unittest.TestCase* with additional features. + + 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` + + Feature 2: Flexible auto-removable temporary dirs which are guaranteed to get removed at the end of test. + + 1. Create a unique temporary dir: + + ```python + def test_whatever(self): + tmp_dir = self.get_auto_remove_tmp_dir() + ``` + + `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. + + ```python + def test_whatever(self): + tmp_dir = self.get_auto_remove_tmp_dir("./xxx") + ``` + + 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. + + 3. You can override the first two options by directly overriding the `before` and `after` args, leading to the + following behavior: + + `before=True`: the temporary dir will always be cleared at the beginning of the test. + + `before=False`: if the temporary dir already existed, any existing files will remain there. + + `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. + + 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 `./` + + Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested + otherwise. + + 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. + + + ```python + def test_whatever(self): + env = self.get_env() + ```""" + + def setUp(self): + # get_auto_remove_tmp_dir feature: + self.teardown_tmp_dirs = [] + + # 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.repo_root_dir_str, self.src_dir_str] + if "/examples" in self.test_file_dir_str: + paths.append(self.examples_dir_str) + paths.append(env.get("PYTHONPATH", "")) + + env["PYTHONPATH"] = ":".join(paths) + return env + + def get_auto_remove_tmp_dir(self, tmp_dir=None, before=None, after=None, return_pathlib_obj=False): + """ + Args: + tmp_dir (`string`, *optional*): + if `None`: + + - a unique temporary path will be created + - sets `before=True` if `before` is `None` + - sets `after=True` if `after` is `None` + else: + + - `tmp_dir` will be created + - sets `before=True` if `before` is `None` + - sets `after=False` if `after` is `None` + before (`bool`, *optional*): + 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. + after (`bool`, *optional*): + 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. + return_pathlib_obj (`bool`, *optional*): + If `True` will return a pathlib.Path object + + Returns: + tmp_dir(`string`): either the same value as passed via *tmp_dir* or the path to the auto-selected tmp dir + """ + if tmp_dir is not None: + # 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 + + # 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: + # 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 + + # 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 Path(tmp_dir).resolve() if return_pathlib_obj else tmp_dir + + 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: + + ``` + one_liner_str = 'from transformers import AutoModel; AutoModel.from_pretrained("google-t5/t5-large")' + 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 + + def tearDown(self): + # get_auto_remove_tmp_dir feature: remove registered temp dirs + for path in self.teardown_tmp_dirs: + shutil.rmtree(path, ignore_errors=True) + self.teardown_tmp_dirs = [] + if is_accelerate_available(): + AcceleratorState._reset_state() + PartialState._reset_state() + + # delete all the env variables having `ACCELERATE` in them + for k in list(os.environ.keys()): + if "ACCELERATE" in k: + del os.environ[k] + + +def mockenv(**kwargs): + """ + 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) + + """ + return mock.patch.dict(os.environ, kwargs) + + +# 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] + + +# --- 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 + + +def pytest_terminal_summary_main(tr, id): + """ + 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. + + This function emulates --duration and -rA pytest arguments. + + This function is to be called from `conftest.py` via `pytest_terminal_summary` wrapper that has to be defined + there. + + Args: + - tr: `terminalreporter` passed from `conftest.py` + - 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. + + 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 = f"reports/{id}" + Path(dir).mkdir(parents=True, exist_ok=True) + report_files = { + k: f"{dir}/{k}.txt" + for k in [ + "durations", + "errors", + "failures_long", + "failures_short", + "failures_line", + "passes", + "stats", + "summary_short", + "warnings", + ] + } + + # 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") + + 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.MULTILINE | re.DOTALL) + tr._tw.line(longrepr) + # note: not printing out any rep.sections to keep the report short + + # 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) + + # report failures with line/short/long styles + config.option.tbstyle = "auto" # full tb + with open(report_files["failures_long"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_failures() + + # 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: + tr._tw = create_terminal_writer(config, f) + tr.summary_errors() + + with open(report_files["warnings"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_warnings() # normal warnings + tr.summary_warnings() # final warnings + + tr.reportchars = "wPpsxXEf" # emulate -rA (used in summary_passes() and short_test_summary()) + + # 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() + + with open(report_files["summary_short"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.short_test_summary() + + with open(report_files["stats"], "w") as f: + tr._tw = create_terminal_writer(config, f) + tr.summary_stats() + + # restore: + tr._tw = orig_writer + tr.reportchars = orig_reportchars + config.option.tbstyle = orig_tbstyle + + +# --- distributed testing functions --- # + +# 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( + [ + asyncio.create_task(_read_stream(p.stdout, lambda l: tee(l, out, sys.stdout, label="stdout:"))), + asyncio.create_task(_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: + stderr = "\n".join(result.stderr) + raise RuntimeError( + f"'{cmd_str}' failed with returncode {result.returncode}\n\n" + f"The combined stderr from workers follows:\n{stderr}" + ) + + # 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: + raise RuntimeError(f"'{cmd_str}' produced no output.") + + return result + + +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.MULTILINE) + return int(worker) + + +def get_torch_dist_unique_port(): + """ + Returns a free port number that can be fed to `torch.distributed.launch`'s `--master_port` argument. + + Binds to port 0 to let the OS assign an available port, avoiding collisions from hardcoded ports + and TCP TIME_WAIT issues between sequential subprocess launches. + """ + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +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. + """ + import numpy as np + + if isinstance(obj, list): + return [nested_simplify(item, decimals) for item in obj] + if isinstance(obj, tuple): + return tuple(nested_simplify(item, decimals) for item in obj) + elif isinstance(obj, np.ndarray): + return nested_simplify(obj.tolist()) + elif isinstance(obj, Mapping): + return {nested_simplify(k, decimals): nested_simplify(v, decimals) for k, v in obj.items()} + elif isinstance(obj, (str, int, np.int64)) or obj is None: + return obj + elif is_torch_available() and isinstance(obj, torch.Tensor): + return nested_simplify(obj.tolist(), decimals) + elif isinstance(obj, float): + return round(obj, decimals) + elif isinstance(obj, (np.int32, np.float32, np.float16)): + return nested_simplify(obj.item(), decimals) + else: + raise Exception(f"Not supported: {type(obj)}") + + +def check_json_file_has_correct_format(file_path): + with open(file_path) 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() == "}" + + +def to_2tuple(x): + if isinstance(x, collections.abc.Iterable): + return x + return (x, x) + + +# 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 occurred 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 + + +class RequestCounter: + """ + Helper class that will count all requests made online. + + 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 + ``` + """ + + def __enter__(self): + self._counter = defaultdict(int) + self._thread_id = threading.get_ident() + self._extra_info = [] + + def patched_with_thread_info(func): + def wrap(*args, **kwargs): + self._extra_info.append(threading.get_ident()) + return func(*args, **kwargs) + + return wrap + + import urllib3 + + self.patcher = patch.object( + urllib3.connectionpool.log, "debug", side_effect=patched_with_thread_info(urllib3.connectionpool.log.debug) + ) + self.mock = self.patcher.start() + return self + + def __exit__(self, *args, **kwargs) -> None: + assert len(self.mock.call_args_list) == len(self._extra_info) + for thread_id, call in zip(self._extra_info, self.mock.call_args_list): + if thread_id != self._thread_id: + continue + # code 307: the URL being requested by the user has moved to a temporary location + if call.args[-2] == 307: + continue + 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() + + def __getitem__(self, key: str) -> int: + return self._counter[key] + + @property + def total_calls(self) -> int: + return sum(self._counter.values()) + + +def is_flaky(max_attempts: int = 5, wait_before_retry: float | None = None, description: str | None = None): + """ + To decorate flaky tests. They will be retried on failures. + + Please note that our push tests use `pytest-rerunfailures`, which prompts the CI to rerun certain types of + failed tests. More specifically, if the test exception contains any substring in `FLAKY_TEST_FAILURE_PATTERNS` + (in `.circleci/create_circleci_config.py`), it will be rerun. If you find a recurrent pattern of failures, + expand `FLAKY_TEST_FAILURE_PATTERNS` in our CI configuration instead of using `is_flaky`. + + 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. + description (`str`, *optional*): + A string to describe the situation (what / where / why is flaky, link to GH issue/PR comments, errors, + etc.) + """ + + 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: + logger.error(f"Test failed with {err} at try {retry_count}/{max_attempts}.") + if wait_before_retry is not None: + time.sleep(wait_before_retry) + retry_count += 1 + + return test_func_ref(*args, **kwargs) + + return unittest.skipUnless(_run_flaky_tests, "test is flaky")(wrapper) + + return decorator + + +def hub_retry(max_attempts: int = 5, wait_before_retry: float | None = 2): + """ + To decorate tests that download from the Hub. They can fail due to a + variety of network issues such as timeouts, connection resets, etc. + + Args: + max_attempts (`int`, *optional*, defaults to 5): + The maximum number of attempts to retry the flaky test. + wait_before_retry (`float`, *optional*, defaults to 2): + If provided, will wait that number of seconds before retrying the test. + """ + + 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) + # We catch all exceptions related to network issues from httpx + except ( + httpx.HTTPError, + httpx.RequestError, + httpx.TimeoutException, + httpx.ReadTimeout, + httpx.ConnectError, + httpx.NetworkError, + ) as err: + logger.error( + f"Test failed with {err} at try {retry_count}/{max_attempts} as it couldn't connect to the specified Hub repository." + ) + 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 + + +def run_first(test_case): + """ + Decorator marking a test with order(1). When pytest-order plugin is installed, tests marked with this decorator + are guaranteed to run first. + + This is especially useful in some test settings like on a Gaudi instance where a Gaudi device can only be used by a + single process at a time. So we make sure all tests that run in a subprocess are launched first, to avoid device + allocation conflicts. + """ + # Without this check, we get unwanted warnings when it's not installed + if is_pytest_order_available(): + import pytest + + return pytest.mark.order(1)(test_case) + else: + return test_case + + +def run_test_in_subprocess(test_case, target_func, inputs=None, timeout=None): + """ + 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. + 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`. + """ + if timeout is None: + timeout = int(os.environ.get("PYTEST_TIMEOUT", "600")) + + 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']}") + + +def run_test_using_subprocess(func): + """ + To decorate a test to run in a subprocess using the `subprocess` module. This could avoid potential GPU memory + issues (GPU OOM or a test that causes many subsequential failing with `CUDA error: device-side assert triggered`). + """ + import pytest + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if os.getenv("_INSIDE_SUB_PROCESS", None) == "1": + func(*args, **kwargs) + else: + test = " ".join(os.environ.get("PYTEST_CURRENT_TEST").split(" ")[:-1]) + try: + env = copy.deepcopy(os.environ) + env["_INSIDE_SUB_PROCESS"] = "1" + # This prevents the entries in `short test summary info` given by the subprocess being truncated. so the + # full information can be passed to the parent pytest process. + # See: https://docs.pytest.org/en/stable/explanation/ci.html + env["CI"] = "true" + + # If not subclass of `unitTest.TestCase` and `pytestconfig` is used: try to grab and use the arguments + if "pytestconfig" in kwargs: + command = list(kwargs["pytestconfig"].invocation_params.args) + for idx, x in enumerate(command): + if x in kwargs["pytestconfig"].args: + test = test.split("::")[1:] + command[idx] = "::".join([f"{func.__globals__['__file__']}"] + test) + command = [f"{sys.executable}", "-m", "pytest"] + command + command = [x for x in command if x != "--no-summary"] + # Otherwise, simply run the test with no option at all + else: + command = [f"{sys.executable}", "-m", "pytest", f"{test}"] + + subprocess.run(command, env=env, check=True, capture_output=True) + except subprocess.CalledProcessError as e: + exception_message = e.stdout.decode() + lines = exception_message.split("\n") + # Add a first line with more informative information instead of just `= test session starts =`. + # This makes the `short test summary info` section more useful. + if "= test session starts =" in lines[0]: + text = "" + for line in lines[1:]: + if line.startswith("FAILED "): + text = line[len("FAILED ") :] + text = "".join(text.split(" - ")[1:]) + elif line.startswith("=") and line.endswith("=") and " failed in " in line: + break + elif len(text) > 0: + text += f"\n{line}" + text = "(subprocess) " + text + lines = [text] + lines + exception_message = "\n".join(lines) + raise pytest.fail(exception_message, pytrace=False) + + return wrapper + + +""" +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 corresponding line. + +To skip cuda tests, make sure to call `SKIP_CUDA_DOCTEST=1 pytest --doctest-modules +""" + + +def preprocess_string(string, skip_cuda_tests): + """Prepare a docstring or a `.md` file to be run by doctest. + + The argument `string` would be the whole file content if it is a `.md` file. For a python file, it would be one of + 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*>>> )(.*?```)" + codeblocks = re.split(codeblock_pattern, string, flags=re.DOTALL) + 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 + + modified_string = "" + if not is_cuda_found: + modified_string = "".join(codeblocks) + + 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 + (?:^(?P [ ]*) >>> .*) # PS1 line + (?:\n [ ]* \.\.\. .*)*) # PS2 lines + \n? + # Want consists of any non-blank lines that do not start with PS1. + (?P (?:(?![ ]*$) # 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 = os.environ.get("SKIP_CUDA_DOCTEST", "0") == "1" + # !!!!!!!!!!! HF Specific !!!!!!!!!!! + + def parse(self, string, name=""): + """ + 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) + + +def _device_agnostic_dispatch(device: str, dispatch_table: dict[str, Callable], *args, **kwargs): + if device not in dispatch_table: + if not callable(dispatch_table["default"]): + return dispatch_table["default"] + + return dispatch_table["default"](*args, **kwargs) + + fn = dispatch_table[device] + + # Some device agnostic functions return values or None, will return then directly. + if not callable(fn): + return fn + + 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, + } + BACKEND_RESET_MAX_MEMORY_ALLOCATED = { + "cuda": torch.cuda.reset_max_memory_allocated, + "cpu": None, + "default": None, + } + BACKEND_MAX_MEMORY_ALLOCATED = { + "cuda": torch.cuda.max_memory_allocated, + "cpu": 0, + "default": 0, + } + BACKEND_RESET_PEAK_MEMORY_STATS = { + "cuda": torch.cuda.reset_peak_memory_stats, + "cpu": None, + "default": None, + } + BACKEND_MEMORY_ALLOCATED = { + "cuda": torch.cuda.memory_allocated, + "cpu": 0, + "default": 0, + } + BACKEND_SYNCHRONIZE = { + "cuda": torch.cuda.synchronize, + "cpu": None, + "default": None, + } + BACKEND_TORCH_ACCELERATOR_MODULE = { + "cuda": torch.cuda, + "cpu": None, + "default": None, + } +else: + BACKEND_MANUAL_SEED = {"default": None} + BACKEND_EMPTY_CACHE = {"default": None} + BACKEND_DEVICE_COUNT = {"default": lambda: 0} + BACKEND_RESET_MAX_MEMORY_ALLOCATED = {"default": None} + BACKEND_RESET_PEAK_MEMORY_STATS = {"default": None} + BACKEND_MAX_MEMORY_ALLOCATED = {"default": 0} + BACKEND_MEMORY_ALLOCATED = {"default": 0} + BACKEND_SYNCHRONIZE = {"default": None} + BACKEND_TORCH_ACCELERATOR_MODULE = {"default": None} + + +if is_torch_hpu_available(): + BACKEND_MANUAL_SEED["hpu"] = torch.hpu.manual_seed + BACKEND_DEVICE_COUNT["hpu"] = torch.hpu.device_count + BACKEND_TORCH_ACCELERATOR_MODULE["hpu"] = torch.hpu + +if is_torch_mlu_available(): + BACKEND_EMPTY_CACHE["mlu"] = torch.mlu.empty_cache + BACKEND_MANUAL_SEED["mlu"] = torch.mlu.manual_seed + BACKEND_DEVICE_COUNT["mlu"] = torch.mlu.device_count + BACKEND_TORCH_ACCELERATOR_MODULE["mlu"] = torch.mlu + +if is_torch_npu_available(): + BACKEND_EMPTY_CACHE["npu"] = torch.npu.empty_cache + BACKEND_MANUAL_SEED["npu"] = torch.npu.manual_seed + BACKEND_DEVICE_COUNT["npu"] = torch.npu.device_count + BACKEND_TORCH_ACCELERATOR_MODULE["npu"] = torch.npu + +if is_torch_xpu_available(): + BACKEND_EMPTY_CACHE["xpu"] = torch.xpu.empty_cache + BACKEND_MANUAL_SEED["xpu"] = torch.xpu.manual_seed + BACKEND_DEVICE_COUNT["xpu"] = torch.xpu.device_count + BACKEND_RESET_MAX_MEMORY_ALLOCATED["xpu"] = torch.xpu.reset_peak_memory_stats + BACKEND_RESET_PEAK_MEMORY_STATS["xpu"] = torch.xpu.reset_peak_memory_stats + BACKEND_MAX_MEMORY_ALLOCATED["xpu"] = torch.xpu.max_memory_allocated + BACKEND_MEMORY_ALLOCATED["xpu"] = torch.xpu.memory_allocated + BACKEND_SYNCHRONIZE["xpu"] = torch.xpu.synchronize + BACKEND_TORCH_ACCELERATOR_MODULE["xpu"] = torch.xpu + + +if is_torch_xla_available(): + BACKEND_EMPTY_CACHE["xla"] = torch.cuda.empty_cache + BACKEND_MANUAL_SEED["xla"] = torch.cuda.manual_seed + BACKEND_DEVICE_COUNT["xla"] = torch.cuda.device_count + + +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) + + +def backend_reset_max_memory_allocated(device: str): + return _device_agnostic_dispatch(device, BACKEND_RESET_MAX_MEMORY_ALLOCATED) + + +def backend_reset_peak_memory_stats(device: str): + return _device_agnostic_dispatch(device, BACKEND_RESET_PEAK_MEMORY_STATS) + + +def backend_max_memory_allocated(device: str): + return _device_agnostic_dispatch(device, BACKEND_MAX_MEMORY_ALLOCATED) + + +def backend_memory_allocated(device: str): + return _device_agnostic_dispatch(device, BACKEND_MEMORY_ALLOCATED) + + +def backend_synchronize(device: str): + return _device_agnostic_dispatch(device, BACKEND_SYNCHRONIZE) + + +def backend_torch_accelerator_module(device: str): + return _device_agnostic_dispatch(device, BACKEND_TORCH_ACCELERATOR_MODULE) + + +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. + device_spec_dir, _ = os.path.split(os.path.realpath(device_spec_path)) + sys.path.append(device_spec_dir) + 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") + + +def compare_pipeline_output_to_hub_spec(output, hub_spec): + missing_keys = [] + unexpected_keys = [] + all_field_names = {field.name for field in fields(hub_spec)} + matching_keys = sorted([key for key in output if key in all_field_names]) + + # Fields with a MISSING default are required and must be in the output + for field in fields(hub_spec): + if field.default is MISSING and field.name not in output: + missing_keys.append(field.name) + + # All output keys must match either a required or optional field in the Hub spec + for output_key in output: + if output_key not in all_field_names: + unexpected_keys.append(output_key) + + if missing_keys or unexpected_keys: + error = ["Pipeline output does not match Hub spec!"] + if matching_keys: + error.append(f"Matching keys: {matching_keys}") + if missing_keys: + error.append(f"Missing required keys in pipeline output: {missing_keys}") + if unexpected_keys: + error.append(f"Keys in pipeline output that are not in Hub spec: {unexpected_keys}") + raise KeyError("\n".join(error)) + + +@require_torch +def cleanup(device: str, gc_collect=False): + if gc_collect: + gc.collect() + backend_empty_cache(device) + torch.compiler.reset() + + +# Type definition of key used in `Expectations` class. +DeviceProperties = tuple[str | None, int | None, int | None] +# Helper type. Makes creating instances of `Expectations` smoother. +PackedDeviceProperties = tuple[str | None, None | int | tuple[int, int]] + + +@cache +def get_device_properties() -> DeviceProperties: + """ + Get environment device properties. + """ + if IS_CUDA_SYSTEM or IS_ROCM_SYSTEM: + import torch + + major, minor = torch.cuda.get_device_capability() + if IS_ROCM_SYSTEM: + return ("rocm", major, minor) + else: + return ("cuda", major, minor) + elif IS_XPU_SYSTEM: + import torch + + # To get more info of the architecture meaning and bit allocation, refer to https://github.com/intel/llvm/blob/sycl/sycl/include/sycl/ext/oneapi/experimental/device_architecture.def + arch = torch.xpu.get_device_capability()["architecture"] + gen_mask = 0x000000FF00000000 + gen = (arch & gen_mask) >> 32 + return ("xpu", gen, None) + elif IS_NPU_SYSTEM: + return ("npu", None, None) + else: + return (torch_device, None, None) + + +def unpack_device_properties( + properties: PackedDeviceProperties | None = None, +) -> DeviceProperties: + """ + Unpack a `PackedDeviceProperties` tuple into consistently formatted `DeviceProperties` tuple. If properties is None, it is fetched. + """ + if properties is None: + return get_device_properties() + device_type, major_minor = properties + if major_minor is None: + major, minor = None, None + elif isinstance(major_minor, int): + major, minor = major_minor, None + else: + major, minor = major_minor + return device_type, major, minor + + +class Expectations(UserDict[PackedDeviceProperties, Any]): + def get_expectation(self) -> Any: + """ + Find best matching expectation based on environment device properties. We look at device_type, major and minor + versions of the drivers. Expectations are stored as a dictionary with keys of the form + (device_type, (major, minor)). If the major and minor versions are not provided, we use None. + """ + return self.find_expectation(get_device_properties()) + + def unpacked(self) -> list[tuple[DeviceProperties, Any]]: + return [(unpack_device_properties(k), v) for k, v in self.data.items()] + + @staticmethod + def is_default(expectation_key: PackedDeviceProperties) -> bool: + """ + This function returns True if the expectation_key is the Default expectation (None, None). + When an Expectation dict contains a Default value, it is generally because the test existed before Expectations. + When we modify a test to use Expectations for a specific hardware, we don't want to affect the tests on other + hardwares. Thus we set the previous value as the Default expectation with key (None, None) and add a value for + the specific hardware with key (hardware_type, (major, minor)). + """ + return all(p is None for p in expectation_key) + + @staticmethod + def score(properties: DeviceProperties, other: DeviceProperties) -> float: + """ + Returns score indicating how similar two instances of the `Properties` tuple are. + Rules are as follows: + * Matching `type` adds one point, semi-matching `type` adds 0.1 point (e.g. cuda and rocm). + * If types match, matching `major` adds another point, and then matching `minor` adds another. + * The Default expectation (None, None) is worth 0.5 point, which is better than semi-matching. More on this + in the `is_default` function. + """ + device_type, major, minor = properties + other_device_type, other_major, other_minor = other + + score = 0 + # Matching device type, maybe major and minor + if device_type is not None and device_type == other_device_type: + score += 1 + if major is not None and major == other_major: + score += 1 + if minor is not None and minor == other_minor: + score += 1 + # Semi-matching device type, which carries less importance than the default expectation + elif device_type in ["cuda", "rocm"] and other_device_type in ["cuda", "rocm"]: + score = 0.1 + + # Default expectation + if Expectations.is_default(other): + score = 0.5 + + return score + + def find_expectation(self, properties: DeviceProperties = (None, None, None)) -> Any: + """ + Find best matching expectation based on provided device properties. We score each expectation, and to + distinguish between expectations with the same score, we use the major and minor version numbers, prioritizing + most recent versions. + """ + (result_key, result) = max( + self.unpacked(), + key=lambda x: ( + Expectations.score(properties, x[0]), # x[0] is a device properties tuple (device_type, major, minor) + x[0][1] if x[0][1] is not None else -1, # This key is the major version, -1 if major is None + x[0][2] if x[0][2] is not None else -1, # This key is the minor version, -1 if minor is None + ), + ) + + if Expectations.score(properties, result_key) == 0: + raise ValueError(f"No matching expectation found for {properties}") + + return result + + def __repr__(self): + return f"{self.data}" + + +def patch_torch_compile_force_graph(): + """ + Patch `torch.compile` to always use `fullgraph=True`. + + This is useful when some `torch.compile` tests are running with `fullgraph=False` and we want to be able to run + them with `fullgraph=True` in some occasion (without introducing new tests) to make sure there is no graph break. + + After PR #40137, `CompileConfig.fullgraph` is `False` by default, this patch is necessary. + """ + + force_fullgraph = os.environ.get("TORCH_COMPILE_FORCE_FULLGRAPH", "") + force_fullgraph = force_fullgraph.lower() in ("yes", "true", "on", "t", "y", "1") + + if force_fullgraph: + import torch + + orig_method = torch.compile + + def patched(*args, **kwargs): + # In `torch_compile`, all arguments except `model` is keyword only argument. + kwargs["fullgraph"] = True + return orig_method(*args, **kwargs) + + torch.compile = patched + + +def _get_test_info(): + """ + Collect some information about the current test. + + For example, test full name, line number, stack, traceback, etc. + """ + + full_test_name = os.environ.get("PYTEST_CURRENT_TEST", "").split(" ")[0] + test_file, test_class, test_name = full_test_name.split("::") + + # from the most recent frame to the top frame + stack_from_inspect = inspect.stack() + # but visit from the top frame to the most recent frame + + actual_test_file, _actual_test_class = test_file, test_class + test_frame, test_obj, test_method = None, None, None + for frame in reversed(stack_from_inspect): + # if test_file in str(frame).replace(r"\\", "/"): + # check frame's function + if it has `self` as locals; double check if self has the (function) name + # TODO: Question: How about expanded? + if ( + test_name.startswith(frame.function) + and "self" in frame.frame.f_locals + and hasattr(frame.frame.f_locals["self"], test_name) + ): + # if test_name == frame.frame.f_locals["self"]._testMethodName: + test_frame = frame + # The test instance + test_obj = frame.frame.f_locals["self"] + # TODO: Do we get the (relative?) path or it's just a file name? + # TODO: Does `test_obj` always have `tearDown` object? + actual_test_file = frame.filename + # TODO: check `test_method` will work used at the several places! + test_method = getattr(test_obj, test_name) + break + + if test_frame is not None: + line_number = test_frame.lineno + + # The frame of `patched` being called (the one and the only one calling `_get_test_info`) + # This is used to get the original method being patched in order to get the context. + frame_of_patched_obj = None + + captured_frames = [] + to_capture = False + # From the most outer (i.e. python's `runpy.py`) frame to most inner frame (i.e. the frame of this method) + # Between `the test method being called` and `before entering `patched``. + for frame in reversed(stack_from_inspect): + if ( + test_name.startswith(frame.function) + and "self" in frame.frame.f_locals + and hasattr(frame.frame.f_locals["self"], test_name) + ): + to_capture = True + # TODO: check simply with the name is not robust. + elif frame.frame.f_code.co_name == "patched": + frame_of_patched_obj = frame + to_capture = False + break + if to_capture: + captured_frames.append(frame) + + tb_next = None + for frame_info in reversed(captured_frames): + tb = types.TracebackType(tb_next, frame_info.frame, frame_info.frame.f_lasti, frame_info.frame.f_lineno) + tb_next = tb + test_traceback = tb + + origin_method_being_patched = frame_of_patched_obj.frame.f_locals["orig_method"] + + # An iterable of type `traceback.StackSummary` with each element of type `FrameSummary` + stack = traceback.extract_stack() + # The frame which calls `the original method being patched` + caller_frame = None + # From the most inner (i.e. recent) frame to the most outer frame + for frame in reversed(stack): + if origin_method_being_patched.__name__ in frame.line: + caller_frame = frame + + caller_path = os.path.relpath(caller_frame.filename) + caller_lineno = caller_frame.lineno + + test_lineno = line_number + + # Get the code context in the test function/method. + from _pytest._code.source import Source + + with open(actual_test_file) as fp: + s = fp.read() + source = Source(s) + test_code_context = "\n".join(source.getstatement(test_lineno - 1).lines) + + # Get the code context in the caller (to the patched function/method). + with open(caller_path) as fp: + s = fp.read() + source = Source(s) + caller_code_context = "\n".join(source.getstatement(caller_lineno - 1).lines) + + test_info = f"test:\n\n{full_test_name}\n\n{'-' * 80}\n\ntest context: {actual_test_file}:{test_lineno}\n\n{test_code_context}" + test_info = f"{test_info}\n\n{'-' * 80}\n\ncaller context: {caller_path}:{caller_lineno}\n\n{caller_code_context}" + + return ( + full_test_name, + test_file, + test_lineno, + test_obj, + test_method, + test_frame, + test_traceback, + test_code_context, + caller_path, + caller_lineno, + caller_code_context, + test_info, + ) + + +def _get_call_arguments(code_context): + """ + Analyze the positional and keyword arguments in a call expression. + + This will extract the expressions of the positional and kwyword arguments, and associate them to the positions and + the keyword argument names. + """ + + def get_argument_name(node): + """Extract the name/expression from an AST node""" + if isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Attribute): + return ast.unparse(node) + elif isinstance(node, ast.Constant): + return repr(node.value) + else: + return ast.unparse(node) + + indent = len(code_context) - len(code_context.lstrip()) + code_context = code_context.replace(" " * indent, "") + + try: + # Parse the line + tree = ast.parse(code_context, mode="eval") + + assert isinstance(tree.body, ast.Call) + call_node = tree.body + + if call_node: + result = { + "positional_args": [], + "keyword_args": {}, + "starargs": None, # *args + "kwargs": None, # **kwargs + } + + # Extract positional arguments + for arg in call_node.args: + arg_name = get_argument_name(arg) + result["positional_args"].append(arg_name) + + # Extract keyword arguments + for keyword in call_node.keywords: + if keyword.arg is None: + # This is **kwargs + result["kwargs"] = get_argument_name(keyword.value) + else: + # Regular keyword argument + arg_name = get_argument_name(keyword.value) + result["keyword_args"][keyword.arg] = arg_name + + return result + + except (SyntaxError, AttributeError) as e: + print(f"Error parsing: {e}") + + return None + + +def _prepare_debugging_info(test_info, info): + """Combine the information about the test and the call information to a patched function/method within it.""" + + info = f"{test_info}\n\n{info}" + p = os.path.join(os.environ.get("_PATCHED_TESTING_METHODS_OUTPUT_DIR", ""), "captured_info.txt") + # TODO (ydshieh): This is not safe when we use pytest-xdist with more than 1 worker. + with open(p, "a") as fp: + fp.write(f"{info}\n\n{'=' * 120}\n\n") + + return info + + +def _patched_tearDown(self, *args, **kwargs): + """Used to report a test that has failures captured and handled by patched functions/methods (without re-raise). + + The patched functions/methods refer to the `patched` defined in `_patch_with_call_info`, which is applied to + `torch.testing.assert_close` and `unittest.case.TestCase.assertEqual`. + + The objective is to avoid a failure being silence after being processed. + + If there is any failure that is not handled by the patched functions/methods, we add custom error message for them + along with the usual pytest failure report. + """ + + # Check for regular failures before clearing: + # when `_patched_tearDown` is called, the current test fails due to an assertion error given by a method being + # patched by `_patch_with_call_info`. The patched method catches such an error and continue running the remaining + # statements within the test. If the test fails with another error not handled by the patched methods, we don't let + # pytest to fail and report it but the original failure (the first one that was processed) instead. + # We still record those failures not handled by the patched methods, and add custom messages along with the usual + # pytest failure report. + regular_failures_info = [] + + errors = None + if hasattr(self._outcome, "errors"): + errors = self._outcome.errors + elif hasattr(self._outcome, "result") and hasattr(self._outcome.result, "errors"): + errors = self._outcome.result.errors + + if hasattr(self, "_outcome") and errors: + for error_entry in errors: + test_instance, (exc_type, exc_obj, exc_tb) = error_entry + # breakpoint() + regular_failures_info.append( + { + "message": f"{str(exc_obj)}\n\n", + "type": exc_type.__name__, + "file": "test_modeling_vit.py", + "line": 237, # get_deepest_frame_line(exc_tb) # Your helper function + } + ) + + # Clear the regular failure (i.e. that is not from any of our patched assertion methods) from pytest's records. + if hasattr(self._outcome, "errors"): + self._outcome.errors.clear() + elif hasattr(self._outcome, "result") and hasattr(self._outcome.result, "errors"): + self._outcome.result.errors.clear() + + # reset back to the original tearDown method, so `_patched_tearDown` won't be run by the subsequent tests if they + # have only test failures that are not handle by the patched methods (or no test failure at all). + orig_tearDown = _patched_tearDown.orig_tearDown + type(self).tearDown = orig_tearDown + + # Call the original tearDown + orig_tearDown(self, *args, **kwargs) + + # Get the failure + test_method = getattr(self, self._testMethodName) + captured_failures = test_method.__func__.captured_failures[id(test_method)] + + # TODO: How could we show several exceptions in a sinigle test on the terminal? (Maybe not a good idea) + captured_exceptions = captured_failures[0]["exception"] + captured_traceback = captured_failures[0]["traceback"] + # Show the captured information on the terminal. + capturued_info = [x["info"] for x in captured_failures] + capturued_info_str = f"\n\n{'=' * 80}\n\n".join(capturued_info) + + # Enhance the exception message if there were suppressed failures + if regular_failures_info: + enhanced_message = f"""{str(captured_exceptions)} + +{"=" * 80} +Handled Failures: ({len(capturued_info)} handled): +{"-" * 80}\n +{capturued_info_str} + +{"=" * 80} +Unhandled Failures: ({len(regular_failures_info)} unhandled): +{"-" * 80}\n +{", ".join(f"{info['type']}: {info['message']}{info['file']}:{info['line']}" for info in regular_failures_info)} + +{"-" * 80} +Note: This failure occurred after other failures analyzed by the patched assertion methods. +To see the full details, temporarily disable assertion patching. +{"=" * 80}""" + + # Create new exception with enhanced message + enhanced_exception = type(captured_exceptions)(enhanced_message) + enhanced_exception.__cause__ = captured_exceptions.__cause__ + enhanced_exception.__context__ = captured_exceptions.__context__ + + # Raise with your existing traceback reconstruction + captured_exceptions = enhanced_exception + + # clean up the recorded status + del test_method.__func__.captured_failures + + raise captured_exceptions.with_traceback(captured_traceback) + + +def _patch_with_call_info(module_or_class, attr_name, _parse_call_info_func, target_args): + """ + Patch a callerable `attr_name` of a module or class `module_or_class`. + + This will allow us to collect the call information, e.g. the argument names and values, also the literal expressions + passed as the arguments. + """ + orig_method = getattr(module_or_class, attr_name) + if not callable(orig_method): + return + + def patched(*args, **kwargs): + # If the target callable is not called within a test, simply call it without modification. + if not os.environ.get("PYTEST_CURRENT_TEST", ""): + return orig_method(*args, **kwargs) + + try: + orig_method(*args, **kwargs) + except AssertionError as e: + captured_exception = e + # captured_traceback = e.__traceback__ + ( + full_test_name, + test_file, + test_lineno, + test_obj, + test_method, + test_frame, + test_traceback, + test_code_context, + caller_path, + caller_lineno, + caller_code_context, + test_info, + ) = _get_test_info() + test_info = f"{test_info}\n\n{'-' * 80}\n\npatched method: {orig_method.__module__}.{orig_method.__name__}" + call_argument_expressions = _get_call_arguments(caller_code_context) + + # This is specific + info = _parse_call_info_func(orig_method, args, kwargs, call_argument_expressions, target_args) + info = _prepare_debugging_info(test_info, info) + + # If the test is running in a CI environment (e.g. not a manual run), let's raise and fail the test, so it + # behaves as usual. + # On Github Actions or CircleCI, this is set automatically. + # When running manually, it's the user to determine if to set it. + # This is to avoid the patched function being called `with self.assertRaises(AssertionError):` and fails + # because of the missing expected `AssertionError`. + # TODO (ydshieh): If there is way to raise only when we are inside such context managers? + # TODO (ydshieh): How not to record the failure if it happens inside `self.assertRaises(AssertionError)`? + if os.getenv("CI") == "true": + raise captured_exception.with_traceback(test_traceback) + + # Save this, so we can raise at the end of the current test + captured_failure = { + "result": "failed", + "exception": captured_exception, + "traceback": test_traceback, + "info": info, + } + + # Record the failure status and its information, so we can raise it later. + # We are modifying the (unbound) function at class level: not its logic but only adding a new extra + # attribute. + if getattr(test_method.__func__, "captured_failures", None) is None: + test_method.__func__.captured_failures = {} + if id(test_method) not in test_method.__func__.captured_failures: + test_method.__func__.captured_failures[id(test_method)] = [] + test_method.__func__.captured_failures[id(test_method)].append(captured_failure) + + # This modifies the `tearDown` which will be called after every tests, but we reset it back inside + # `_patched_tearDown`. + if not hasattr(type(test_obj).tearDown, "orig_tearDown"): + orig_tearDown = type(test_obj).tearDown + _patched_tearDown.orig_tearDown = orig_tearDown + type(test_obj).tearDown = _patched_tearDown + + setattr(module_or_class, attr_name, patched) + + +def _parse_call_info(func, args, kwargs, call_argument_expressions, target_args): + """ + Prepare a string containing the call info to `func`, e.g. argument names/values/expressions. + """ + signature = inspect.signature(func) + signature_names = [param.name for param_name, param in signature.parameters.items()] + + # called as `self.method_name()` or `xxx.method_name()`. + if len(args) == len(call_argument_expressions["positional_args"]) + 1: + # We simply add "self" as the expression despite it might not be the actual argument name. + # (This part is very unlikely what a user would be interest to know) + call_argument_expressions["positional_args"] = ["self"] + call_argument_expressions["positional_args"] + + param_position_mapping = {param_name: idx for idx, param_name in enumerate(signature_names)} + + arg_info = {} + for arg_name in target_args: + if arg_name in kwargs: + arg_value = kwargs[arg_name] + arg_expr = call_argument_expressions["keyword_args"][arg_name] + else: + arg_pos = param_position_mapping[arg_name] + arg_value = args[arg_pos] + arg_expr = call_argument_expressions["positional_args"][arg_pos] + + arg_value_str = _format_py_obj(arg_value) + arg_info[arg_name] = {"arg_expr": arg_expr, "arg_value_str": arg_value_str} + + info = "" + for arg_name in arg_info: + arg_expr, arg_value_str = arg_info[arg_name]["arg_expr"], arg_info[arg_name]["arg_value_str"] + info += f"{'-' * 80}\n\nargument name: `{arg_name}`\nargument expression: `{arg_expr}`\n\nargument value:\n\n{arg_value_str}\n\n" + + # remove the trailing \n\n + info = info[:-2] + + return info + + +def patch_testing_methods_to_collect_info(): + """ + Patch some methods (`torch.testing.assert_close`, `unittest.case.TestCase.assertEqual`, etc). + + This will allow us to collect the call information, e.g. the argument names and values, also the literal expressions + passed as the arguments. + """ + p = os.path.join(os.environ.get("_PATCHED_TESTING_METHODS_OUTPUT_DIR", ""), "captured_info.txt") + Path(p).unlink(missing_ok=True) + + if is_torch_available(): + import torch + + _patch_with_call_info(torch.testing, "assert_close", _parse_call_info, target_args=("actual", "expected")) + + _patch_with_call_info(unittest.case.TestCase, "assertEqual", _parse_call_info, target_args=("first", "second")) + _patch_with_call_info(unittest.case.TestCase, "assertListEqual", _parse_call_info, target_args=("list1", "list2")) + _patch_with_call_info( + unittest.case.TestCase, "assertTupleEqual", _parse_call_info, target_args=("tuple1", "tuple2") + ) + _patch_with_call_info(unittest.case.TestCase, "assertSetEqual", _parse_call_info, target_args=("set1", "set1")) + _patch_with_call_info(unittest.case.TestCase, "assertDictEqual", _parse_call_info, target_args=("d1", "d2")) + _patch_with_call_info(unittest.case.TestCase, "assertIn", _parse_call_info, target_args=("member", "container")) + _patch_with_call_info(unittest.case.TestCase, "assertNotIn", _parse_call_info, target_args=("member", "container")) + _patch_with_call_info(unittest.case.TestCase, "assertLess", _parse_call_info, target_args=("a", "b")) + _patch_with_call_info(unittest.case.TestCase, "assertLessEqual", _parse_call_info, target_args=("a", "b")) + _patch_with_call_info(unittest.case.TestCase, "assertGreater", _parse_call_info, target_args=("a", "b")) + _patch_with_call_info(unittest.case.TestCase, "assertGreaterEqual", _parse_call_info, target_args=("a", "b")) + + +def torchrun(script: str, nproc_per_node: int, is_torchrun: bool = True, env: dict | None = None): + """Run the `script` using `torchrun` command for multi-processing in a subprocess. Captures errors as necessary.""" + with tempfile.NamedTemporaryFile(mode="w+", suffix=".py") as tmp: + tmp.write(script) + tmp.flush() + tmp.seek(0) + if is_torchrun: + cmd = ( + f"torchrun --nproc_per_node {nproc_per_node} --master_port {get_torch_dist_unique_port()} {tmp.name}" + ).split() + else: + cmd = ["python3", tmp.name] + + # Note that the subprocess will be waited for here, and raise an error if not successful + try: + _ = subprocess.run(cmd, capture_output=True, env=env, text=True, check=True) + except subprocess.CalledProcessError as e: + raise Exception(f"The following error was captured: {e.stderr}") + + +def _format_tensor(t, indent_level=0, sci_mode=None): + """Format torch's tensor in a pretty way to be shown 👀 in the test report.""" + + # `torch.testing.assert_close` could accept python int/float numbers. + if not isinstance(t, torch.Tensor): + t = torch.tensor(t) + + # Simply make the processing below simpler (not to handle both cases) + is_scalar = False + if t.ndim == 0: + t = torch.tensor([t]) + is_scalar = True + + # For scalar or one-dimensional tensor, keep it as one-line. If there is only one element along any dimension except + # the last one, we also keep it as one-line. + if t.ndim <= 1 or set(t.shape[0:-1]) == {1}: + # Use `detach` to remove `grad_fn=<...>`, and use `to("cpu")` to remove `device='...'` + t = t.detach().to("cpu") + + # We work directly with the string representation instead the tensor itself + t_str = str(t) + + # remove `tensor( ... )` so keep only the content + t_str = t_str.replace("tensor(", "").replace(")", "") + + # Sometimes there are extra spaces between `[` and the first digit of the first value (for alignment). + # For example `[[ 0.06, -0.51], [-0.76, -0.49]]`. It may have multiple consecutive spaces. + # Let's remove such extra spaces. + while "[ " in t_str: + t_str = t_str.replace("[ ", "[") + + # Put everything in a single line. We replace `\n` by a space ` ` so we still keep `,\n` as `, `. + t_str = t_str.replace("\n", " ") + + # Remove repeated spaces (introduced by the previous step) + while " " in t_str: + t_str = t_str.replace(" ", " ") + + # remove leading `[` and `]` for scalar tensor + if is_scalar: + t_str = t_str[1:-1] + + t_str = " " * 4 * indent_level + t_str + + return t_str + + # Otherwise, we separate the representations of each element along an outer dimension by new lines (after a `,`). + # The representation of each element is obtained by calling this function recursively with current `indent_level`. + else: + t_str = str(t) + + # (For the recursive calls should receive this value) + if sci_mode is None: + sci_mode = "e+" in t_str or "e-" in t_str + + # Use the original content to determine the scientific mode to use. This is required as the representation of + # t[index] (computed below) maybe have different format regarding scientific notation. + torch.set_printoptions(sci_mode=sci_mode) + + t_str = " " * 4 * indent_level + "[\n" + # Keep the ending `,` for all outer dimensions whose representations are not put in one-line, even if there is + # only one element along that dimension. + t_str += ",\n".join(_format_tensor(x, indent_level=indent_level + 1, sci_mode=sci_mode) for x in t) + t_str += ",\n" + " " * 4 * indent_level + "]" + + torch.set_printoptions(sci_mode=None) + + return t_str + + +def _quote_string(s): + """Given a string `s`, return a python literal expression that give `s` when it is used in a python source code. + + For example, if `s` is the string `abc`, the return value is `"abc"`. + + We choice double quotes over single quote despite `str(s)` would give `'abc'` instead of `"abc"`. + """ + has_single_quote = "'" in s + has_double_quote = '"' in s + + if has_single_quote and has_double_quote: + # replace any double quote by the raw string r'\"'. + s = s.replace('"', r"\"") + return f'"{s}"' + elif has_single_quote: + return f'"{s}"' + elif has_double_quote: + return f"'{s}'" + else: + return f'"{s}"' + + +def _format_py_obj(obj, indent=0, mode="", cache=None, prefix=""): + """Format python objects of basic built-in type in a pretty way so we could copy-past them to code editor easily. + + Currently, this support int, float, str, list, tuple, and dict. + + It also works with `torch.Tensor` via calling `format_tesnor`. + """ + + if cache is None: + cache = {} + else: + if (id(obj), indent, mode, prefix) in cache: + return cache[(id(obj), indent, mode, prefix)] + + # special format method for `torch.Tensor` + if str(obj.__class__) == "": + return _format_tensor(obj) + + elif obj.__class__.__name__ == "str": + quoted_string = _quote_string(obj) + # we don't want the newline being interpreted + quoted_string = quoted_string.replace("\n", r"\n") + output = quoted_string + + elif obj.__class__.__name__ in ["int", "float"]: + # for float like `1/3`, we will get `0.3333333333333333` + output = str(obj) + + elif obj.__class__.__name__ in ["list", "tuple", "dict"]: + parenthesis = { + "list": "[]", + "tuple": "()", + "dict": "{}", + } + p1, p2 = parenthesis[obj.__class__.__name__] + + elements_without_indent = [] + if isinstance(obj, dict): + for idx, (k, v) in enumerate(obj.items()): + last_element = idx == len(obj) - 1 + ok = _format_py_obj(k, indent=indent + 1, mode="one-line", cache=cache) + ov = _format_py_obj( + v, + indent=indent + 1, + mode=mode, + cache=cache, + prefix=ok.lstrip() + ": " + "," if not last_element else "", + ) + # Each element could be multiple-line, but the indent of its first line is removed + elements_without_indent.append(f"{ok.lstrip()}: {ov.lstrip()}") + + else: + for idx, x in enumerate(obj): + last_element = idx == len(obj) - 1 + o = _format_py_obj( + x, indent=indent + 1, mode=mode, cache=cache, prefix="," if not last_element else "" + ) + # Each element could be multiple-line, but the indent of its first line is removed + elements_without_indent.append(o.lstrip()) + + groups = [] + buf = [] + for idx, x in enumerate(elements_without_indent): + buf.append(x) + + x_expanded = "\n" in buf[-1] + not_last_element = idx != len(elements_without_indent) - 1 + # if `x` should be separated from subsequent elements + should_finalize_x = x_expanded or len(f"{' ' * (4 * (indent + 1))}") + len( + ", ".join(buf[-1:]) + ) > 120 - int(not_last_element) + + # if `buf[:-1]` (i.e. without `x`) should be combined together (into one line) + should_finalize_buf = x_expanded + + # the recursive call returns single line, so we can use it to determine if we can fit the width limit + if not should_finalize_buf: + buf_not_fit_into_one_line = len(f"{' ' * (4 * (indent + 1))}") + len(", ".join(buf)) > 120 - int( + not_last_element + ) + should_finalize_buf = buf_not_fit_into_one_line + + # any element of iterable type need to be on its own line + if (type(obj[idx]) if type(obj) is not dict else type(list(obj.values())[idx])) in [list, tuple, dict]: + should_finalize_x = True + should_finalize_buf = True + + # any type change --> need to be added after a new line + prev_type = None + current_type = type(obj[idx]) if type(obj) is not dict else type(list(obj.values())[idx]) + if len(buf) > 1: + prev_type = type(obj[idx - 1]) if type(obj) is not dict else type(list(obj.values())[idx - 1]) + type_changed = current_type != prev_type + if type_changed: + should_finalize_buf = True + + # all elements in the buf are string --> don't finalize the buf by width limit + if prev_type is None or (prev_type is str and current_type is str): + should_finalize_buf = False + + # collect as many elements of string type as possible (without width limit). + # These will be examined as a whole (if not fit into the width, each element would be in its own line) + if current_type is str: + should_finalize_x = False + # `len(buf) == 1` or `obj[idx-1]` is a string + if prev_type in [None, str]: + should_finalize_buf = False + + if should_finalize_buf: + orig_buf_len = len(buf) + + if orig_buf_len > 1: + not_fit_into_one_line = None + + # all elements in `obj` that give `buf[:-1]` are string. + if prev_type is str: + # `-1` at the end: because buf[-2] is not the last element + not_fit_into_one_line = len(f"{' ' * (4 * (indent + 1))}") + len(", ".join(buf[:-1])) > 120 - 1 + + if not_fit_into_one_line: + for x in buf[:-1]: + groups.append([x]) + else: + groups.append(buf[:-1]) + + buf = buf[-1:] + + if should_finalize_x: + groups.append(buf) + buf = [] + + # The last buf + if len(buf) > 0: + not_fit_into_one_line = None + if current_type is str: + # no `-1` at the end: because buf[-1] is the last element + not_fit_into_one_line = len(f"{' ' * (4 * (indent + 1))}") + len(", ".join(buf)) > 120 + + if not_fit_into_one_line: + for x in buf: + groups.append([x]) + else: + groups.append(buf) + + output = f"{' ' * 4 * indent}{p1}\n" + element_strings = [f"{' ' * (4 * (indent + 1))}" + ", ".join(buf) for buf in groups] + output += ",\n".join(element_strings) + output += f"\n{' ' * 4 * indent}{p2}" + + # if all elements are in one-line + no_new_line_in_elements = all("\n" not in x for x in element_strings) + # if yes, we can form a one-line representation of `obj` + could_use_one_line = no_new_line_in_elements + + # if mode == "one-line", this function always returns one-line representation, so `no_new_line_in_elements` + # will be `True`. + if could_use_one_line: + one_line_form = ", ".join([x.lstrip() for x in element_strings]) + one_line_form = f"{p1}{one_line_form}{p2}" + + if mode == "one-line": + return output + + # check with the width limit + could_use_one_line = len(f"{' ' * 4 * indent}") + len(prefix) + len(one_line_form) <= 120 + + # extra conditions for returning one-line representation + def use_one_line_repr(obj): + # iterable types + if type(obj) in (list, tuple, dict): + # get all types + element_types = [] + if type(obj) is dict: + element_types.extend(type(x) for x in obj.values()) + elif type(obj) in [list, tuple]: + element_types.extend(type(x) for x in obj) + + # At least one element is of iterable type + if any(x in (list, tuple, dict) for x in element_types): + # If `obj` has more than one element and at least one of them is iterable --> no one line repr. + if len(obj) > 1: + return False + + # only one element that is iterable, but not the same type as `obj` --> no one line repr. + if type(obj) is not type(obj[0]): + return False + + # one-line repr. if possible, without width limit + return no_new_line_in_elements + + # all elements are of simple types, but more than one type --> no one line repr. + if len(set(element_types)) > 1: + return False + + # all elements are of the same simple type + if element_types[0] in [int, float]: + # one-line repr. without width limit + return no_new_line_in_elements + elif element_types[0] is str: + if len(obj) == 1: + # one single string element --> one-line repr. without width limit + return no_new_line_in_elements + else: + # multiple string elements --> one-line repr. if fit into width limit + return could_use_one_line + + # simple types (int, flat, string) + return True + + # width condition combined with specific mode conditions + if use_one_line_repr(obj): + output = f"{' ' * 4 * indent}{one_line_form}" + + cache[(id(obj), indent, mode, prefix)] = output + + return output + + +def write_file(file, content): + with open(file, "w") as f: + f.write(content) + + +def read_json_file(file): + with open(file, "r") as fh: + return json.load(fh) + + +# ============================================================================= +# Training CI Utilities - Logging and Memory Monitoring +# ============================================================================= + + +# ANSI color codes for terminal output +class Colors: + """ANSI color codes for terminal output formatting.""" + + RESET = "\033[0m" + BOLD = "\033[1m" + DIM = "\033[2m" + + # Foreground colors + RED = "\033[31m" + GREEN = "\033[32m" + YELLOW = "\033[33m" + BLUE = "\033[34m" + MAGENTA = "\033[35m" + CYAN = "\033[36m" + WHITE = "\033[37m" + + # Bright variants + BRIGHT_RED = "\033[91m" + BRIGHT_GREEN = "\033[92m" + BRIGHT_YELLOW = "\033[93m" + BRIGHT_BLUE = "\033[94m" + BRIGHT_CYAN = "\033[96m" + + +class ColoredFormatter(logging.Formatter): + """Custom formatter that adds colors based on log level.""" + + LEVEL_COLORS = { + logging.DEBUG: Colors.DIM + Colors.CYAN, + logging.INFO: Colors.WHITE, + logging.WARNING: Colors.BRIGHT_YELLOW, + logging.ERROR: Colors.BRIGHT_RED, + logging.CRITICAL: Colors.BOLD + Colors.BRIGHT_RED, + } + + # Loggers that should be dimmed (less important/verbose) + DIMMED_LOGGERS = {"httpx", "httpcore", "urllib3", "requests"} + + def __init__(self, fmt: str | None = None, datefmt: str | None = None): + super().__init__(fmt, datefmt) + + def format(self, record: logging.LogRecord) -> str: + # Check if this logger should be dimmed + is_dimmed = record.name in self.DIMMED_LOGGERS + + if is_dimmed: + # Dim the entire log line for httpx and similar + timestamp = self.formatTime(record, self.datefmt) + message = record.getMessage() + return f"{Colors.DIM}{timestamp} - {record.name} - {record.levelname:8} - {message}{Colors.RESET}" + + # Get color for this level + color = self.LEVEL_COLORS.get(record.levelno, Colors.RESET) + + # Color the level name + levelname = record.levelname + colored_levelname = f"{color}{levelname:8}{Colors.RESET}" + + # Color the timestamp + colored_time = f"{Colors.DIM}{self.formatTime(record, self.datefmt)}{Colors.RESET}" + + # Color the logger name + colored_name = f"{Colors.BLUE}{record.name}{Colors.RESET}" + + # Get message + message = record.getMessage() + + return f"{colored_time} - {colored_name} - {colored_levelname} - {message}" + + +_warn_once_logged: set[str] = set() + + +def init_test_logger() -> logging.Logger: + """Initialize a test-specific logger with colored stderr handler and INFO level for tests. + + Uses a named logger instead of root logger to avoid conflicts with pytest-xdist parallel execution. + Uses stderr instead of stdout to avoid deadlocks with pytest-xdist output capture. + """ + logger = logging.getLogger("transformers.training_test") + logger.setLevel(logging.INFO) + + # Only add handler if not already present (avoid duplicate handlers on repeated calls) + if not logger.handlers: + # Use stderr instead of stdout - pytest-xdist captures stdout which can cause deadlocks + ch = logging.StreamHandler(sys.stderr) + ch.setLevel(logging.INFO) + + # Use colored formatter if terminal supports it, plain otherwise + if sys.stderr.isatty(): + formatter = ColoredFormatter(datefmt="%Y-%m-%d %H:%M:%S") + else: + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S" + ) + + ch.setFormatter(formatter) + logger.addHandler(ch) + + logger.propagate = False # Don't propagate to root logger to avoid duplicate output + return logger + + +def warn_once(logger_instance: logging.Logger, msg: str) -> None: + """Log a warning message only once per unique message. + + Uses a global set to track messages that have already been logged + to prevent duplicate warning messages from cluttering the output. + + Args: + logger_instance: The logger instance to use for warning. + msg: The warning message to log. + """ + if msg not in _warn_once_logged: + logger_instance.warning(msg) + _warn_once_logged.add(msg) + + +# Named tuple for passing memory stats for logging +MemoryStats = collections.namedtuple( + "MemoryStats", + [ + "rss_gib", # Resident Set Size in GiB + "rss_pct", # RSS as percentage of total memory + "vms_gib", # Virtual Memory Size in GiB + "peak_rss_gib", # Peak RSS in GiB + "peak_rss_pct", # Peak RSS as percentage of total memory + "available_gib", # Available system memory in GiB + "total_gib", # Total system memory in GiB + ], +) + + +class CPUMemoryMonitor: + """Monitor CPU memory usage for the current process.""" + + def __init__(self): + self.device_name = "CPU" + self._peak_rss = 0 + self._process = None + self.total_memory = 0 + self.total_memory_gib = 0 + + if is_psutil_available(): + import psutil + + self._process = psutil.Process(os.getpid()) + mem_info = psutil.virtual_memory() + self.total_memory = mem_info.total + self.total_memory_gib = self._to_gib(self.total_memory) + + def _to_gib(self, memory_in_bytes: int) -> float: + """Convert bytes to GiB.""" + return memory_in_bytes / (1024 * 1024 * 1024) + + def _to_pct(self, memory_in_bytes: int) -> float: + """Convert bytes to percentage of total memory.""" + if self.total_memory == 0: + return 0.0 + return 100.0 * memory_in_bytes / self.total_memory + + def _update_peak(self) -> None: + """Update peak memory tracking.""" + if self._process is not None: + current_rss = self._process.memory_info().rss + self._peak_rss = max(self._peak_rss, current_rss) + + def get_stats(self) -> MemoryStats: + """Get current memory statistics.""" + if not is_psutil_available(): + return MemoryStats(0, 0, 0, 0, 0, 0, 0) + + import psutil + + self._update_peak() + + mem_info = self._process.memory_info() + sys_mem = psutil.virtual_memory() + + return MemoryStats( + rss_gib=self._to_gib(mem_info.rss), + rss_pct=self._to_pct(mem_info.rss), + vms_gib=self._to_gib(mem_info.vms), + peak_rss_gib=self._to_gib(self._peak_rss), + peak_rss_pct=self._to_pct(self._peak_rss), + available_gib=self._to_gib(sys_mem.available), + total_gib=self._to_gib(sys_mem.total), + ) + + def reset_peak_stats(self) -> None: + """Reset peak memory tracking.""" + if self._process is not None: + self._peak_rss = self._process.memory_info().rss + + +def build_cpu_memory_monitor(logger_instance: logging.Logger | None = None) -> CPUMemoryMonitor: + """Build and initialize a CPU memory monitor. + + Args: + logger_instance: Optional logger to log initialization info. If None, no logging is done. + + Returns: + CPUMemoryMonitor instance. + """ + monitor = CPUMemoryMonitor() + if logger_instance is not None: + if is_psutil_available(): + logger_instance.info(f"CPU memory monitor initialized: {monitor.total_memory_gib:.2f} GiB total") + else: + logger_instance.warning("psutil not available, memory monitoring disabled") + return monitor + + +def convert_all_safetensors_to_bins(folder: str): + """Convert all safetensors files into torch bin files, to mimic saving with torch (since we still support loading + bin files, but not saving them anymore)""" + for file in os.listdir(folder): + path = os.path.join(folder, file) + if file.endswith(".safetensors"): + new_path = path.replace(".safetensors", ".bin").replace("model", "pytorch_model") + state_dict = load_file(path) + os.remove(path) + torch.save(state_dict, new_path) + # Adapt the index as well + elif file == SAFE_WEIGHTS_INDEX_NAME: + new_path = os.path.join(folder, WEIGHTS_INDEX_NAME) + with open(path) as f: + index = json.loads(f.read()) + os.remove(path) + if "weight_map" in index.keys(): + weight_map = index["weight_map"] + new_weight_map = {} + for k, v in weight_map.items(): + new_weight_map[k] = v.replace(".safetensors", ".bin").replace("model", "pytorch_model") + index["weight_map"] = new_weight_map + with open(new_path, "w") as f: + f.write(json.dumps(index, indent=4)) + + +@contextmanager +def force_serialization_as_bin_files(): + """Since we don't support saving with torch `.bin` files anymore, but still support loading them, we use this context + to easily create the bin files and try to load them back""" + try: + # Monkey patch the method to save as bin files + original_save = PreTrainedModel.save_pretrained + + def new_save(self, save_directory, *args, **kwargs): + original_save(self, save_directory, *args, **kwargs) + convert_all_safetensors_to_bins(save_directory) + + PreTrainedModel.save_pretrained = new_save + + yield + finally: + PreTrainedModel.save_pretrained = original_save