Open bveeramani opened 4 months ago
c21
commented
4 months ago Hi @bveeramani - to better understanding if this is a Ray Data only issue, can you check if same code works with Ray Core actor? If it does not work with Ray Core, it's a more fundamental issue with Tensorflow serializability. Otherwise, I want to understand how easy to fix on Ray Data side to make it work, thanks.
bveeramani
commented
4 months ago Serialization works if you directly use an actor.
@ray.remote
class TFPredictor:
...
TFPredictor.remote()Interestingly, serialization also works when you write a more minimal reproduction. I expected this to fail based on the traceback.
import cloudpickle
from ray.data._internal.planner.plan_udf_map_op import _parse_op_fn
from ray.data._internal.logical.operators.map_operator import AbstractUDFMap
fn, init_fn = _parse_op_fn(AbstractUDFMap("spam", None, TFPredictor))
cloudpickle.dumps(fn)
cloudpickle.dumps(init_fn)Something weird is going on, and it'll probably take more time to identify the root cause.
bveeramani
commented
4 months ago If you move the imports into the class, the program doesn't error.
class TFPredictor:
def __init__(self):
import tensorflow as tf
from tensorflow import keras
# Move the neural network to GPU by specifying the GPU device.
with tf.device("GPU:0"):
input_layer = keras.Input(shape=(100,))
output_layer = keras.layers.Dense(1, activation="sigmoid")
self.model = keras.Sequential([input_layer, output_layer])
def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
import tensorflow as tf
# Move the input batch to GPU by specifying GPU device.
with tf.device("GPU:0"):
return {"output": self.model(batch["data"]).numpy()}
What happened + What you expected to happen
Ran batch inference with a TensorFlow model. Got a cloudpickle error:
Versions / Dependencies
Ray: 77ee044a7e3c4d583f94956d7334fb3d6d7ce014 TensorFlow: 2.13-2.15 (2.12 and 2.16 work)
Reproduction script
Issue Severity
Medium: It is a significant difficulty but I can work around it.