class_docstr: Implement the runtime logic for jobs with checkpointing at the level of
  epoch. Can be used to run either single-host or distributed jobs. Job
  runner is a callable to be called once from the master, passing a session
  as an argument. This call will block until the Job execution is complete.
  
  If a checkpoint_manager is passed, checkpoints will be taken after
  initialization and after each epoch execution. If, in addition,
  `resume_from_epoch` is an epoch number, the corresponding checkpoint will
  be loaded and job execution will continue from the given epoch. In
  this case, the job's init_group will not be run.
  
  Refer to checkpoint_test.py for an example.
class_docstr_tok: ["Implement","the","runtime","logic","for","jobs","with","checkpointing","at","the","level","of","epoch",".","Can","be","used","to","run","either","single-host","or","distributed","jobs",".","Job","runner","is","a","callable","to","be","called","once","from","the","master",",","passing","a","session","as","an","argument",".","This","call","will","block","until","the","Job","execution","is","complete",".","If","a","checkpoint_manager","is","passed",",","checkpoints","will","be","taken","after","initialization","and","after","each","epoch","execution",".","If",",","in","addition",",","`","resume_from_epoch","`","is","an","epoch","number",",","the","corresponding","checkpoint","will","be","loaded","and","job","execution","will","continue","from","the","given","epoch",".","In","this","case",",","the","job","'s","init_group","will","not","be","run",".","Refer","to","checkpoint_test.py","for","an","example","."]
docstr: Runs the training flow.
  
  Args:
      session: A Session object. Valid choises are: LocalSession,
          LocalHostScheduler, and DistributedSession. It is used to
          execute one TaskGroup a time.
docstr_tok: ["Runs","the","training","flow",".","Args",":","session",":","A","Session","object",".","Valid","choises","are",":","LocalSession",",","LocalHostScheduler",",","and","DistributedSession",".","It","is","used","to","execute","one","TaskGroup","a","time","."]
code: def train(self, session):
      """Runs the training flow.
  
      Args:
          session: A Session object. Valid choises are: LocalSession,
              LocalHostScheduler, and DistributedSession. It is used to
              execute one TaskGroup a time.
      """
      # identify the epoch we must resume from
      if self.checkpoint_manager:
          self.checkpoint_manager.set_params(
              nodes=self.job.nodes_to_checkpoint()
          )
          self.resume_from_epoch = (
              self.checkpoint_manager.get_resume_from_epoch_id(
                  self.resume_from_epoch
              )
          )
          if self.resume_from_epoch is not None:
              logger.info(
                  "Resuming from epoch {}".format(
                      self.resume_from_epoch
                  )
              )
  
      # Initialize all the nodes.
      from_scratch = self.resume_from_epoch is None
      if from_scratch:
          session.run(self.job.init_group)
  
      if self.checkpoint_manager:
          logger.info("Preparing checkpoints ...")
          session.run(
              self.checkpoint_manager.init(
                  self.job.nodes_to_checkpoint(),
                  retrieve_from_epoch=self.resume_from_epoch,
              )
          )
          # Save the first checkpoint before training starts, or resume from
          # a previously saved checkpoint.
          if from_scratch:
              self.save_checkpoints(0, session)
          else:
              logger.info(
                  "Loading checkpoints for epoch {} ...".format(
                      self.resume_from_epoch
                  )
              )
              session.run(
                  self.checkpoint_manager.load(self.resume_from_epoch)
              )
              self.checkpoint_manager.report_checkpoint_stats(
                  "checkpoint_load"
              )
              logger.info("Checkpoint loaded")
  
      logger.info("Finished initializing")
  
      # Start training.
      epoch = 1 if from_scratch else self.resume_from_epoch + 1
      while True:
          logger.info("Starting epoch %d" % epoch)
          session.run(self.job.epoch_group)
          logger.info("Finished epoch %d" % epoch)
          stop_conditions = [
              o.fetch() for o in self.job.stop_conditions
          ]
  
          if self.checkpoint_manager:
              self.save_checkpoints(epoch, session)
  
          if any(stop_conditions):
              logger.info("Stopping")
              break
          epoch += 1
      logger.info("Finished training")
      # Upload the checkpoints.
      if self.upload_task_group_builder:
          upload_task_group = self.upload_task_group_builder.build(
              epoch, self.checkpoint_manager
          )
          session.run(upload_task_group)
          logger.info("Finished uploading the checkpoints")
  
      # Download the parameters to save
      session.run(self.job.download_group)
      logger.info("Finished downloading the parameters")
  
      # Finally run the exit step to save nets
      session.run(self.job.exit_group)
      logger.info("Finished running the exit group")
      return epoch
  
imports: import os
  import logging
  from caffe2.python import core, context
  from caffe2.python.net_builder import ops
  from caffe2.python.task import final_output, Node, Task, TaskGroup, TaskOutput, WorkspaceType
label_desc: Use image node_id 2 for calling the JobRunner obj's underlying member method code with example usage: obj.train(session) and returns: epoch