Support for Spark DL notebooks with PyTriton on Databricks/Dataproc #483
Conversation
There was a problem hiding this comment.
Looks good overall. A few comments.
In a future optimization we can look at something like https://github.com/triton-inference-server/client/blob/main/src/python/examples/simple_http_cudashm_client.py or for regular shm to reduce data copy (if I'm interpreting these correctly).
| sudo /databricks/python3/bin/pip3 install --upgrade --force-reinstall -r temp_requirements.txt | ||
| rm temp_requirements.txt | ||
|
|
||
| set +x |
There was a problem hiding this comment.
Add a carriage return at the end of last line in all files this symbol appears.
There was a problem hiding this comment.
Deleted, also merged the tf/torch scripts into one for convenience.
| "df = spark.read.parquet(\"imdb_test\").limit(100).cache()" | ||
| "def _use_stage_level_scheduling(spark, rdd):\n", | ||
| "\n", | ||
| " if spark.version < \"3.4.0\":\n", |
There was a problem hiding this comment.
This check probably not needed since predict_batch_udf is also not in spark < 3.4
| "metadata": {}, | ||
| "outputs": [], | ||
| "source": [ | ||
| "df = spark.read.parquet(data_path).limit(256).repartition(8)" |
There was a problem hiding this comment.
Is limit and repartition needed? And is this the right order? And why these numbers? A comment might be in order. Propagate any changes to other notebooks.
There was a problem hiding this comment.
This was intended to test the minimal scenario of 1 batch per task—especially with tensorflow, too high of a number can be really slow (>1 min). (In previous versions we were limiting to 100 rows: https://github.com/NVIDIA/spark-rapids-examples/blob/branch-23.06/examples/ML%2BDL-Examples/Spark-DL/dl_inference/huggingface/conditional_generation.ipynb?short_path=d3949f8#L1208)
| ] | ||
| }, | ||
| { | ||
| "cell_type": "code", | ||
| "execution_count": 56, | ||
| "execution_count": null, |
There was a problem hiding this comment.
fyi, spark.stop() below might be bad for databricks. It puts the cluster in a bad state. (at least in older versions like 13.3 from what I've seen).
There was a problem hiding this comment.
Yup, issue persists on latest runtime - addressed
| "def stop_triton(it):\n", | ||
| " import docker\n", | ||
| " import time\n", | ||
| "def stop_triton(pids):\n", |
There was a problem hiding this comment.
Can this along with all the other triton related code that is common across the notebooks be moved to a single python file triton_utils.py that gets shipped via pyfiles with each Spark job and then imported in the notebooks? Would avoid a lot of repetition.
Good idea, will definitely follow-up with this improvement. Note per pytriton team—with shm, there still will be an additional inter-process data copy (until Triton 3 release): |
Support for running DL Inference notebooks on CSP environments.
Notebook outputs are saved from running locally, but all notebooks were tested on Databricks/Dataproc.