[tune] Refactored Keras integration callbacks (#10509)

doc/source/conf.py

@@ -57,6 +57,7 @@ MOCK_MODULES = [
     "tensorflow.contrib.slim",
     "tensorflow.core",
     "tensorflow.core.util",
+    "tensorflow.keras",
     "tensorflow.python",
     "tensorflow.python.client",
     "tensorflow.python.util",
@@ -78,7 +79,9 @@ for mod_name in MOCK_MODULES:
 # ray.rllib.models.action_dist.py and
 # ray.rllib.models.lstm.py will use tf.VERSION
 sys.modules["tensorflow"].VERSION = "9.9.9"
+sys.modules["tensorflow.keras.callbacks"] = ChildClassMock()
 sys.modules["pytorch_lightning"] = ChildClassMock()
+
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.

doc/source/tune/api_docs/integration.rst

@@ -7,6 +7,15 @@ External library integrations (tune.integration)
     :local:
     :depth: 1
 
+.. _tune-integration-keras:
+
+Keras (tune.integration.keras)
+------------------------------------------------------
+
+.. autoclass:: ray.tune.integration.keras.TuneReportCallback
+
+.. autoclass:: ray.tune.integration.keras.TuneReportCheckpointCallback
+
 .. _tune-integration-kubernetes:
 
 Kubernetes (tune.integration.kubernetes)