diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index de0d52742..01298a753 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -126,6 +126,7 @@ class Envs: SGLANG_DISABLE_REQUEST_LOGGING = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) SGLANG_SIMULATE_ACC_METHOD = EnvStr("multinomial") + SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp") # Model Parallel SGLANG_USE_MESSAGE_QUEUE_BROADCASTER = EnvBool(True) diff --git a/test/srt/test_start_profile.py b/test/srt/test_start_profile.py index 60f5f7960..41342ef3f 100644 --- a/test/srt/test_start_profile.py +++ b/test/srt/test_start_profile.py @@ -9,6 +9,7 @@ import unittest import requests +from sglang.srt.environ import envs from sglang.srt.utils import kill_process_tree from sglang.test.test_utils import ( DEFAULT_SMALL_MODEL_NAME_FOR_TEST, @@ -25,6 +26,7 @@ class TestStartProfile(CustomTestCase): @classmethod def setUpClass(cls): + envs.SGLANG_TORCH_PROFILER_DIR.set(OUTPUT_DIR) cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( @@ -111,5 +113,4 @@ class TestStartProfile(CustomTestCase): if __name__ == "__main__": - os.environ["SGLANG_TORCH_PROFILER_DIR"] = OUTPUT_DIR unittest.main()