|
|
import os |
|
|
import unittest |
|
|
from argparse import Namespace |
|
|
from unittest import TestCase, mock |
|
|
|
|
|
from trainer import TrainerArgs |
|
|
from trainer.distribute import get_gpus |
|
|
|
|
|
|
|
|
class TestGpusStringParsingMethods(TestCase): |
|
|
@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}) |
|
|
def test_parse_gpus_set_in_env_var_and_args(self): |
|
|
args = Namespace(gpus="0,1") |
|
|
gpus = get_gpus(args) |
|
|
expected_value = ["0"] |
|
|
self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value)) |
|
|
|
|
|
@mock.patch.dict(os.environ, {}) |
|
|
def test_parse_gpus_set_in_args(self): |
|
|
args = Namespace(gpus="0,1") |
|
|
gpus = get_gpus(args) |
|
|
expected_value = ["0", "1"] |
|
|
self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value)) |
|
|
|
|
|
@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}) |
|
|
def test_parse_gpus_set_in_env_var(self): |
|
|
args = Namespace() |
|
|
gpus = get_gpus(args) |
|
|
expected_value = ["0", "1"] |
|
|
self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value)) |
|
|
|
|
|
@mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0, 1 "}) |
|
|
def test_parse_gpus_set_in_env_var_with_spaces(self): |
|
|
args = Namespace() |
|
|
gpus = get_gpus(args) |
|
|
expected_value = ["0", "1"] |
|
|
self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value)) |
|
|
|
|
|
@mock.patch.dict(os.environ, {}) |
|
|
def test_parse_gpus_set_in_args_with_spaces(self): |
|
|
args = Namespace(gpus="0, 1, 2, 3 ") |
|
|
gpus = get_gpus(args) |
|
|
expected_value = ["0", "1", "2", "3"] |
|
|
self.assertEqual(expected_value, gpus, msg_for_test_failure(expected_value)) |
|
|
|
|
|
|
|
|
def msg_for_test_failure(expected_value): |
|
|
return "GPU Values are expected to be " + str(expected_value) |
|
|
|
|
|
|
|
|
def create_args_parser(): |
|
|
parser = TrainerArgs().init_argparse(arg_prefix="") |
|
|
parser.add_argument("--gpus", type=str) |
|
|
return parser |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
unittest.main() |
|
|
|
