[air] add to_air_checkpoint method for inference only workload.

python/ray/ml/predictors/integrations/tensorflow/__init__.py

@@ -1,5 +1,6 @@
 from ray.ml.predictors.integrations.tensorflow.tensorflow_predictor import (
     TensorflowPredictor,
 )
+from ray.ml.predictors.integrations.tensorflow.utils import to_air_checkpoint
 
-__all__ = ["TensorflowPredictor"]
+__all__ = ["TensorflowPredictor", "to_air_checkpoint"]

python/ray/ml/predictors/integrations/tensorflow/utils.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from tensorflow import keras
+
+from ray.ml.checkpoint import Checkpoint
+from ray.ml.constants import MODEL_KEY, PREPROCESSOR_KEY
+from ray.ml.preprocessor import Preprocessor
+
+
+def to_air_checkpoint(
+    model: keras.Model, preprocessor: Optional[Preprocessor] = None
+) -> Checkpoint:
+    """Convert a pretrained model to AIR checkpoint for serve or inference.
+
+    Args:
+        model: A pretrained model.
+        preprocessor: Stateless preprocessor only. The preprocessing logic will
+            be applied to serve/inference.
+        """
+    checkpoint = Checkpoint.from_dict(
+        {PREPROCESSOR_KEY: preprocessor, MODEL_KEY: model.get_weights()}
+    )
+    return checkpoint

python/ray/ml/tests/test_tensorflow_predictor.py

@@ -1,4 +1,6 @@
-from ray.ml.predictors.integrations.tensorflow import TensorflowPredictor
+import ray
+from ray.ml.batch_predictor import BatchPredictor
+from ray.ml.predictors.integrations.tensorflow import TensorflowPredictor, to_air_checkpoint
 from ray.ml.preprocessor import Preprocessor
 from ray.ml.checkpoint import Checkpoint
 from ray.ml.constants import PREPROCESSOR_KEY, MODEL_KEY
@@ -93,6 +95,15 @@ def test_predict_dataframe_with_feature_columns():
     assert predictions.to_numpy().flatten().tolist() == [1, 3]
 
 
+def test_tensorflow_predictor_no_training():
+    model = build_model()
+    checkpoint = to_air_checkpoint(model)
+    batch_predictor = BatchPredictor(checkpoint, TensorflowPredictor, model_definition=build_model)
+    predict_dataset = ray.data.range(3)
+    predictions = batch_predictor.predict(predict_dataset)
+    assert predictions.count() == 3
+
+
 if __name__ == "__main__":
     import pytest
     import sys