MCPcopy Create free account
hub / github.com/AI-Hypercomputer/maxtext / deactivate

Method deactivate

src/MaxText/profiler.py:99–125  ·  view source on GitHub ↗

End the profiler. The result is uploaded to the output bucket.

(self, blocking_object=None)

Source from the content-addressed store, hash-verified

97 self.deactivate(blocking_object=state)
98
99 def deactivate(self, blocking_object=None):
100 """End the profiler.
101 The result is uploaded to the output bucket."""
102 if self.profile_cleanly and blocking_object is not None:
103 jax.block_until_ready(blocking_object)
104
105 if self.managed_mldiagnostics and self.mode == "xplane":
106 # Handle the special profileing logic for managed_mldiagnostics
107 if self.prof is not None:
108 self.prof.stop()
109 return
110
111 if not (self.upload_all_profiler_results or jax.process_index() == 0):
112 return
113 if self.mode == "nsys":
114 if self.libcudart is not None:
115 self.libcudart.cudaProfilerStop()
116 else:
117 max_logging.log("WARNING: library for nsys was not loaded \n" "profiler has no effect")
118 return
119 # Popen() instead of run() for non-blocking behavior
120 if shutil.which("gsutil") is not None:
121 subprocess.Popen(["gsutil", "cp", "*nsys-rep", self.output_path]) # pylint: disable=consider-using-with
122 else:
123 max_logging.log("WARNING: gsutil is not installed or not found in the system's PATH. Skipping upload...")
124 elif self.mode == "xplane":
125 jax.profiler.stop_trace()
126
127 def _set_first_profiler_step(self, skip_steps, start_step):
128 return start_step + skip_steps

Callers 8

train_loopFunction · 0.95
mainFunction · 0.95
ar_benchmark_loopFunction · 0.95
train_loopFunction · 0.95
train_loopFunction · 0.95
generation_worker_fnFunction · 0.80

Calls

no outgoing calls

Tested by

no test coverage detected