class_docstr: A simple callback that reports checkpoints to Ray on train epoch end.
  
  This callback is a subclass of `lightning.pytorch.callbacks.Callback
  <https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.Callback.html#lightning.pytorch.callbacks.Callback>`_.
  
  It fetches the latest `trainer.callback_metrics` and reports together with
  the checkpoint on each training epoch end.
  
  Checkpoints will be saved in the following structure::
  
      checkpoint_00000*/      Ray Train Checkpoint
      └─ checkpoint.ckpt      PyTorch Lightning Checkpoint
  
  For customized reporting and checkpointing logic, implement your own
  `lightning.pytorch.callbacks.Callback` following this user
  guide: :ref:`Saving and Loading Checkpoints <train-dl-saving-checkpoints>`.
class_docstr_tok: ["A","simple","callback","that","reports","checkpoints","to","Ray","on","train","epoch","end",".","This","callback","is","a","subclass","of","`","lightning.pytorch.callbacks.Callback","<","https",":","\/\/lightning.ai\/docs\/pytorch\/stable\/api\/lightning.pytorch.callbacks.Callback.html","#","lightning.pytorch.callbacks.Callback",">","`","_",".","It","fetches","the","latest","`","trainer.callback_metrics","`","and","reports","together","with","the","checkpoint","on","each","training","epoch","end",".","Checkpoints","will","be","saved","in","the","following","structure",":",":","checkpoint_00000","*","\/","Ray","Train","Checkpoint","\u2514\u2500","checkpoint.ckpt","PyTorch","Lightning","Checkpoint","For","customized","reporting","and","checkpointing","logic",",","implement","your","own","`","lightning.pytorch.callbacks.Callback","`","following","this","user","guide",":",":","ref",":","`","Saving","and","Loading","Checkpoints","<","train-dl-saving-checkpoints",">","`","."]
code: def on_train_epoch_end(self, trainer, pl_module) -> None:
      # Creates a checkpoint dir with fixed name
      tmpdir = os.path.join(
          self.tmpdir_prefix, str(trainer.current_epoch)
      )
      os.makedirs(tmpdir, exist_ok=True)
  
      # Fetch metrics
      metrics = trainer.callback_metrics
      metrics = {k: v.item() for k, v in metrics.items()}
  
      # (Optional) Add customized metrics
      metrics["epoch"] = trainer.current_epoch
      metrics["step"] = trainer.global_step
  
      # Save checkpoint to local
      ckpt_path = os.path.join(tmpdir, self.CHECKPOINT_NAME)
      trainer.save_checkpoint(ckpt_path, weights_only=False)
  
      # Report to train session
      checkpoint = Checkpoint.from_directory(tmpdir)
      train.report(metrics=metrics, checkpoint=checkpoint)
  
      # Add a barrier to ensure all workers finished reporting here
      torch.distributed.barrier()
  
      if self.local_rank == 0:
          shutil.rmtree(tmpdir)
  
imports: import logging
  import os
  import shutil
  import tempfile
  from typing import Any, Dict
  import torch
  from packaging.version import Version
  import ray
  from ray import train
  from ray._private.usage.usage_lib import TagKey, record_extra_usage_tag
  from ray.train import Checkpoint
  from ray.util import PublicAPI
label_desc: Use image node_id 2 for calling the RayTrainReportCallback obj's underlying member method code with example usage: obj.on_train_epoch_end(trainer, pl_module) without return types