class_docstr: Records a history of memory allocation and free events, and dumps to a
  file which can be visualized offline. It by default keeps track of
  100000 events, and dumps at user specified step as well as on OOM.
  The profiler is ideally started before model is instantiated in memory
  so that allocations that stay constant in training are accounted for.
  
  Args:
      output_dir: Directory where to save the memory snapshots.
      memory_snapshot_params: Instance of MemorySnapshotParams.
  
  Raises:
      ValueError: If `start_step` is negative, or `stop_step` is less than or equal to zero.
      ValueError: If `start_step` is greater than or equal to `stop_step`.
      ValueError: If `start_step` is set and `stop_step` is not set.
      ValueError: If `stop_step` is set and neither `start_step` nor `enable_oom_observer` are set.
      ValueError: If `enable_oom_observer` is False and neither `start_step` nor `stop_step` is set
  
  Examples::
      memory_snapshot_params = MemorySnapshotParams(start_step=5, stop_step=10, enable_oom_observer=True)
      memory_snapshot_profiler = MemorySnapshotProfiler(output_dir="/tmp", memory_snapshot_params=memory_snapshot_params)
      for batch in dataloader:
          ...
          memory_snapshot_profiler.step()
class_docstr_tok: ["Records","a","history","of","memory","allocation","and","free","events",",","and","dumps","to","a","file","which","can","be","visualized","offline",".","It","by","default","keeps","track","of","100000","events",",","and","dumps","at","user","specified","step","as","well","as","on","OOM",".","The","profiler","is","ideally","started","before","model","is","instantiated","in","memory","so","that","allocations","that","stay","constant","in","training","are","accounted","for",".","Args",":","output_dir",":","Directory","where","to","save","the","memory","snapshots",".","memory_snapshot_params",":","Instance","of","MemorySnapshotParams",".","Raises",":","ValueError",":","If","`","start_step","`","is","negative",",","or","`","stop_step","`","is","less","than","or","equal","to","zero",".","ValueError",":","If","`","start_step","`","is","greater","than","or","equal","to","`","stop_step","`",".","ValueError",":","If","`","start_step","`","is","set","and","`","stop_step","`","is","not","set",".","ValueError",":","If","`","stop_step","`","is","set","and","neither","`","start_step","`","nor","`","enable_oom_observer","`","are","set",".","ValueError",":","If","`","enable_oom_observer","`","is","False","and","neither","`","start_step","`","nor","`","stop_step","`","is","set","Examples",":",":","memory_snapshot_params","=","MemorySnapshotParams","(","start_step=5",",","stop_step=10",",","enable_oom_observer=True",")","memory_snapshot_profiler","=","MemorySnapshotProfiler","(","output_dir=","''","\/tmp","''",",","memory_snapshot_params=memory_snapshot_params",")","for","batch","in","dataloader",":","...","memory_snapshot_profiler.step","(",")"]
code: def __init__(
      self,
      output_dir: str,
      memory_snapshot_params: Optional[MemorySnapshotParams] = None,
  ) -> None:
      self.output_dir: str = output_dir
      self.params: MemorySnapshotParams = (
          memory_snapshot_params or MemorySnapshotParams()
      )
      start_step = self.params.start_step
      stop_step = self.params.stop_step
      if start_step is not None:
          if start_step < 0:
              raise ValueError("start_step must be nonnegative.")
          elif stop_step is None:
              raise ValueError(
                  "stop_step must be specified when start_step is set."
              )
          elif start_step >= stop_step:
              raise ValueError("start_step must be < stop_step.")
      if stop_step is not None:
          if stop_step <= 0:
              raise ValueError("stop_step must be positive.")
          elif (
              start_step is None and not self.params.enable_oom_observer
          ):
              raise ValueError(
                  "stop_step must be enabled with either start_step or enable_oom_observer."
              )
      if (
          start_step is None
          and stop_step is None
          and not self.params.enable_oom_observer
      ):
          raise ValueError(
              "At least one of start_step/stop_step or enable_oom_observer must be set."
          )
  
      self.step_num: int = 0
  
      if not is_torch_version_geq_2_0():
          raise RuntimeError("CUDA memory snapshot requires torch>=2.0")
      if self.params.enable_oom_observer:
          attach_oom_observer(
              output_dir=output_dir,
              trace_max_entries=self.params.max_entries,
          )
      if start_step is not None and start_step == 0:
          self.start()
  
      logger.info(
          f"Created MemorySnapshotProfiler with MemorySnapshotParams={self.params}."
      )
  
file_path: tnt/torchtnt/utils/memory_snapshot_profiler.py
imports: import logging
  from dataclasses import dataclass
  from types import TracebackType
  from typing import Optional, Type
  import torch
  from torchtnt.utils.oom import attach_oom_observer, log_memory_snapshot
  from torchtnt.utils.version import is_torch_version_geq_2_0
label_desc: Use image node_id 1 to create a new MemorySnapshotProfiler object  with example: obj = MemorySnapshotProfiler(output_dir, memory_snapshot_params)