Method oom_observer

composer/callbacks/oom_observer.py:135–186

(device: int, alloc: int, device_alloc: int, device_free: int)

                ensure_folder_is_empty(self.folder_name)

        def oom_observer(device: int, alloc: int, device_alloc: int, device_free: int):
            # Snapshot right after an OOM happened
            log.warning('Out Of Memory (OOM) observed')

            assert self.filename
            assert self.folder_name, 'folder_name must be set in init'
            filename = Path(self.folder_name) / Path(
                format_name_with_dist_and_time(self.filename, run_name=state.run_name, timestamp=state.timestamp),
            )

            try:
                self.filename_config = SnapshotFileNameConfig.from_file_name(str(filename))
                log.info(f'Dumping OOMObserver visualizations')

                snapshot = torch.cuda.memory._snapshot()
                # No data was recorded - avoids a `ValueError` in `trace_plot`
                if all(len(t) == 0 for t in snapshot['device_traces']):
                    log.info(f'No allocation is recorded in memory snapshot)')
                    return

                with open(self.filename_config.snapshot_file, 'wb') as fd:
                    pickle.dump(snapshot, fd)

                with open(self.filename_config.trace_plot_file, 'w+') as fd:
                    fd.write(torch.cuda._memory_viz.trace_plot(snapshot))  # type: ignore

                with open(self.filename_config.segment_plot_file, 'w+') as fd:
                    fd.write(torch.cuda._memory_viz.segment_plot(snapshot))  # type: ignore

                with open(self.filename_config.segment_flamegraph_file, 'w+') as fd:
                    fd.write(torch.cuda._memory_viz.segments(snapshot))  # type: ignore

                with open(self.filename_config.memory_flamegraph_file, 'w+') as fd:
                    fd.write(torch.cuda._memory_viz.memory(snapshot))  # type: ignore

                log.info(f'Saved memory visualizations to local files with prefix = {filename} during OOM')

                if self.remote_path_in_bucket is not None:
                    for f in self.filename_config.list_filenames():
                        base_file_name = os.path.basename(f)
                        remote_file_name = os.path.join(self.remote_path_in_bucket, base_file_name)
                        remote_file_name = remote_file_name.lstrip('/')  # remove leading slashes
                        log.info(f'Uploading memory visualization to remote: {remote_file_name} from {f}')
                        try:
                            logger.upload_file(remote_file_name=remote_file_name, file_path=f, overwrite=self.overwrite)
                        except FileExistsError as e:
                            raise FileExistsError(
                                f'Uploading memory visualizations failed with error: {e}. overwrite was set to {self.overwrite}. To overwrite memory visualizations with Trainer, set save_overwrite to True.',
                            ) from e

            except Exception as e:
                log.error(f'Failed to capture memory snapshot {e}')

        if self._enabled:
            torch.cuda.memory._record_memory_history(