Method batch_end

composer/profiler/json_trace_handler.py:278–367

(self, state: State, logger: Logger)

            self._save_at_batch_end = True

    def batch_end(self, state: State, logger: Logger) -> None:
        assert state.profiler is not None
        timestamp = state.timestamp
        trace_folder = format_name_with_dist(self.folder, run_name=state.run_name)
        if self._save_at_batch_end:
            # no longer active, but was previously active.
            # Epty the queue and save the trace file
            trace_filename = os.path.join(
                trace_folder,
                format_name_with_dist_and_time(self.filename, state.run_name, timestamp),
            )
            trace_dirname = os.path.dirname(trace_filename)
            if trace_dirname:
                os.makedirs(trace_dirname, exist_ok=True)
            with open(trace_filename, 'w+') as f:
                is_first_line = True
                f.write('[\n')
                while True:
                    try:
                        s = self._queue.get_nowait()
                    except queue.Empty:
                        break
                    if not is_first_line:
                        s = ',\n' + s
                    is_first_line = False
                    f.write(s)
                f.write('\n]\n')

            if self.remote_file_name is not None:
                remote_file_name = format_name_with_dist_and_time(self.remote_file_name, state.run_name, timestamp)
                logger.upload_file(
                    remote_file_name=remote_file_name,
                    file_path=trace_filename,
                    overwrite=self.overwrite,
                )
            # Gather the filenames
            trace_files = [pathlib.Path(x) for x in dist.all_gather_object(trace_filename)]
            self.saved_traces.append((timestamp, trace_files))

            # Ensure that all traces have been saved.
            dist.barrier()

            if self.merged_trace_filename is not None and dist.get_local_rank() == 0:
                # Merge together all traces from the node into one file
                start_rank = dist.get_global_rank()
                end_rank = dist.get_global_rank() + dist.get_local_world_size()
                trace_files_to_merge = trace_files[start_rank:end_rank]
                merged_trace_filename = os.path.join(
                    trace_folder,
                    format_name_with_dist(
                        self.merged_trace_filename,
                        state.run_name,
                    ),
                )
                merged_trace_dirname = os.path.dirname(merged_trace_filename)
                if merged_trace_dirname:
                    os.makedirs(merged_trace_dirname, exist_ok=True)