[Profiler] Add SGLANG_PROFILE_RECORD_SHAPES for recording shapes when profiling (#11641)
Signed-off-by: zejunchen-zejun <zejun.chen@amd.com> Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
@@ -99,6 +99,7 @@ SGLang supports various environment variables that can be used to configure its
|
||||
| --- | --- | --- |
|
||||
| `SGLANG_TORCH_PROFILER_DIR` | Directory for PyTorch profiler output | `/tmp` |
|
||||
| `SGLANG_PROFILE_WITH_STACK` | Set `with_stack` option (bool) for PyTorch profiler (capture stack trace) | `true` |
|
||||
| `SGLANG_PROFILE_RECORD_SHAPES` | Set `record_shapes` option (bool) for PyTorch profiler (record shapes) | `true` |
|
||||
| `SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS` | Config BatchSpanProcessor.schedule_delay_millis if tracing is enabled | `500` |
|
||||
| `SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE` | Config BatchSpanProcessor.max_export_batch_size if tracing is enabled | `64` |
|
||||
|
||||
|
||||
@@ -136,6 +136,7 @@ class Envs:
|
||||
SGLANG_IS_IN_CI_AMD = EnvBool(False)
|
||||
SGLANG_SET_CPU_AFFINITY = EnvBool(False)
|
||||
SGLANG_PROFILE_WITH_STACK = EnvBool(True)
|
||||
SGLANG_PROFILE_RECORD_SHAPES = EnvBool(True)
|
||||
SGLANG_RECORD_STEP_TIME = EnvBool(False)
|
||||
SGLANG_FORCE_SHUTDOWN = EnvBool(False)
|
||||
SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
|
||||
|
||||
@@ -319,6 +319,10 @@ class TokenizerCommunicatorMixin:
|
||||
self.auto_create_handle_loop()
|
||||
env_with_stack: bool = get_bool_env_var("SGLANG_PROFILE_WITH_STACK", "true")
|
||||
with_stack = False if with_stack is False or env_with_stack is False else True
|
||||
env_record_shapes: bool = get_bool_env_var(
|
||||
"SGLANG_PROFILE_RECORD_SHAPES", "true"
|
||||
)
|
||||
record_shapes = (record_shapes is not False) and env_record_shapes
|
||||
req = ProfileReq(
|
||||
type=ProfileReqType.START_PROFILE,
|
||||
output_dir=output_dir,
|
||||
|
||||
Reference in New Issue
Block a user