a callback to intentionally interrupt training def initself stepsint epochint self batchcount 0 self epochcount 0 self stepsint stepsint self epochint epochint def onepochendself epoch lognone self epochcount 1 if self epochint is not none and self epochcount self epochint raise runtimeerrorepochinterruption def onbatchendself batch logsnone self batchcount 1 if self stepsint is not none and self batchcount self stepsint raise runtimeerrorstepsinterruption class canarylayerlayers layer def initself super init self counter self addweight shape initializerzeros dtypefloat32 trainablefalse def callself x self counter assignadd1 return x class backupandrestorecallbacktesttesting testcase def makemodelself model sequential canarylayer layers dense1 model compile lossmse optimizersgd metricsmse return model check invalid savefreq both string and non integer def testsavefrequnknownerrorself with self assertraisesregexvalueerror expectedregexinvalid value callbacks backupandrestore backupdirbackupdir savefreqbatch with self assertraisesregexvalueerror expectedregexinvalid value callbacks backupandrestorebackupdirbackupdir savefreq0 15 checking if after interruption correct model params and weights are loaded in stepwise backup pytest mark requirestrainablebackend def testbestcasestepself tempdir self gettempdir backupdir fileutils jointempdir subdir self assertfalsefileutils existsbackupdir model self makemodel cbk callbacks backupandrestorebackupdir savefreq1 xtrain np random random10 3 ytrain np random random10 1 try model fit xtrain ytrain batchsize4 callbacks cbk interruptingcallbackstepsint2 epochintnone epochs2 verbose0 except runtimeerror self asserttruefileutils existsbackupdir self assertequalcbk currentepoch 0 self assertequalcbk lastbatchseen 1 self assertequalintmodel layers0 counter value 2 hist model fit xtrain ytrain batchsize4 callbackscbk epochs5 self assertequalcbk currentepoch 4 self assertequalhist epoch1 4 self assertequalintmodel layers0 counter value 17 checking if after interruption correct model params and weights are loaded in epochwise backup pytest mark requirestrainablebackend def testbestcaseepochself tempdir self gettempdir backupdir fileutils jointempdir subdir self assertfalsefileutils existsbackupdir model self makemodel self assertequalintmodel layers0 counter value 0 cbk callbacks backupandrestore backupdirbackupdir savefreqepoch xtrain np random random10 3 ytrain np random random10 1 try model fit xtrain ytrain batchsize4 callbacks cbk interruptingcallbackstepsintnone epochint2 epochs6 verbose0 except runtimeerror self assertequalcbk currentepoch 1 self asserttruefileutils existsbackupdir self assertequalintmodel layers0 counter value 6 hist model fit xtrain ytrain batchsize4 callbackscbk epochs5 self assertequalcbk currentepoch 4 self assertequalhist epoch1 4 self assertequalintmodel layers0 counter value 21 checking if after interruption when model is deleted pytest mark requirestrainablebackend def testmodeldeletedcaseepochself tempdir self gettempdir backupdir fileutils jointempdir subdir self assertfalsefileutils existsbackupdir model self makemodel cbk callbacks backupandrestorebackupdir savefreqepoch xtrain np random random10 3 ytrain np random random10 1 model fit xtrain ytrain batchsize4 callbackscbk epochs2 verbose0 self assertfalsefileutils existsbackupdir def testbackupdiremptyerrorself with self assertraisesregex valueerror expectedregexempty backupdir argument passed callbacks backupandrestorebackupdir savefreqepoch def testbackupdirnoneerrorself with self assertraisesregex valueerror expectedregexempty backupdir argument passed callbacks backupandrestorebackupdirnone savefreqepoch a callback to intentionally interrupt training check invalid save_freq both string and non integer checking if after interruption correct model params and weights are loaded in step wise backup checking if after interruption correct model params and weights are loaded in epoch wise backup checking if after interruption when model is deleted
import numpy as np
import pytest

from keras import callbacks
from keras import layers
from keras import testing
from keras.models import Sequential
from keras.utils import file_utils


class InterruptingCallback(callbacks.Callback):
    

    def __init__(self, steps_int, epoch_int):
        self.batch_count = 0
        self.epoch_count = 0
        self.steps_int = steps_int
        self.epoch_int = epoch_int

    def on_epoch_end(self, epoch, log=None):
        self.epoch_count += 1
        if self.epoch_int is not None and self.epoch_count == self.epoch_int:
            raise RuntimeError("EpochInterruption")

    def on_batch_end(self, batch, logs=None):
        self.batch_count += 1
        if self.steps_int is not None and self.batch_count == self.steps_int:
            raise RuntimeError("StepsInterruption")


class CanaryLayer(layers.Layer):
    def __init__(self):
        super().__init__()
        self.counter = self.add_weight(
            shape=(), initializer="zeros", dtype="float32", trainable=False
        )

    def call(self, x):
        self.counter.assign_add(1)
        return x


class BackupAndRestoreCallbackTest(testing.TestCase):
    def make_model(self):
        model = Sequential(
            [
                CanaryLayer(),
                layers.Dense(1),
            ]
        )
        model.compile(
            loss="mse",
            optimizer="sgd",
            metrics=["mse"],
        )
        return model

    
    def test_save_freq_unknown_error(self):
        with self.assertRaisesRegex(ValueError, expected_regex="Invalid value"):
            callbacks.BackupAndRestore(
                backup_dir="backup_dir", save_freq="batch"
            )

        with self.assertRaisesRegex(ValueError, expected_regex="Invalid value"):
            callbacks.BackupAndRestore(backup_dir="backup_dir", save_freq=0.15)

    
    
    @pytest.mark.requires_trainable_backend
    def test_best_case_step(self):
        temp_dir = self.get_temp_dir()
        backup_dir = file_utils.join(temp_dir, "subdir")
        self.assertFalse(file_utils.exists(backup_dir))

        model = self.make_model()
        cbk = callbacks.BackupAndRestore(backup_dir, save_freq=1)

        x_train = np.random.random((10, 3))
        y_train = np.random.random((10, 1))

        try:
            model.fit(
                x_train,
                y_train,
                batch_size=4,
                callbacks=[
                    cbk,
                    InterruptingCallback(steps_int=2, epoch_int=None),
                ],
                epochs=2,
                verbose=0,
            )
        except RuntimeError:
            self.assertTrue(file_utils.exists(backup_dir))
            self.assertEqual(cbk._current_epoch, 0)
            self.assertEqual(cbk._last_batch_seen, 1)
            self.assertEqual(int(model.layers[0].counter.value), 2)

            hist = model.fit(
                x_train, y_train, batch_size=4, callbacks=[cbk], epochs=5
            )

            self.assertEqual(cbk._current_epoch, 4)
            self.assertEqual(hist.epoch[-1], 4)
            self.assertEqual(int(model.layers[0].counter.value), 17)

    
    
    @pytest.mark.requires_trainable_backend
    def test_best_case_epoch(self):
        temp_dir = self.get_temp_dir()
        backup_dir = file_utils.join(temp_dir, "subdir")
        self.assertFalse(file_utils.exists(backup_dir))

        model = self.make_model()
        self.assertEqual(int(model.layers[0].counter.value), 0)
        cbk = callbacks.BackupAndRestore(
            backup_dir=backup_dir, save_freq="epoch"
        )

        x_train = np.random.random((10, 3))
        y_train = np.random.random((10, 1))

        try:
            model.fit(
                x_train,
                y_train,
                batch_size=4,
                callbacks=[
                    cbk,
                    InterruptingCallback(steps_int=None, epoch_int=2),
                ],
                epochs=6,
                verbose=0,
            )
        except RuntimeError:
            self.assertEqual(cbk._current_epoch, 1)
            self.assertTrue(file_utils.exists(backup_dir))
            self.assertEqual(int(model.layers[0].counter.value), 6)

            hist = model.fit(
                x_train, y_train, batch_size=4, callbacks=[cbk], epochs=5
            )
            self.assertEqual(cbk._current_epoch, 4)
            self.assertEqual(hist.epoch[-1], 4)
            self.assertEqual(int(model.layers[0].counter.value), 21)

    
    @pytest.mark.requires_trainable_backend
    def test_model_deleted_case_epoch(self):
        temp_dir = self.get_temp_dir()
        backup_dir = file_utils.join(temp_dir, "subdir")
        self.assertFalse(file_utils.exists(backup_dir))

        model = self.make_model()
        cbk = callbacks.BackupAndRestore(backup_dir, save_freq="epoch")

        x_train = np.random.random((10, 3))
        y_train = np.random.random((10, 1))
        model.fit(
            x_train,
            y_train,
            batch_size=4,
            callbacks=[cbk],
            epochs=2,
            verbose=0,
        )
        self.assertFalse(file_utils.exists(backup_dir))

    def test_backup_dir_empty_error(self):
        with self.assertRaisesRegex(
            ValueError, expected_regex="Empty `backup_dir` argument passed"
        ):
            callbacks.BackupAndRestore(backup_dir="", save_freq="epoch")

    def test_backup_dir_none_error(self):
        with self.assertRaisesRegex(
            ValueError, expected_regex="Empty `backup_dir` argument passed"
        ):
            callbacks.BackupAndRestore(backup_dir=None, save_freq="epoch")