52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
# Owner(s): ["oncall: profiler"]
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
from torch.testing._internal.common_utils import run_tests, TestCase
|
|
|
|
|
|
class SimpleKinetoInitializationTest(TestCase):
|
|
@patch.dict(os.environ, {"KINETO_USE_DAEMON": "1"})
|
|
def test_kineto_profiler_with_environment_variable(self):
|
|
"""
|
|
This test checks whether kineto works with torch in daemon mode, please refer to issue #112389 and #131020.
|
|
Besides that, this test will also check that kineto will not be initialized when user loads the shared library
|
|
directly.
|
|
"""
|
|
script = """
|
|
import torch
|
|
if torch.cuda.is_available() > 0:
|
|
torch.cuda.init()
|
|
"""
|
|
try:
|
|
subprocess.check_output(
|
|
[sys.executable, "-W", "always", "-c", script],
|
|
cwd=os.path.dirname(os.path.realpath(__file__)),
|
|
)
|
|
except subprocess.CalledProcessError as e:
|
|
if e.returncode != 0:
|
|
self.assertTrue(
|
|
False,
|
|
"Kineto is not working properly with the Dynolog environment variable",
|
|
)
|
|
# import the shared library directly - it triggers static init but doesn't call kineto_init
|
|
env = os.environ.copy()
|
|
env["KINETO_USE_DAEMON"] = "1"
|
|
if "KINETO_DAEMON_INIT_DELAY_S" in env:
|
|
env.pop("KINETO_DAEMON_INIT_DELAY_S")
|
|
_, stderr = TestCase.run_process_no_exception(
|
|
f"from ctypes import CDLL; CDLL('{torch._C.__file__}')"
|
|
)
|
|
self.assertNotRegex(
|
|
stderr.decode("ascii"),
|
|
"Registering daemon config loader",
|
|
"kineto should not be initialized when the shared library is imported directly",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|