1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
# Owner(s): ["oncall: jit"]
import os
import re
import tempfile
import unittest
import torch._lazy
import torch._lazy.ts_backend
import torch.nn as nn
from torch.testing._internal.common_utils import IS_WINDOWS, run_tests, TestCase
torch._lazy.ts_backend.init()
@unittest.skipIf(IS_WINDOWS, "To be fixed")
class DebugUtilTest(TestCase):
def _run_linear(self):
device = "lazy"
model = nn.Linear(5, 5).to(device)
output = model(torch.randn(1, 5).to(device))
torch._lazy.mark_step()
def test_get_python_frames(self):
# We only care about the first "Python Stacktrace" part of the saved
# graph. However, we cannot save the whole stack for comparison given
# it depends on a lot of things.
partial_graph = (
r"Python Stacktrace:.*"
r"mark_step \(.*/_lazy/__init__.py:[0-9]+\).*"
r"_run_linear \(.*lazy/test_debug_util.py:[0-9]+\).*"
r"test_get_python_frames \(.*lazy/test_debug_util.py:[0-9]+\)"
)
with tempfile.NamedTemporaryFile(mode="r+", encoding="utf-8") as graph_file:
os.environ["LTC_SAVE_TENSORS_FILE"] = graph_file.name
self._run_linear()
file = graph_file.read()
if re.search(partial_graph, file, re.DOTALL) is None:
print(file)
self.assertTrue(False)
if __name__ == "__main__":
run_tests()
|
