Add an easy way to run a script for a few steps only

I've wanted this tool for a while, figured I should just propose it.
Often I need to test out a script or colab I did not write, and just
want to run a few train steps without for every fit call without
finding every call to fit in the script. This adds a debugging tool
to do just that.

```
KERAS_MAX_EPOCHS=1 KERAS_MAX_STEPS=5 python train.py
```

Author: mattdangerw

keras/api/_tf_keras/keras/config/__init__.py

@@ -17,11 +17,15 @@
 from keras.src.backend.config import (
     is_flash_attention_enabled as is_flash_attention_enabled,
 )
+from keras.src.backend.config import max_epochs as max_epochs
+from keras.src.backend.config import max_steps as max_steps
 from keras.src.backend.config import set_epsilon as set_epsilon
 from keras.src.backend.config import set_floatx as set_floatx
 from keras.src.backend.config import (
     set_image_data_format as set_image_data_format,
 )
+from keras.src.backend.config import set_max_epochs as set_max_epochs
+from keras.src.backend.config import set_max_steps as set_max_steps
 from keras.src.dtype_policies.dtype_policy import dtype_policy as dtype_policy
 from keras.src.dtype_policies.dtype_policy import (
     set_dtype_policy as set_dtype_policy,

keras/api/config/__init__.py

@@ -17,11 +17,15 @@
 from keras.src.backend.config import (
     is_flash_attention_enabled as is_flash_attention_enabled,
 )
+from keras.src.backend.config import max_epochs as max_epochs
+from keras.src.backend.config import max_steps as max_steps
 from keras.src.backend.config import set_epsilon as set_epsilon
 from keras.src.backend.config import set_floatx as set_floatx
 from keras.src.backend.config import (
     set_image_data_format as set_image_data_format,
 )
+from keras.src.backend.config import set_max_epochs as set_max_epochs
+from keras.src.backend.config import set_max_steps as set_max_steps
 from keras.src.dtype_policies.dtype_policy import dtype_policy as dtype_policy
 from keras.src.dtype_policies.dtype_policy import (
     set_dtype_policy as set_dtype_policy,

keras/src/backend/config.py

@@ -15,6 +15,10 @@
 # Default backend: TensorFlow.
 _BACKEND = "tensorflow"
 
+# Cap run duration for debugging.
+_MAX_EPOCHS = None
+_MAX_STEPS = None
+
 
 @keras_export(["keras.config.floatx", "keras.backend.floatx"])
 def floatx():
@@ -304,7 +308,10 @@ def keras_home():
     _backend = os.environ["KERAS_BACKEND"]
     if _backend:
         _BACKEND = _backend
-
+if "KERAS_MAX_EPOCHS" in os.environ:
+    _MAX_EPOCHS = int(os.environ["KERAS_MAX_EPOCHS"])
+if "KERAS_MAX_STEPS" in os.environ:
+    _MAX_STEPS = int(os.environ["KERAS_MAX_STEPS"])
 
 if _BACKEND != "tensorflow":
     # If we are not running on the tensorflow backend, we should stop tensorflow
@@ -333,3 +340,66 @@ def backend():
 
     """
     return _BACKEND
+
+
+@keras_export(["keras.config.set_max_epochs"])
+def set_max_epochs(max_epochs):
+    """Limit the maximum number of epochs for any call to fit.
+
+    This will cap the number of epochs for any training run using `model.fit()`.
+    This is purely for debugging, and can also be set via the `KERAS_MAX_EPOCHS`
+    environment variable to quickly run a script without modifying its source.
+
+    Args:
+        max_epochs: The integer limit on the number of epochs or `None`. If
+            `None`, no limit is applied.
+    """
+    global _MAX_EPOCHS
+    _MAX_EPOCHS = max_epochs
+
+
+@keras_export(["keras.config.set_max_steps"])
+def set_max_steps(max_steps):
+    """Limit the maximum number of steps for any call to fit/evaluate/predict.
+
+    This will cap the number of steps for single epoch of a call to `fit()`,
+    `evaluate()`, or `predict()`. This is purely for debugging, and can also be
+    set via the `KERAS_MAX_STEPS` environment variable to quickly run a script
+    without modifying its source.
+
+    Args:
+        max_epochs: The integer limit on the number of epochs or `None`. If
+            `None`, no limit is applied.
+    """
+    global _MAX_STEPS
+    _MAX_STEPS = max_steps
+
+
+@keras_export(["keras.config.max_epochs"])
+def max_epochs():
+    """Get the maximum number of epochs for any call to fit.
+
+    Retrieves the limit on the number of epochs set by
+    `keras.config.set_max_epochs` or the `KERAS_MAX_EPOCHS` environment
+    variable.
+
+    Returns:
+        The integer limit on the number of epochs or `None`, if no limit has
+        been set.
+    """
+    return _MAX_EPOCHS
+
+
+@keras_export(["keras.config.max_steps"])
+def max_steps():
+    """Get the maximum number of steps for any call to fit/evaluate/predict.
+
+    Retrieves the limit on the number of epochs set by
+    `keras.config.set_max_steps` or the `KERAS_MAX_STEPS` environment
+    variable.
+
+    Args:
+        max_epochs: The integer limit on the number of epochs or `None`. If
+            `None`, no limit is applied.
+    """
+    return _MAX_STEPS

keras/src/backend/jax/trainer.py

@@ -9,6 +9,7 @@
 from keras.src import callbacks as callbacks_module
 from keras.src import optimizers as optimizers_module
 from keras.src import tree
+from keras.src.backend import config
 from keras.src.backend import distribution_lib as jax_distribution_lib
 from keras.src.distribution import distribution_lib
 from keras.src.trainers import trainer as base_trainer
@@ -341,6 +342,7 @@ def fit(
         validation_freq=1,
     ):
         self._assert_compile_called("fit")
+        epochs = config.max_epochs() or epochs
         # TODO: respect compiled trainable state
         self._eval_epoch_iterator = None
         if validation_split and validation_data is None:

keras/src/backend/tensorflow/trainer.py

@@ -9,6 +9,7 @@
 from keras.src import metrics as metrics_module
 from keras.src import optimizers as optimizers_module
 from keras.src import tree
+from keras.src.backend import config
 from keras.src.losses import loss as loss_module
 from keras.src.trainers import trainer as base_trainer
 from keras.src.trainers.data_adapters import array_slicing
@@ -309,6 +310,7 @@ def fit(
         validation_freq=1,
     ):
         self._assert_compile_called("fit")
+        epochs = config.max_epochs() or epochs
         # TODO: respect compiled trainable state
         self._eval_epoch_iterator = None
         if validation_split and validation_data is None:

keras/src/backend/torch/trainer.py

@@ -8,6 +8,7 @@
 from keras.src import callbacks as callbacks_module
 from keras.src import optimizers as optimizers_module
 from keras.src import tree
+from keras.src.backend import config
 from keras.src.trainers import trainer as base_trainer
 from keras.src.trainers.data_adapters import array_slicing
 from keras.src.trainers.data_adapters import data_adapter_utils
@@ -187,6 +188,7 @@ def fit(
             raise ValueError(
                 "You must call `compile()` before calling `fit()`."
             )
+        epochs = config.max_epochs() or epochs
 
         # TODO: respect compiled trainable state
         self._eval_epoch_iterator = None

keras/src/trainers/epoch_iterator.py

@@ -42,6 +42,7 @@
 import contextlib
 import warnings
 
+from keras.src.backend import config
 from keras.src.trainers import data_adapters
 
 
@@ -57,7 +58,7 @@ def __init__(
         class_weight=None,
         steps_per_execution=1,
     ):
-        self.steps_per_epoch = steps_per_epoch
+        self.steps_per_epoch = config.max_steps() or steps_per_epoch
         self.steps_per_execution = steps_per_execution
         self._current_iterator = None
         self._epoch_iterator = None

keras/src/trainers/trainer_test.py

@@ -14,6 +14,7 @@
 from keras.src import ops
 from keras.src import optimizers
 from keras.src import testing
+from keras.src.backend import config
 from keras.src.backend.common.symbolic_scope import in_symbolic_scope
 from keras.src.callbacks.callback import Callback
 from keras.src.optimizers.rmsprop import RMSprop
@@ -1506,6 +1507,48 @@ def test_steps_per_epoch(self, steps_per_epoch_test, mode):
             )
             self.assertAllClose(model.evaluate(x, y), model_2.evaluate(x, y))
 
+    @pytest.mark.requires_trainable_backend
+    def test_max_epochs_and_steps(self):
+        batch_size = 8
+        epochs = 4
+        num_batches = 10
+        data_size = num_batches * batch_size
+        x, y = np.ones((data_size, 4)), np.ones((data_size, 1))
+        model = ExampleModel(units=1)
+        model.compile(
+            loss="mse",
+            optimizer="sgd",
+            metrics=[EpochAgnosticMeanSquaredError()],
+        )
+        step_observer = StepObserver()
+        model.fit(
+            x=x,
+            y=y,
+            batch_size=batch_size,
+            epochs=epochs,
+            callbacks=[step_observer],
+            verbose=0,
+        )
+        self.assertEqual(step_observer.epoch_begin_count, epochs)
+        self.assertEqual(step_observer.begin_count, num_batches * epochs)
+        try:
+            config.set_max_epochs(2)
+            config.set_max_steps(3)
+            step_observer = StepObserver()
+            model.fit(
+                x=x,
+                y=y,
+                batch_size=batch_size,
+                epochs=epochs,
+                callbacks=[step_observer],
+                verbose=0,
+            )
+            self.assertEqual(step_observer.epoch_begin_count, 2)
+            self.assertEqual(step_observer.begin_count, 6)
+        finally:
+            config.set_max_epochs(None)
+            config.set_max_steps(None)
+
     @parameterized.named_parameters(
         named_product(
             steps_per_epoch_test=[