wehub-resource-sync
55 lines
1.4 KiB
Python
55 lines
1.4 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import logging
|
|
import signal
|
|
import time
|
|
import unittest
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from fairseq.distributed import DistributedTimeoutWrapper
|
|
|
|
|
|
class ModuleWithDelay(nn.Module):
|
|
|
|
def __init__(self, delay):
|
|
super().__init__()
|
|
self.delay = delay
|
|
|
|
def forward(self, x):
|
|
time.sleep(self.delay)
|
|
return x
|
|
|
|
|
|
class TestDistributedTimeoutWrapper(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
logging.disable(logging.CRITICAL)
|
|
|
|
def tearDown(self):
|
|
logging.disable(logging.NOTSET)
|
|
|
|
def test_no_timeout(self):
|
|
module = DistributedTimeoutWrapper(ModuleWithDelay(1), 0, signal.SIGINT)
|
|
module(torch.rand(5))
|
|
module.stop_timeout()
|
|
|
|
def test_timeout_safe(self):
|
|
module = DistributedTimeoutWrapper(ModuleWithDelay(1), 10, signal.SIGINT)
|
|
module(torch.rand(5))
|
|
module.stop_timeout()
|
|
|
|
def test_timeout_killed(self):
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
module = DistributedTimeoutWrapper(ModuleWithDelay(5), 1, signal.SIGINT)
|
|
module(torch.rand(5))
|
|
module.stop_timeout()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|