EMR compatibility issue with OnnxWrapper

[anqini]

Is there an existing issue for this?

• I have searched the existing issues and did not find a match.

Who can help?

I can help to contribute to fix the issue in the short term and collaborate in a long term.

What are you working on?

Testing Spark-nlp functionality with Onnx and Tensorflow components in the ML pipeline on EMR with multiple task nodes. We hope to onboard a few models trained by ourselves to spark nlp pipeline by using Onnx and tensorflow integration.

Current Behavior

I ran into the "the SparkSession should only be created on driver" issue caused by OnnxWrapper when testing spark-nlp on EMR. I suspected there was an issue in the code. Basically, the spark worker node was trying to get configuration from hadoop to initialize OnnxWrapper session. but It they cannot get sparkSession since it is not the driver.

Caused by: java.lang.IllegalStateException: SparkSession should only be created and accessed on the driver.
	at org.apache.spark.sql.SparkSession$.org$apache$spark$sql$SparkSession$$assertOnDriver(SparkSession.scala:1158)
	at org.apache.spark.sql.SparkSession$Builder.getOrCreate(SparkSession.scala:927)
	at com.johnsnowlabs.nlp.util.io.ResourceHelper$.$anonfun$getActiveSparkSession$1(ResourceHelper.scala:56)
	at scala.Option.getOrElse(Option.scala:189)
	at com.johnsnowlabs.nlp.util.io.ResourceHelper$.getActiveSparkSession(ResourceHelper.scala:57)
	at com.johnsnowlabs.nlp.util.io.ResourceHelper$.spark$lzycompute(ResourceHelper.scala:104)
	at com.johnsnowlabs.nlp.util.io.ResourceHelper$.spark(ResourceHelper.scala:104)
	at com.johnsnowlabs.util.ConfigHelper$.sparkSession$lzycompute(ConfigHelper.scala:23)
	at com.johnsnowlabs.util.ConfigHelper$.sparkSession(ConfigHelper.scala:23)
	at com.johnsnowlabs.util.ConfigHelper$.getConfigValueOrElse(ConfigHelper.scala:79)
	at com.johnsnowlabs.util.ConfigLoader$.getConfigInfo(ConfigLoader.scala:70)
	at com.johnsnowlabs.util.ConfigLoader$.configData$lzycompute(ConfigLoader.scala:36)
	at com.johnsnowlabs.util.ConfigLoader$.configData(ConfigLoader.scala:60)
	at com.johnsnowlabs.util.ConfigLoader$.getConfigIntValue(ConfigLoader.scala:79)
	at com.johnsnowlabs.ml.onnx.OnnxWrapper$.getCPUSessionConfig(OnnxWrapper.scala:195)

Expected Behavior

should run without exception either in worker node or driver node.

Steps To Reproduce

Step 1 create EMR cluster

we recommend create the cluster with EC2 RSA key, if you don't have it, you may want to create it first.

aws emr create-cluster \
    --release-label emr-6.8.0 \
    --use-default-roles \
    --instance-groups InstanceGroupType=MASTER,InstanceCount=1,InstanceType=m4.large InstanceGroupType=CORE,InstanceCount=1,InstanceType=m4.large InstanceGroupType=TASK,InstanceCount=5,InstanceType=m4.large \
    --applications Name=Spark Name=Zeppelin \
    --ec2-attributes KeyName=<yourKeyName>

Step 2 Find DNS in AWS Console

find the Primary node public DNS of your cluster. It looks like hadoop@*.compute-1.amazonaws.com.

Step 3 SSH into the cluster

ssh into the emr cluster

ssh -i <key_path> <nodepublicDNS>

Step 4

install and run pyspark shell

pip install spark-nlp==5.1.4

pyspark --packages com.johnsnowlabs.nlp:spark-nlp_2.12:5.1.4

Step 5

run code

from sparknlp.base import *
from sparknlp.annotator import *
from pyspark.ml import Pipeline

documentAssembler = DocumentAssembler() \
    .setInputCol("text") \
    .setOutputCol("document")

sentence = SentenceDetector() \
    .setInputCols(["document"]) \
    .setOutputCol("sentence")

tokenizer = Tokenizer() \
    .setInputCols(["sentence"]) \
    .setOutputCol("token")

# Use the transformer embeddings
embeddings = BertEmbeddings.pretrained(name='bert_base_cased', lang='en') \
    .setInputCols(['document', 'token']) \
    .setOutputCol('embeddings')

# This pretrained model requires those specific transformer embeddings
ner_model = NerDLModel.pretrained("ner_dl_bert", "en") \
    .setInputCols(["document", "token", "embeddings"]) \
    .setOutputCol("ner")

pipeline = Pipeline().setStages([
    documentAssembler,
    sentence,
    tokenizer,
    embeddings,
    ner_model
])

data = spark.createDataFrame([["U.N. official Ekeus heads for Baghdad."]]).toDF("text")
result = pipeline.fit(data).transform(data)

result.select("ner.result").show(truncate=False)

exception will be thrown when executing the lass line of code.

Spark NLP version and Apache Spark

'5.1.4'
'3.3.0-amzn-0'

Type of Spark Application

No response

Java Version

No response

Java Home Directory

No response

Setup and installation

No response

Operating System and Version

No response

Link to your project (if available)

No response

Additional Information

No response

[maziyarpanahi]

Thanks @anqini for the perfect step by step for us to reproduce this.

@danilojsl is this related to the new Spark NLP Config feature added recently to allow users set ONNX related configs?

PS: I have tested ONNX on EMR and Glue before we added this feature, so that's the first place I would look

[anqini]

FYI, the issue is related to running the onnx config on Task node instead of master node.
Here is the EMR node specification(Primary/Core/Task): https://docs.aws.amazon.com/emr/latest/ManagementGuide/emr-master-core-task-nodes.html

If there is only Master/Core node in the cluster such as small cluster or local testing environment, it probably works fine. however if you have multiple task nodes, the onnx configuration will fail.

[danilojsl]

Thanks @anqini for the perfect step by step for us to reproduce this.

@danilojsl is this related to the new Spark NLP Config feature added recently to allow users set ONNX related configs?

@maziyarpanahi Yes it raises in the code lines introduced for ONNX configs through spark session. According to @anqini explanation, the bug happens when there are multiple task nodes.

[anqini]

Hi team, i have run some ONNX relevant spark jobs including Embedding Annotator and Classifiers on my multi-tasknode EMR cluster. The spark jobs contain the example code snippets from spark-nlp websites. Some of them failed as shown in the screenshot. My interpretation is that the failed steps indicate the corresponding annotators use the default ONNX model while the successfully ones are using Tensorflow model. We could add fixing those annotators in the scope of this change. Thanks.

[maziyarpanahi]

Hi team, i have run some ONNX relevant spark jobs including Embedding Annotator and Classifiers on my multi-tasknode EMR cluster. The spark jobs contain the example code snippets from spark-nlp websites. Some of them failed as shown in the screenshot. My interpretation is that the failed steps indicate the corresponding annotators use the default ONNX model while the successfully ones are using Tensorflow model. We could add fixing those annotators in the scope of this change. Thanks.

Hi @anqini

First of all, thanks a lot for running these tests! You are very correct that they are by default ONNX. We have debugged and confirmed that the ONNX Config introduced in the last release is not compatible with multi-tasknode EMR cluster. I have never used this type of EMR Cluster myself, but we are on it to see what's different and how we can avoid this.

We are also working on addressing some bugs related to ONNX (E5 and MPNet accuracy and average_pooling), ONNX Config, and the Zip util we have for ONNX models for the upcoming release.

@danilojsl since this error doesn't happen for the rest of configs inside ConfigHelper, let's refactor the ONNX parts to avoid this error.

[rrader]

Hi! I have the same issue with the EMR.
Do you have a workaround? Can I downgrade to some version to get it to work?
Thank you!

[keithdewar]

@danilojsl @maziyarpanahi

Hi JSL, I'm having the same issue in Databricks on AWS. Sounds like you're on it but in case there's something different about my issue...thanks and best, Keith

I'm running on DBR 12.2 LTS ML (Spark = 3.3.2, Scala = 2.12, GPU). Spark-nlp = 5.1.4 + spark_nlp_gpu_assembly_5_1_4.jar. g5.xlarge, 2 to 8 workers

CODE

from sparknlp.annotator import *
from sparknlp.common import *
from sparknlp.base import *
from sparknlp.pretrained import *
from sparknlp.training import *

spark=SparkSession.builder.appName("pysparkdf").getOrCreate()

documentAssembler = DocumentAssembler() \
    .setInputCol("text") \
    .setOutputCol("document")

tokenizer = Tokenizer() \
    .setInputCols(["document"]) \
    .setOutputCol("token")

bert_loaded = BertEmbeddings.pretrained("bert_embeddings_bge_large","en") \
    .setInputCols(["document", "token"]) \
    .setOutputCol("embeddings") \
    .setCaseSensitive(True)
    
pipeline = Pipeline(stages=[documentAssembler, tokenizer, bert_loaded])

data = spark.createDataFrame([["I love Spark NLP"]]).toDF("text")

predictionDF = model.transform(data)

result = pipeline.fit(data).transform(data)

result.display()`

ERROR (exception thrown on the last line of code)

bert_embeddings_bge_large download started this may take some time.
Approximate size to download 758.2 MB
[OK!]

IllegalStateException: SparkSession should only be created and accessed on the driver.
org.apache.spark.SparkException: Job aborted due to stage failure: Task 2 in stage 26.0 failed 4 times, most recent failure: Lost task 2.3 in stage 26.0 (TID 139) (100.65.8.104 executor 0): java.lang.IllegalStateException: SparkSession should only be created and accessed on the driver.
at org.apache.spark.sql.SparkSession$.org$apache$spark$sql$SparkSession$$assertOnDriver(SparkSession.scala:1457)
at org.apache.spark.sql.SparkSession$Builder.getOrCreate(SparkSession.scala:1217)
at com.johnsnowlabs.nlp.util.io.ResourceHelper$.$anonfun$getActiveSparkSession$1(ResourceHelper.scala:56)
at scala.Option.getOrElse(Option.scala:189)
at com.johnsnowlabs.nlp.util.io.ResourceHelper$.getActiveSparkSession(ResourceHelper.scala:57)
at com.johnsnowlabs.nlp.util.io.ResourceHelper$.spark$lzycompute(ResourceHelper.scala:104)
at com.johnsnowlabs.nlp.util.io.ResourceHelper$.spark(ResourceHelper.scala:104)
at com.johnsnowlabs.util.ConfigHelper$.sparkSession$lzycompute(ConfigHelper.scala:23)
at com.johnsnowlabs.util.ConfigHelper$.sparkSession(ConfigHelper.scala:23)
at com.johnsnowlabs.util.ConfigHelper$.getConfigValueOrElse(ConfigHelper.scala:79)
at com.johnsnowlabs.util.ConfigLoader$.getConfigInfo(ConfigLoader.scala:70)
at com.johnsnowlabs.util.ConfigLoader$.configData$lzycompute(ConfigLoader.scala:36)
at com.johnsnowlabs.util.ConfigLoader$.configData(ConfigLoader.scala:60)
at com.johnsnowlabs.util.ConfigLoader$.getConfigIntValue(ConfigLoader.scala:79)
at com.johnsnowlabs.ml.onnx.OnnxWrapper$.getCUDASessionConfig(OnnxWrapper.scala:151)
at com.johnsnowlabs.ml.onnx.OnnxWrapper$.com$johnsnowlabs$ml$onnx$OnnxWrapper$$withSafeOnnxModelLoader(OnnxWrapper.scala:90)
at com.johnsnowlabs.ml.onnx.OnnxWrapper.getSession(OnnxWrapper.scala:47)
at com.johnsnowlabs.ml.ai.Bert.tag(Bert.scala:93)
at com.johnsnowlabs.ml.ai.Bert.$anonfun$predict$1(Bert.scala:360)
at scala.collection.Iterator$$anon$11.nextCur(Iterator.scala:486)

[keithdewar]

@anqini did you find a work around in the meantime?