Keras fit() lacks native timeout parameter

Implementing a Time-Based Timeout

Because keras.fit() lacks a native timeout parameter, the most reliable way to enforce a time limit across Keras 2.x and 3.x—regardless of the backend (TensorFlow, JAX, or PyTorch)—is to implement a custom Callback that modifies the model.stop_training attribute.

The Implementation

The following implementation tracks the start time and signals the training loop to terminate at the end of the first epoch that exceeds the specified duration.

import time
from keras.callbacks import Callback

class TimeoutCallback(Callback):
    def __init__(self, seconds):
        super().__init__()
        self.seconds = seconds
        self.start_time = None

    def on_train_begin(self, logs=None):
        self.start_time = time.time()

    def on_epoch_end(self, epoch, logs=None):
        if self.start_time and (time.time() - self.start_time) > self.seconds:
            print(f"\nTimeout reached. Stopping training at epoch {epoch}.")
            self.model.stop_training = True

Ensuring Model Safety and Resource Release

To prevent the loss of progress and ensure a clean exit, you must combine the timeout callback with a checkpointing strategy. Simply stopping the loop does not automatically save the final state of the weights.

• Weight Persistence: Use keras.callbacks.ModelCheckpoint. This ensures that if a timeout occurs, the weights from the last completed epoch are already persisted to disk.
• Graceful Termination: Setting model.stop_training = True is the standard API method to halt training. It allows the fit() loop to finish its current iteration and execute any remaining on_epoch_end or on_train_end logic, preventing the inconsistent states associated with KeyboardInterrupt or OS-level process kills.
• Resource Management: Once fit() returns, the Keras backend handles the release of memory and device handles according to the specific backend's (TF, JAX, PyTorch) garbage collection and memory management rules.

Technical Assumptions and Constraints

Constraint	Behavior
Granularity	The timeout is checked only at the end of an epoch. If a single epoch takes longer than the timeout duration, the training will not stop until that epoch completes.
Backend Latency	In JAX or TensorFlow with asynchronous execution, there may be a slight delay between the stop_training trigger and the actual cessation of GPU/TPU activity.

Diagnostic Detail Needed: Are you using a very large batch size or a dataset where a single epoch exceeds your desired timeout window? If so, you may need to implement a custom training loop using GradientTape to check the timer at the batch level rather than the epoch level.