Class JobRunner

caffe2/python/checkpoint.py:655–816

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 com

class JobRunner:
    """
    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.
    """
    def __init__(self, job, checkpoint_manager=None, resume_from_epoch=None,
                 upload_task_group_builder=None):
        """Initializes the JobRunner.

        Args:
            job: A Job object. The job to be executed.
            checkpoint_manager: Can be a CheckpointManager for single machine
                or a MultiNodeCheckpointManager for multi-machine. The manager
                that initializes/saves/loads checkpoints.
            resume_from_epoch: An integer. The epoch to resume from.
            upload_task_group_builder: A subclass of the
                UploadTaskGroupBuilder. Creates a task group to upload
                checkpoints.
        """
        self.resume_from_epoch = resume_from_epoch
        self.checkpoint_manager = checkpoint_manager
        self.job = job
        self.upload_task_group_builder = upload_task_group_builder

    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(