TensorFlowAOT/test/testAOTTools.py

0001 # coding: utf-8
0002
0003 """
0004 AOT compilation and dev workflow tests.
0005 """
0006
0007 import os
0008 import re
0009 import shlex
0010 import subprocess
0011 import tempfile
0012 import functools
0013 import unittest
0014 import platform
0015
0016
0017 this_dir = os.path.dirname(os.path.abspath(__file__))
0018
0019
0020 def run_cmd(cmd):
0021     if not isinstance(cmd, str):
0022         cmd = shlex.join(cmd)
0023     return subprocess.run(cmd, shell=True, check=True)
0024
0025
0026 def run_in_tmp(func):
0027     @functools.wraps(func)
0028     def wrapper(self):
0029         with tempfile.TemporaryDirectory() as tmp_dir:
0030             func(self, tmp_dir)
0031     return wrapper
0032
0033
0034 class TFAOTTests(unittest.TestCase):
0035
0036     @run_in_tmp
0037     def test_dev_workflow(self, tmp_dir):
0038         import cms_tfaot
0039
0040         # find the cms_tfaot install dir to locate the test model
0041         m = re.match(r"(.+/\d+\.\d+\.\d+\-[^/]+)/lib/.+$", cms_tfaot.__file__)
0042         self.assertIsNotNone(m)
0043         config_file = os.path.join(m.group(1), "share", "test_models", "simple", "aot_config.yaml")
0044         self.assertTrue(os.path.exists(config_file))
0045
0046         arch = "{0}-pc-linux".format(platform.processor())
0047
0048         # run the dev workflow
0049         # create the test model
0050         cmd = [
0051             "cms_tfaot_compile",
0052             "-c", config_file,
0053             "-o", tmp_dir,
0054             "--tool-name", "tfaot-model-test",
0055             "--dev",
0056             "--additional-flags=--target_triple=" + arch
0057         ]
0058         run_cmd(cmd)
0059
0060         # check files
0061         exists = lambda *p: os.path.exists(os.path.join(tmp_dir, *p))
0062         self.assertTrue(exists("tfaot-model-test.xml"))
0063         self.assertTrue(exists("include", "tfaot-model-test"))
0064         self.assertTrue(exists("include", "tfaot-model-test", "test_simple_bs1.h"))
0065         self.assertTrue(exists("include", "tfaot-model-test", "test_simple_bs2.h"))
0066         self.assertTrue(exists("include", "tfaot-model-test", "model.h"))
0067         self.assertTrue(exists("lib", "test_simple_bs1.o"))
0068         self.assertTrue(exists("lib", "test_simple_bs2.o"))
0069
0070
0071 if __name__ == "__main__":
0072     unittest.main()