This is an automated email from the ASF dual-hosted git repository. gurwls223 pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push: new 400573d33af [SPARK-45093][CONNECT][PYTHON] Properly support error handling and conversion for AddArtifactHandler 400573d33af is described below commit 400573d33af39d61748dd0b2960c47161277f85e Author: Martin Grund <martin.gr...@databricks.com> AuthorDate: Mon Dec 4 08:22:13 2023 +0900 [SPARK-45093][CONNECT][PYTHON] Properly support error handling and conversion for AddArtifactHandler ### What changes were proposed in this pull request? This patch improves the error handling when errors are happening in the `AddArtifact` path. In particular, the `AddArtifactHandler` would not properly return exceptions but all exceptions would end up yielding `UNKNOWN` errors. This patch makes sure we wrap the errors in the add artifact path the same way as we're wrapping the errors in the normal query execution path. In addition, it adds tests and verification for this behavior in Scala and in Python. ### Why are the changes needed? Stability ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Added UT for Scala and Python ### Was this patch authored or co-authored using generative AI tooling? No Closes #44092 from grundprinzip/SPARK-ADD_ARTIFACT_CRAP. Authored-by: Martin Grund <martin.gr...@databricks.com> Signed-off-by: Hyukjin Kwon <gurwls...@apache.org> --- .../service/SparkConnectAddArtifactsHandler.scala | 26 +++++++++++++++++++--- .../spark/sql/connect/utils/ErrorUtils.scala | 25 +++++++++++++++++++-- .../connect/service/AddArtifactsHandlerSuite.scala | 20 +++++++++++++++-- python/pyspark/sql/connect/client/core.py | 15 ++++++++----- .../sql/tests/connect/client/test_artifact.py | 20 +++++++++++++++++ 5 files changed, 93 insertions(+), 13 deletions(-) diff --git a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala index e664e07dce1..ea3b578be3b 100644 --- a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala +++ b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAddArtifactsHandler.scala @@ -24,7 +24,6 @@ import scala.collection.mutable import scala.util.control.NonFatal import com.google.common.io.CountingOutputStream -import io.grpc.StatusRuntimeException import io.grpc.stub.StreamObserver import org.apache.spark.connect.proto @@ -32,6 +31,7 @@ import org.apache.spark.connect.proto.{AddArtifactsRequest, AddArtifactsResponse import org.apache.spark.connect.proto.AddArtifactsResponse.ArtifactSummary import org.apache.spark.sql.artifact.ArtifactManager import org.apache.spark.sql.artifact.util.ArtifactUtils +import org.apache.spark.sql.connect.utils.ErrorUtils import org.apache.spark.util.Utils /** @@ -51,7 +51,7 @@ class SparkConnectAddArtifactsHandler(val responseObserver: StreamObserver[AddAr private var chunkedArtifact: StagedChunkedArtifact = _ private var holder: SessionHolder = _ - override def onNext(req: AddArtifactsRequest): Unit = { + override def onNext(req: AddArtifactsRequest): Unit = try { if (this.holder == null) { this.holder = SparkConnectService.getOrCreateIsolatedSession( req.getUserContext.getUserId, @@ -78,6 +78,17 @@ class SparkConnectAddArtifactsHandler(val responseObserver: StreamObserver[AddAr } else { throw new UnsupportedOperationException(s"Unsupported data transfer request: $req") } + } catch { + ErrorUtils.handleError( + "addArtifacts.onNext", + responseObserver, + holder.userId, + holder.sessionId, + None, + false, + Some(() => { + cleanUpStagedArtifacts() + })) } override def onError(throwable: Throwable): Unit = { @@ -128,7 +139,16 @@ class SparkConnectAddArtifactsHandler(val responseObserver: StreamObserver[AddAr responseObserver.onNext(builder.build()) responseObserver.onCompleted() } catch { - case e: StatusRuntimeException => onError(e) + ErrorUtils.handleError( + "addArtifacts.onComplete", + responseObserver, + holder.userId, + holder.sessionId, + None, + false, + Some(() => { + cleanUpStagedArtifacts() + })) } } diff --git a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala index 7cb555ca47e..703b11c0c73 100644 --- a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala +++ b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/utils/ErrorUtils.scala @@ -153,7 +153,15 @@ private[connect] object ErrorUtils extends Logging { .build() } - private def buildStatusFromThrowable( + /** + * This is a helper method that can be used by any GRPC handler to convert existing Throwables + * into GRPC conform status objects. + * + * @param st + * @param sessionHolderOpt + * @return + */ + private[connect] def buildStatusFromThrowable( st: Throwable, sessionHolderOpt: Option[SessionHolder]): RPCStatus = { val errorInfo = ErrorInfo @@ -221,6 +229,17 @@ private[connect] object ErrorUtils extends Logging { * String value indicating the operation type (analysis, execution) * @param observer * The GRPC response observer. + * @param userId + * The user id. + * @param sessionId + * The session id. + * @param events + * The ExecuteEventsManager if present to report about failures. + * @param isInterrupted + * Whether the error is caused by an interruption or during execution. + * @param callback + * Optional callback to be called after the error has been sent that allows to caller to + * execute additional cleanup logic. * @tparam V * @return */ @@ -230,7 +249,8 @@ private[connect] object ErrorUtils extends Logging { userId: String, sessionId: String, events: Option[ExecuteEventsManager] = None, - isInterrupted: Boolean = false): PartialFunction[Throwable, Unit] = { + isInterrupted: Boolean = false, + callback: Option[() => Unit] = None): PartialFunction[Throwable, Unit] = { // SessionHolder may not be present, e.g. if the session was already closed. // When SessionHolder is not present error details will not be available for FetchErrorDetails. @@ -281,6 +301,7 @@ private[connect] object ErrorUtils extends Logging { executeEventsManager.postFailed(wrapped.getMessage) } } + callback.foreach(_.apply()) observer.onError(wrapped) } } diff --git a/connector/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala b/connector/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala index 2a65032beef..e681aa4726f 100644 --- a/connector/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala +++ b/connector/connect/server/src/test/scala/org/apache/spark/sql/connect/service/AddArtifactsHandlerSuite.scala @@ -26,6 +26,10 @@ import scala.concurrent.duration._ import scala.jdk.CollectionConverters._ import com.google.protobuf.ByteString +import com.google.rpc.ErrorInfo +import io.grpc.Status.Code +import io.grpc.StatusRuntimeException +import io.grpc.protobuf.StatusProto import io.grpc.stub.StreamObserver import org.apache.spark.connect.proto @@ -371,9 +375,15 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper { val name = "/absolute/path/" val request = createDummyArtifactRequests(name) request.foreach { req => - intercept[IllegalArgumentException] { + val e = intercept[StatusRuntimeException] { handler.onNext(req) } + assert(e.getStatus.getCode == Code.INTERNAL) + val statusProto = StatusProto.fromThrowable(e) + assert(statusProto.getDetailsCount == 1) + val details = statusProto.getDetails(0) + val info = details.unpack(classOf[ErrorInfo]) + assert(info.getReason.contains("java.lang.IllegalArgumentException")) } handler.onCompleted() } finally { @@ -388,9 +398,15 @@ class AddArtifactsHandlerSuite extends SharedSparkSession with ResourceHelper { val names = Seq("..", "../sibling", "../nephew/directory", "a/../../b", "x/../y/../..") val request = names.flatMap(createDummyArtifactRequests) request.foreach { req => - intercept[IllegalArgumentException] { + val e = intercept[StatusRuntimeException] { handler.onNext(req) } + assert(e.getStatus.getCode == Code.INTERNAL) + val statusProto = StatusProto.fromThrowable(e) + assert(statusProto.getDetailsCount == 1) + val details = statusProto.getDetails(0) + val info = details.unpack(classOf[ErrorInfo]) + assert(info.getReason.contains("java.lang.IllegalArgumentException")) } handler.onCompleted() } finally { diff --git a/python/pyspark/sql/connect/client/core.py b/python/pyspark/sql/connect/client/core.py index f037e968be0..e36b7d74a78 100644 --- a/python/pyspark/sql/connect/client/core.py +++ b/python/pyspark/sql/connect/client/core.py @@ -1575,12 +1575,15 @@ class SparkConnectClient(object): raise SparkConnectGrpcException(str(rpc_error)) from None def add_artifacts(self, *paths: str, pyfile: bool, archive: bool, file: bool) -> None: - for path in paths: - for attempt in self._retrying(): - with attempt: - self._artifact_manager.add_artifacts( - path, pyfile=pyfile, archive=archive, file=file - ) + try: + for path in paths: + for attempt in self._retrying(): + with attempt: + self._artifact_manager.add_artifacts( + path, pyfile=pyfile, archive=archive, file=file + ) + except Exception as error: + self._handle_error(error) def copy_from_local_to_fs(self, local_path: str, dest_path: str) -> None: for attempt in self._retrying(): diff --git a/python/pyspark/sql/tests/connect/client/test_artifact.py b/python/pyspark/sql/tests/connect/client/test_artifact.py index 7fde0958e38..69a9525c65e 100644 --- a/python/pyspark/sql/tests/connect/client/test_artifact.py +++ b/python/pyspark/sql/tests/connect/client/test_artifact.py @@ -20,6 +20,7 @@ import tempfile import unittest import os +from pyspark.errors.exceptions.connect import SparkConnectGrpcException from pyspark.sql import SparkSession from pyspark.testing.connectutils import ReusedConnectTestCase, should_test_connect from pyspark.testing.utils import SPARK_HOME @@ -56,6 +57,25 @@ class ArtifactTestsMixin: SparkSession.builder.remote(f"sc://localhost:{ChannelBuilder.default_port()}").create() ) + def test_artifacts_cannot_be_overwritten(self): + with tempfile.TemporaryDirectory() as d: + pyfile_path = os.path.join(d, "my_pyfile.py") + with open(pyfile_path, "w+") as f: + f.write("my_func = lambda: 10") + + self.spark.addArtifacts(pyfile_path, pyfile=True) + + # Writing the same file twice is fine, and should not throw. + self.spark.addArtifacts(pyfile_path, pyfile=True) + + with open(pyfile_path, "w+") as f: + f.write("my_func = lambda: 11") + + with self.assertRaisesRegex( + SparkConnectGrpcException, "\\(java.lang.RuntimeException\\) Duplicate Artifact" + ): + self.spark.addArtifacts(pyfile_path, pyfile=True) + def check_add_zipped_package(self, spark_session): with tempfile.TemporaryDirectory() as d: package_path = os.path.join(d, "my_zipfile") --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org