text: # Copyright (c) OpenMMLab. All rights reserved.
  import multiprocessing as mp
  import os
  import platform
  import sys
  
  import cv2
  import mmcv
  import pytest
  from mmcv import Config
  
  import mmflow
  from mmflow.utils import setup_multi_processes
  
  
  @pytest.mark.parametrize('workers_per_gpu', (0, 2))
  @pytest.mark.parametrize(('valid', 'env_cfg'), [(True,
                                                   dict(
                                                       mp_start_method='fork',
                                                       opencv_num_threads=0,
                                                       omp_num_threads=1,
                                                       mkl_num_threads=1)),
                                                  (False,
                                                   dict(
                                                       mp_start_method=1,
                                                       opencv_num_threads=0.1,
                                                       omp_num_threads='s',
                                                       mkl_num_threads='1'))])
  def test_setup_multi_processes(workers_per_gpu, valid, env_cfg):
      # temp save system setting
      sys_start_mehod = mp.get_start_method(allow_none=True)
      sys_cv_threads = cv2.getNumThreads()
      # pop and temp save system env vars
      sys_omp_threads = os.environ.pop('OMP_NUM_THREADS', default=None)
      sys_mkl_threads = os.environ.pop('MKL_NUM_THREADS', default=None)
  
      config = dict(
          data=dict(
              train_dataloader=dict(dict(workers_per_gpu=workers_per_gpu))))
      config.update(env_cfg)
      cfg = Config(config)
      setup_multi_processes(cfg)
  
      # test when cfg is valid and workers_per_gpu > 0
      # setup_multi_processes will work
      if valid and workers_per_gpu > 0:
          # test config without setting env
  
          assert os.getenv('OMP_NUM_THREADS') == str(env_cfg['omp_num_threads'])
          assert os.getenv('MKL_NUM_THREADS') == str(env_cfg['mkl_num_threads'])
          # when set to 0, the num threads will be 1
          assert cv2.getNumThreads() == env_cfg[
              'opencv_num_threads'] if env_cfg['opencv_num_threads'] > 0 else 1
          if platform.system() != 'Windows':
              assert mp.get_start_method() == env_cfg['mp_start_method']
  
          # revert setting to avoid affecting other programs
          if sys_start_mehod:
              mp.set_start_method(sys_start_mehod, force=True)
          cv2.setNumThreads(sys_cv_threads)
          if sys_omp_threads:
              os.environ['OMP_NUM_THREADS'] = sys_omp_threads
          else:
              os.environ.pop('OMP_NUM_THREADS')
          if sys_mkl_threads:
              os.environ['MKL_NUM_THREADS'] = sys_mkl_threads
          else:
              os.environ.pop('MKL_NUM_THREADS')
  
      elif valid and workers_per_gpu == 0:
  
          if platform.system() != 'Windows':
              assert mp.get_start_method() == env_cfg['mp_start_method']
          assert cv2.getNumThreads() == env_cfg[
              'opencv_num_threads'] if env_cfg['opencv_num_threads'] > 0 else 1
          assert 'OMP_NUM_THREADS' not in os.environ
          assert 'MKL_NUM_THREADS' not in os.environ
          if sys_start_mehod:
              mp.set_start_method(sys_start_mehod, force=True)
          cv2.setNumThreads(sys_cv_threads)
          if sys_omp_threads:
              os.environ['OMP_NUM_THREADS'] = sys_omp_threads
          if sys_mkl_threads:
              os.environ['MKL_NUM_THREADS'] = sys_mkl_threads
  
      else:
          assert mp.get_start_method(allow_none=True) == sys_start_mehod
          assert cv2.getNumThreads() == sys_cv_threads
          assert 'OMP_NUM_THREADS' not in os.environ
          assert 'MKL_NUM_THREADS' not in os.environ
  
  
  def test_collect_env():
      try:
          import torch  # noqa: F401
      except ModuleNotFoundError:
          pytest.skip('skipping tests that require PyTorch')
  
      from mmflow.utils import collect_env
      env_info = collect_env()
      expected_keys = [
          'sys.platform', 'Python', 'CUDA available', 'PyTorch',
          'PyTorch compiling details', 'OpenCV', 'MMCV', 'MMCV Compiler',
          'MMCV CUDA Compiler', 'MMFlow', 'GCC'
      ]
      for key in expected_keys:
          assert key in env_info
  
      if env_info['CUDA available']:
          for key in ['CUDA_HOME', 'NVCC']:
              assert key in env_info
  
      assert env_info['sys.platform'] == sys.platform
      assert env_info['Python'] == sys.version.replace('\n', '')
      assert env_info['MMCV'] == mmcv.__version__
      assert mmflow.__version__ in env_info['MMFlow']
  