test.py

"""test.py
Setup and Run hub models.

Make sure to enable an https proxy if necessary, or the setup steps may hang.
"""
# This file shows how to use the benchmark suite from user end.
import gc
import functools
import os
import traceback
import unittest
from unittest.mock import patch
import yaml

import torch
from torchbenchmark import _list_model_paths, _list_canary_model_paths, ModelTask, get_metadata_from_yaml
from torchbenchmark.util.metadata_utils import skip_by_metadata


# Some of the models have very heavyweight setup, so we have to set a very
# generous limit. That said, we don't want the entire test suite to hang if
# a single test encounters an extreme failure, so we give up after a test is
# unresponsive to 5 minutes by default. (Note: this does not require that the
# entire test case completes in 5 minutes. It requires that if the worker is
# unresponsive for 5 minutes the parent will presume it dead / incapacitated.)
TIMEOUT = int(os.getenv("TIMEOUT", 300))  # Seconds

class TestBenchmark(unittest.TestCase):

    def setUp(self):
        gc.collect()

    def tearDown(self):
        gc.collect()

def _create_example_model_instance(task: ModelTask, device: str):
    skip = False
    try:
        task.make_model_instance(test="eval", device=device, extra_args=["--accuracy"])
    except NotImplementedError:
        try:
            task.make_model_instance(test="train", device=device, extra_args=["--accuracy"])
        except NotImplementedError:
            skip = True
    finally:
        if skip:
            raise NotImplementedError(f"Model is not implemented on the device {device}")

def _load_test(path, device):

    def _skip_cuda_memory_check_p(metadata):
        if device != "cuda":
            return True
        if "skip_cuda_memory_leak" in metadata and metadata["skip_cuda_memory_leak"]:
            return True
        return False

    def example_fn(self):
        task = ModelTask(path, timeout=TIMEOUT)
        with task.watch_cuda_memory(skip=_skip_cuda_memory_check_p(metadata), assert_equal=self.assertEqual):
            try:
                _create_example_model_instance(task, device)
                accuracy = task.get_model_attribute("accuracy")
                assert accuracy == "pass" or accuracy == "eager_1st_run_OOM" or accuracy == "eager_2nd_run_OOM", f"Expected accuracy pass, get {accuracy}"
                task.del_model_instance()
            except NotImplementedError as e:
                self.skipTest(f'Method `get_module()` on {device} is not implemented because "{e}", skipping...')

    def train_fn(self):
        metadata = get_metadata_from_yaml(path)
        task = ModelTask(path, timeout=TIMEOUT)
        allow_customize_batch_size = task.get_model_attribute("ALLOW_CUSTOMIZE_BSIZE", classattr=True)
        # to speedup test, use batch size 1 if possible
        batch_size = 1 if allow_customize_batch_size else None
        with task.watch_cuda_memory(skip=_skip_cuda_memory_check_p(metadata), assert_equal=self.assertEqual):
            try:
                task.make_model_instance(test="train", device=device, batch_size=batch_size)
                task.invoke()
                task.check_details_train(device=device, md=metadata)
                task.del_model_instance()
            except NotImplementedError as e:
                self.skipTest(f'Method train on {device} is not implemented because "{e}", skipping...')

    def eval_fn(self):
        metadata = get_metadata_from_yaml(path)
        task = ModelTask(path, timeout=TIMEOUT)
        allow_customize_batch_size = task.get_model_attribute("ALLOW_CUSTOMIZE_BSIZE", classattr=True)
        # to speedup test, use batch size 1 if possible
        batch_size = 1 if allow_customize_batch_size else None
        with task.watch_cuda_memory(skip=_skip_cuda_memory_check_p(metadata), assert_equal=self.assertEqual):
            try:
                task.make_model_instance(test="eval", device=device, batch_size=batch_size)
                task.invoke()
                task.check_details_eval(device=device, md=metadata)
                task.check_eval_output()
                task.del_model_instance()
            except NotImplementedError as e:
                self.skipTest(f'Method eval on {device} is not implemented because "{e}", skipping...')

    def check_device_fn(self):
        task = ModelTask(path, timeout=TIMEOUT)
        with task.watch_cuda_memory(skip=_skip_cuda_memory_check_p(metadata), assert_equal=self.assertEqual):
            try:
                task.make_model_instance(test="eval", device=device)
                task.check_device()
                task.del_model_instance()
            except NotImplementedError as e:
                self.skipTest(f'Method check_device on {device} is not implemented because "{e}", skipping...')

    name = os.path.basename(path)
    metadata = get_metadata_from_yaml(path)
    for fn, fn_name in zip([example_fn, train_fn, eval_fn, check_device_fn],
                           ["example", "train", "eval", "check_device"]):
        # set exclude list based on metadata
        setattr(TestBenchmark, f'test_{name}_{fn_name}_{device}',
                (unittest.skipIf(skip_by_metadata(test=fn_name, device=device, extra_args=[], metadata=metadata), \
                                 "This test is skipped by its metadata")(fn)))


def _load_tests():
    devices = ['cpu']
    if torch.cuda.is_available():
        devices.append('cuda')
    if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        devices.append('mps')
    if device := os.getenv('ACCELERATOR'):
        devices.append(device)
    model_paths = _list_model_paths()
    if os.getenv('USE_CANARY_MODELS'):
        model_paths.extend(_list_canary_model_paths())
    for path in model_paths:
        # TODO: skipping quantized tests for now due to BC-breaking changes for prepare
        # api, enable after PyTorch 1.13 release
        if "quantized" in path:
            continue
        for device in devices:
            _load_test(path, device)


_load_tests()
if __name__ == '__main__':
    unittest.main()