Flink Torch
I want to explore if we can leverage PyFlink to run ML batch prediction performant. Here we go.
What're the files?
ml-100k/
: Movielens 100K dataset. It is pretty much the "Hello world" in recsys world.model_def.py
: A naive Matrix Factorization implementation, shamelessly copied from Kaggle, and polished following PyTorch quickstart.pytorch_movielens.py
: It trains the model, saves it, loads it back and runs a prediction.model.pth
: A saved model from runningpytorch_movielens.py
.udf_def.py
: PyFlink UDF that loads the model againsts which it runs batch prediction.test_udf_predict.py
: Unit test of the PyFlink UDF.
Try it!
Steps:
- Make sure you have
conda
installed. - Create a conda env
flink-ml
withconda env create -f environment.yml
, and activate withconda activate flink-ml
. - (Optional) Re-use the existing trained model
model.pth
, or re-train the model withpython pytorch_movielens.py
. - Run batch prediction with
python batch_prediction.py
. You should find predictions in your standard out. It says, user 269 may love movies[316,486,664,729,1020]
best, and user 270 may love movies[77,83,306,800,1014]
best.
...
5> +U(269,316,486,664,729,1020)
5> +U(270,77,83,306,800,1014)
Beyond local
Note all commands are run under the root path of the project.
bundle your resources
zip -r resources.zip resources
bundle your conda env
(cd /usr/local/anaconda3/envs/flink-ml && zip -r - .) > flink-ml.zip
submit
~/softwares/flink-1.12.0/bin/flink run -d \
-pyexec flink-ml/bin/python \
-pyarch resources.zip,flink-ml.zip#flink-ml \
-pyfs udf_def.py,model_def.py \
-py batch_prediction.py