Unverified Commit 876312f0 authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

[Core] Fix `gpu_worker.py` pre-commit errors (#35312)


Signed-off-by: default avatarNick Hill <nickhill123@gmail.com>
parent 5de98abc
...@@ -744,7 +744,8 @@ class Worker(WorkerBase): ...@@ -744,7 +744,8 @@ class Worker(WorkerBase):
# Create the profiler wrapper only on the first start call # Create the profiler wrapper only on the first start call
if self.profiler is None: if self.profiler is None:
if self.profiler_config.profiler == "torch": profiler_type = self.profiler_config.profiler
if profiler_type == "torch":
self.profiler = TorchProfilerWrapper( self.profiler = TorchProfilerWrapper(
self.profiler_config, self.profiler_config,
worker_name=trace_name, worker_name=trace_name,
...@@ -754,9 +755,12 @@ class Worker(WorkerBase): ...@@ -754,9 +755,12 @@ class Worker(WorkerBase):
logger.debug( logger.debug(
"Starting torch profiler with trace name: %s", trace_name "Starting torch profiler with trace name: %s", trace_name
) )
elif self.profiler_config.profiler == "cuda": elif profiler_type == "cuda":
self.profiler = CudaProfilerWrapper(self.profiler_config) self.profiler = CudaProfilerWrapper(self.profiler_config)
logger.debug("Starting CUDA profiler") logger.debug("Starting CUDA profiler")
else:
logger.warning("Unrecognized profiler: %s", profiler_type)
return
self.profiler.start() self.profiler.start()
else: else:
# Profiler already initialized. Restart profiling but keep # Profiler already initialized. Restart profiling but keep
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment