Start the profiler. nsys profiler becomes no-op when libcudart.so is not available on the system.
(self, blocking_object=None, optional_postfix="")
| 59 | self.activate(blocking_object=state, optional_postfix=optional_postfix) |
| 60 | |
| 61 | def activate(self, blocking_object=None, optional_postfix=""): |
| 62 | """Start the profiler. |
| 63 | nsys profiler becomes no-op when libcudart.so is not available on the system.""" |
| 64 | if self.profile_cleanly and blocking_object is not None: |
| 65 | jax.block_until_ready(blocking_object) |
| 66 | |
| 67 | if self.managed_mldiagnostics and self.mode == "xplane": |
| 68 | # Handle the special profiling logic for managed_mldiagnostics |
| 69 | if self.prof is None: |
| 70 | # Starts xprof collector. |
| 71 | # Only profiling on the first device, if not upload_all_profiler_results. None is for all devices. |
| 72 | self.prof = mldiag.xprof(process_index_list=None if self.upload_all_profiler_results else [0]) |
| 73 | self.prof.start() |
| 74 | return |
| 75 | |
| 76 | if not (self.upload_all_profiler_results or jax.process_index() == 0): |
| 77 | return |
| 78 | if self.mode != "": |
| 79 | self.output_path = os.path.join(self.base_output_dir, optional_postfix) |
| 80 | if self.mode == "nsys": |
| 81 | try: |
| 82 | self.libcudart = cdll.LoadLibrary("libcudart.so") |
| 83 | except Exception as e: # pylint: disable=broad-except |
| 84 | max_logging.log(f"WARNING: Failed to load library for nsys: {e}\n" "profiler has no effect") |
| 85 | return |
| 86 | self.libcudart.cudaProfilerStart() |
| 87 | elif self.mode == "xplane": |
| 88 | jax.profiler.start_trace(self.output_path) |
| 89 | |
| 90 | def maybe_deactivate_profiler(self, step, state): |
| 91 | """Conditionally deactivates the profiler based on the current step. |
no test coverage detected