sven-weber-db commented on code in PR #55683:
URL: https://github.com/apache/spark/pull/55683#discussion_r3328331895
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/externalUDF/PythonUDFWorkerSpecificationSuite.scala:
##########
@@ -26,72 +26,24 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.api.python.SimplePythonFunction
import org.apache.spark.sql.IntegratedUDFTestUtils
import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.udf.worker.UDFWorkerSpecification
-import org.apache.spark.udf.worker.core.{TestDirectWorkerDispatcher,
- UnixSocketWorkerConnection}
-
-/**
- * A test [[UnixSocketWorkerConnection]] that opens a real Unix
- * domain socket channel to the worker.
- */
-private class RealSocketConnection(
- socketPath: String,
- private val channel: java.nio.channels.SocketChannel)
- extends UnixSocketWorkerConnection(socketPath) {
-
- override def isActive: Boolean = channel.isOpen
-
- override def close(): Unit = {
- channel.close()
- super.close()
- }
-}
-
-private object RealSocketConnection {
- def connect(socketPath: String): RealSocketConnection = {
- val address =
- java.net.UnixDomainSocketAddress.of(socketPath)
- val channel = java.nio.channels.SocketChannel.open(
- java.net.StandardProtocolFamily.UNIX)
- channel.connect(address)
- new RealSocketConnection(socketPath, channel)
- }
-}
-
-/**
- * Extends [[TestDirectWorkerDispatcher]] to use a real Unix
- * domain socket connection instead of just checking file
- * existence.
- */
-private class RealSocketTestDispatcher(
- spec: UDFWorkerSpecification)
- extends TestDirectWorkerDispatcher(spec) {
-
- // Use /tmp to avoid UDS path length limits in deep
- // worktree paths.
- override protected def newEndpointAddress(
- workerId: String): String = {
- val socketDir = java.nio.file.Files.createTempDirectory(
- java.nio.file.Paths.get("/tmp"), "udf-test-")
- socketDir.toFile.deleteOnExit()
- socketDir.resolve(s"w-$workerId.sock").toString
- }
-
- override protected def createConnection(
- socketPath: String): UnixSocketWorkerConnection =
- RealSocketConnection.connect(socketPath)
-}
+import org.apache.spark.udf.worker.core.testing.TestDirectGrpcDispatcher
Review Comment:
Just to make sure this is intended: The checks in this test are slightly
weekend now. Before, we verified that a UDS can be created and will be
connectable. Now we only verify that the Python process is able to create the
socket. I think this is fine and verifies the properties we are looking for
here but I still want to point out this slight behavior change.
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/DirectGrpcDispatcher.scala:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.udf.worker.core.grpc
+
+import java.io.{File, IOException}
+import java.nio.file.{FileVisitResult, Files, Path, SimpleFileVisitor}
+import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions}
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern,
UDFWorkerSpecification,
+ WorkerConnectionSpec}
+import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle,
WorkerLogger,
+ WorkerSession}
+import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher,
DirectWorkerException,
+ DirectWorkerTimeoutException}
+import
org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS
+
+/**
+ * :: Experimental ::
+ * A concrete [[DirectWorkerDispatcher]] that spawns workers and talks to
+ * them over the UDF gRPC protocol on a Unix domain socket. Allocates a
+ * private 0700 socket directory at construction; each worker is given a
+ * UDS path inside it.
+ *
+ * gRPC over UDS is the only protocol/transport combination shipped today.
Review Comment:
nit: I don't think this comment is necessary. While this statement is true,
it will be very clear since this is the only implemented Dispatcher. If this
statement becomes false at some point, it is easy to miss it for updates.
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala:
##########
@@ -72,20 +81,56 @@ class DirectWorkerProcess(
/** Number of sessions currently using this worker. */
def activeSessions: Int = activeSessionCount.get()
- /** Increments the active session count. */
- def acquireSession(): Unit = activeSessionCount.incrementAndGet()
+ /**
+ * Increments the active session count.
+ *
+ * Preconditions: must not be called after [[close]]. The dispatcher
+ * arbitrates acquire-vs-close ordering; calling this on a closed worker
+ * indicates a dispatcher-side bug.
+ *
+ * @throws IllegalStateException if the worker has already been closed.
+ */
+ def acquireSession(): Unit = {
+ if (closed.get()) {
+ throw new IllegalStateException(
+ s"cannot acquire session: worker $id is already closed")
+ }
+ activeSessionCount.incrementAndGet()
+ }
+
+ /**
+ * Returns true if a session has marked this worker as unsafe to reuse.
+ * Pool implementations MUST treat [[isInvalid]] as a hard rejection in
+ * [[onLastSessionReleased]] -- the worker must be torn down, not
+ * recycled. Once set, the flag is sticky for the worker's lifetime.
+ */
+ def isInvalid: Boolean = invalid.get()
+
+ /**
+ * Marks this worker as unsafe to return to a reuse pool. Idempotent;
+ * sticky once set. Sessions call this when they observe a transport
+ * failure, hung worker, or any termination path that leaves the worker
+ * in an unknown state.
+ */
+ override def markInvalid(): Unit = invalid.set(true)
/**
* Decrements the active session count. Fires [[onLastSessionReleased]]
- * on the 0-transition. A negative count indicates an unbalanced
- * acquire/release; we log and reset to 0 rather than silently mask it.
+ * on the 0-transition.
+ *
+ * A negative count indicates an unbalanced acquire/release and is
+ * surfaced as a warning. We deliberately do NOT reset the counter:
+ * a concurrent `acquireSession()` between the decrement here and a
+ * `set(0)` would be silently overwritten, and the next `releaseSession`
+ * would prematurely fire `onLastSessionReleased` and tear the worker
+ * down underneath the racing caller. The counter recovers naturally on
+ * subsequent balanced calls.
*/
- def releaseSession(): Unit = {
+ override def releaseSession(): Unit = {
val c = activeSessionCount.decrementAndGet()
if (c < 0) {
logger.warn(
Review Comment:
I know this is not from this PR but why is this not a fatal failure? IMO
unbalanced acquire/release are unexpected and point to a fatal problem in our
code. A warning will easily be missed - especially if this code is only
exercised via tests where no one will read the logs.
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/DirectGrpcDispatcher.scala:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.udf.worker.core.grpc
+
+import java.io.{File, IOException}
+import java.nio.file.{FileVisitResult, Files, Path, SimpleFileVisitor}
+import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions}
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern,
UDFWorkerSpecification,
+ WorkerConnectionSpec}
+import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle,
WorkerLogger,
+ WorkerSession}
+import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher,
DirectWorkerException,
+ DirectWorkerTimeoutException}
+import
org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS
+
+/**
+ * :: Experimental ::
+ * A concrete [[DirectWorkerDispatcher]] that spawns workers and talks to
+ * them over the UDF gRPC protocol on a Unix domain socket. Allocates a
+ * private 0700 socket directory at construction; each worker is given a
+ * UDS path inside it.
+ *
+ * gRPC over UDS is the only protocol/transport combination shipped today.
+ * Lives in the `core.grpc` package alongside [[GrpcWorkerSession]] and
+ * [[GrpcWorkerChannel]] -- the two collaborators it composes.
+ */
+@Experimental
+class DirectGrpcDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger = WorkerLogger.NoOp)
+ extends DirectWorkerDispatcher(workerSpec, logger) {
+
+ // Upper bound on the rename-on-collision retry loop in newEndpointAddress.
+ // 16 is well above any realistic concurrent-worker count and keeps the
+ // failure mode bounded if something pathological occupies the directory.
+ private val MAX_SOCKET_LEAF_RETRIES = 16
+
+ // Removed explicitly in closeTransport(). deleteOnExit is avoided because
+ // the JDK retains the path for the JVM lifetime, which leaks in
+ // long-lived drivers.
+ //
+ // Initialisation order: the parent constructor runs validateTransportSupport
+ // (against the spec passed as a parameter, not subclass state) BEFORE this
+ // val is evaluated. If validation throws, this field is never assigned and
+ // no temp directory is created, which is the desired behaviour. Keep
+ // future field initialisers below this line aware of the same ordering
+ // contract -- the parent constructor cannot read them, and they cannot
+ // assume validateTransportSupport has not already failed.
+ private val socketDir: Path = createPrivateTempDirectory()
+
+ /**
+ * Returns the UDS path the worker should bind. Uses a short leaf so the
+ * full path stays within the 108-byte UDS `sun_path` limit; the worker's
+ * stable id (full UUID) is still passed via `--id` and logged separately.
+ *
+ * '''Why the truncation:''' macOS resolves `$TMPDIR` to
+ * `/var/folders/xx/yyyy/T/` (~47 chars) and `Files.createTempDirectory`
+ * appends another ~29 chars for the dispatcher's private dir. With a
+ * full UUID leaf (`worker-<36>.sock` = 49 chars) the total path is
+ * ~126 chars and `bind(2)` would silently fail with `ENAMETOOLONG`.
+ * 16 hex chars (64 bits) keeps the leaf at ~25 chars (total ~100) while
+ * making collisions negligible at any realistic concurrent-worker count.
+ *
+ * '''Defensive collision check:''' a 64-bit hash collision is
+ * astronomically unlikely (~2^-32 per pair) but if it ever fired the
+ * second worker's `bind(2)` would surface as an opaque "Worker exited
+ * with code N" via [[waitForReady]] rather than a clear filename
+ * collision. We probe the path here and append a counter suffix on
+ * the off-chance the leaf already exists so the error message is
+ * recognisable in the wild.
+ */
+ override protected def newEndpointAddress(workerId: String): String = {
+ val short = workerId.replace("-", "").take(16)
+ var candidate = socketDir.resolve(s"w-$short.sock")
+ var suffix = 0
+ while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) {
+ suffix += 1
+ candidate = socketDir.resolve(s"w-$short-$suffix.sock")
+ }
+ if (Files.exists(candidate)) {
+ throw new IllegalStateException(
+ s"could not allocate a free UDS path under $socketDir after " +
+ s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)")
+ }
+ candidate.toString
+ }
+
+ override protected def waitForReady(
+ address: String,
+ process: Process,
+ outputFile: File): Unit = {
+ val file = new File(address)
+ // At least one poll so very small initTimeouts don't trip a premature
+ // timeout before the worker has any chance to create the socket.
+ val maxAttempts = math.max(1, (initTimeoutMs /
SOCKET_POLL_INTERVAL_MS).toInt)
+ var attempts = 0
+ while (!file.exists() && attempts < maxAttempts) {
+ if (!process.isAlive) throwWorkerExitedBeforeSocket(process, address,
outputFile)
+ Thread.sleep(SOCKET_POLL_INTERVAL_MS)
+ attempts += 1
+ }
+ if (!file.exists()) {
+ if (process.isAlive) {
+ DirectWorkerDispatcher.destroyForciblyAndReap(
+ process, logger, s"init timeout $address")
+ val tail = readOutputTail(outputFile)
+ throw new DirectWorkerTimeoutException(
+ s"Worker did not create socket at $address within
${initTimeoutMs}ms\n$tail")
+ } else {
+ // Worker exited after the last poll without creating the socket;
+ // prefer the exit-code message over the ambiguous "did not create".
+ throwWorkerExitedBeforeSocket(process, address, outputFile)
+ }
+ }
+ }
+
+ override protected def cleanupEndpointAddress(address: String): Unit = {
+ Files.deleteIfExists(new File(address).toPath)
+ }
+
+ override protected def closeTransport(): Unit = {
+ if (!Files.exists(socketDir)) return
+ // Recursive post-order delete: today socketDir contains only socket files
+ // at the top level, but a future change that namespaces workers into
+ // subdirectories should not silently leak them.
+ Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] {
+ override def visitFile(file: Path, attrs: BasicFileAttributes):
FileVisitResult = {
+ try Files.deleteIfExists(file) catch { case _: IOException => () }
+ FileVisitResult.CONTINUE
+ }
+ override def postVisitDirectory(dir: Path, exc: IOException):
FileVisitResult = {
+ try Files.deleteIfExists(dir) catch { case _: IOException => () }
+ FileVisitResult.CONTINUE
+ }
+ })
+ }
+
+ // `spec` is the same object as the `workerSpec` field but passed
+ // explicitly: at the point this runs (parent constructor body), `this`
+ // is only partially constructed and reading subclass fields is unsafe.
+ // See the contract on the abstract method in [[DirectWorkerDispatcher]].
+ override protected def validateTransportSupport(spec:
UDFWorkerSpecification): Unit = {
+ val props = spec.getDirect.getProperties
+ require(props.hasConnection,
+ "DirectWorker.properties.connection must be set")
+ val conn = props.getConnection
+ require(conn.getTransportCase ==
WorkerConnectionSpec.TransportCase.UNIX_DOMAIN_SOCKET,
+ "DirectGrpcDispatcher requires UNIX domain socket transport, " +
+ s"got ${conn.getTransportCase}")
+ // BIDI is the only pattern the gRPC `Execute` RPC speaks. If the spec
+ // declares capabilities, it MUST include BIDI; an unset capabilities
+ // block is treated as "no constraint" and accepted.
+ if (spec.hasCapabilities) {
Review Comment:
Why do we silently accept an empty capabilities block here?
This is part of the spec already so any future instances of the worker
should set and support it
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerDispatcher.scala:
##########
@@ -117,6 +127,10 @@ abstract class DirectWorkerDispatcher(
private[this] val workers = new ConcurrentHashMap[String,
DirectWorkerProcess]()
private[this] val closed = new AtomicBoolean(false)
+ // TODO: extract the env state machine + JVM shutdown hook into a
Review Comment:
nit: Can we link the JIRA parent ticket here? That way we have a easy way of
finding open TODOs once we productionize this
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/DirectGrpcDispatcher.scala:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.udf.worker.core.grpc
+
+import java.io.{File, IOException}
+import java.nio.file.{FileVisitResult, Files, Path, SimpleFileVisitor}
+import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions}
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern,
UDFWorkerSpecification,
+ WorkerConnectionSpec}
+import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle,
WorkerLogger,
+ WorkerSession}
+import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher,
DirectWorkerException,
+ DirectWorkerTimeoutException}
+import
org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS
+
+/**
+ * :: Experimental ::
+ * A concrete [[DirectWorkerDispatcher]] that spawns workers and talks to
+ * them over the UDF gRPC protocol on a Unix domain socket. Allocates a
+ * private 0700 socket directory at construction; each worker is given a
+ * UDS path inside it.
+ *
+ * gRPC over UDS is the only protocol/transport combination shipped today.
+ * Lives in the `core.grpc` package alongside [[GrpcWorkerSession]] and
+ * [[GrpcWorkerChannel]] -- the two collaborators it composes.
+ */
+@Experimental
+class DirectGrpcDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger = WorkerLogger.NoOp)
+ extends DirectWorkerDispatcher(workerSpec, logger) {
+
+ // Upper bound on the rename-on-collision retry loop in newEndpointAddress.
+ // 16 is well above any realistic concurrent-worker count and keeps the
+ // failure mode bounded if something pathological occupies the directory.
+ private val MAX_SOCKET_LEAF_RETRIES = 16
+
+ // Removed explicitly in closeTransport(). deleteOnExit is avoided because
+ // the JDK retains the path for the JVM lifetime, which leaks in
+ // long-lived drivers.
+ //
+ // Initialisation order: the parent constructor runs validateTransportSupport
+ // (against the spec passed as a parameter, not subclass state) BEFORE this
+ // val is evaluated. If validation throws, this field is never assigned and
+ // no temp directory is created, which is the desired behaviour. Keep
+ // future field initialisers below this line aware of the same ordering
+ // contract -- the parent constructor cannot read them, and they cannot
+ // assume validateTransportSupport has not already failed.
+ private val socketDir: Path = createPrivateTempDirectory()
+
+ /**
+ * Returns the UDS path the worker should bind. Uses a short leaf so the
+ * full path stays within the 108-byte UDS `sun_path` limit; the worker's
+ * stable id (full UUID) is still passed via `--id` and logged separately.
+ *
+ * '''Why the truncation:''' macOS resolves `$TMPDIR` to
+ * `/var/folders/xx/yyyy/T/` (~47 chars) and `Files.createTempDirectory`
+ * appends another ~29 chars for the dispatcher's private dir. With a
+ * full UUID leaf (`worker-<36>.sock` = 49 chars) the total path is
+ * ~126 chars and `bind(2)` would silently fail with `ENAMETOOLONG`.
+ * 16 hex chars (64 bits) keeps the leaf at ~25 chars (total ~100) while
+ * making collisions negligible at any realistic concurrent-worker count.
+ *
+ * '''Defensive collision check:''' a 64-bit hash collision is
+ * astronomically unlikely (~2^-32 per pair) but if it ever fired the
+ * second worker's `bind(2)` would surface as an opaque "Worker exited
+ * with code N" via [[waitForReady]] rather than a clear filename
+ * collision. We probe the path here and append a counter suffix on
+ * the off-chance the leaf already exists so the error message is
+ * recognisable in the wild.
+ */
+ override protected def newEndpointAddress(workerId: String): String = {
+ val short = workerId.replace("-", "").take(16)
+ var candidate = socketDir.resolve(s"w-$short.sock")
+ var suffix = 0
+ while (Files.exists(candidate) && suffix < MAX_SOCKET_LEAF_RETRIES) {
+ suffix += 1
+ candidate = socketDir.resolve(s"w-$short-$suffix.sock")
+ }
+ if (Files.exists(candidate)) {
+ throw new IllegalStateException(
+ s"could not allocate a free UDS path under $socketDir after " +
+ s"$MAX_SOCKET_LEAF_RETRIES retries (truncated id=$short)")
+ }
+ candidate.toString
+ }
+
+ override protected def waitForReady(
+ address: String,
+ process: Process,
+ outputFile: File): Unit = {
+ val file = new File(address)
+ // At least one poll so very small initTimeouts don't trip a premature
+ // timeout before the worker has any chance to create the socket.
+ val maxAttempts = math.max(1, (initTimeoutMs /
SOCKET_POLL_INTERVAL_MS).toInt)
+ var attempts = 0
+ while (!file.exists() && attempts < maxAttempts) {
+ if (!process.isAlive) throwWorkerExitedBeforeSocket(process, address,
outputFile)
+ Thread.sleep(SOCKET_POLL_INTERVAL_MS)
+ attempts += 1
+ }
+ if (!file.exists()) {
+ if (process.isAlive) {
+ DirectWorkerDispatcher.destroyForciblyAndReap(
+ process, logger, s"init timeout $address")
+ val tail = readOutputTail(outputFile)
+ throw new DirectWorkerTimeoutException(
+ s"Worker did not create socket at $address within
${initTimeoutMs}ms\n$tail")
+ } else {
+ // Worker exited after the last poll without creating the socket;
+ // prefer the exit-code message over the ambiguous "did not create".
+ throwWorkerExitedBeforeSocket(process, address, outputFile)
+ }
+ }
+ }
+
+ override protected def cleanupEndpointAddress(address: String): Unit = {
+ Files.deleteIfExists(new File(address).toPath)
+ }
+
+ override protected def closeTransport(): Unit = {
+ if (!Files.exists(socketDir)) return
+ // Recursive post-order delete: today socketDir contains only socket files
+ // at the top level, but a future change that namespaces workers into
+ // subdirectories should not silently leak them.
+ Files.walkFileTree(socketDir, new SimpleFileVisitor[Path] {
+ override def visitFile(file: Path, attrs: BasicFileAttributes):
FileVisitResult = {
+ try Files.deleteIfExists(file) catch { case _: IOException => () }
+ FileVisitResult.CONTINUE
+ }
+ override def postVisitDirectory(dir: Path, exc: IOException):
FileVisitResult = {
+ try Files.deleteIfExists(dir) catch { case _: IOException => () }
+ FileVisitResult.CONTINUE
+ }
+ })
+ }
+
+ // `spec` is the same object as the `workerSpec` field but passed
+ // explicitly: at the point this runs (parent constructor body), `this`
+ // is only partially constructed and reading subclass fields is unsafe.
+ // See the contract on the abstract method in [[DirectWorkerDispatcher]].
+ override protected def validateTransportSupport(spec:
UDFWorkerSpecification): Unit = {
+ val props = spec.getDirect.getProperties
+ require(props.hasConnection,
+ "DirectWorker.properties.connection must be set")
+ val conn = props.getConnection
+ require(conn.getTransportCase ==
WorkerConnectionSpec.TransportCase.UNIX_DOMAIN_SOCKET,
+ "DirectGrpcDispatcher requires UNIX domain socket transport, " +
+ s"got ${conn.getTransportCase}")
+ // BIDI is the only pattern the gRPC `Execute` RPC speaks. If the spec
Review Comment:
nit: unnecessary abbreviation without definition
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/DirectGrpcDispatcher.scala:
##########
@@ -0,0 +1,237 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.udf.worker.core.grpc
+
+import java.io.{File, IOException}
+import java.nio.file.{FileVisitResult, Files, Path, SimpleFileVisitor}
+import java.nio.file.attribute.{BasicFileAttributes, PosixFilePermissions}
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{UDFProtoCommunicationPattern,
UDFWorkerSpecification,
+ WorkerConnectionSpec}
+import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerHandle,
WorkerLogger,
+ WorkerSession}
+import org.apache.spark.udf.worker.core.direct.{DirectWorkerDispatcher,
DirectWorkerException,
+ DirectWorkerTimeoutException}
+import
org.apache.spark.udf.worker.core.direct.DirectWorkerDispatcher.SOCKET_POLL_INTERVAL_MS
+
+/**
+ * :: Experimental ::
+ * A concrete [[DirectWorkerDispatcher]] that spawns workers and talks to
+ * them over the UDF gRPC protocol on a Unix domain socket. Allocates a
+ * private 0700 socket directory at construction; each worker is given a
+ * UDS path inside it.
+ *
+ * gRPC over UDS is the only protocol/transport combination shipped today.
+ * Lives in the `core.grpc` package alongside [[GrpcWorkerSession]] and
+ * [[GrpcWorkerChannel]] -- the two collaborators it composes.
+ */
+@Experimental
+class DirectGrpcDispatcher(
+ workerSpec: UDFWorkerSpecification,
+ logger: WorkerLogger = WorkerLogger.NoOp)
+ extends DirectWorkerDispatcher(workerSpec, logger) {
+
+ // Upper bound on the rename-on-collision retry loop in newEndpointAddress.
+ // 16 is well above any realistic concurrent-worker count and keeps the
+ // failure mode bounded if something pathological occupies the directory.
+ private val MAX_SOCKET_LEAF_RETRIES = 16
+
+ // Removed explicitly in closeTransport(). deleteOnExit is avoided because
+ // the JDK retains the path for the JVM lifetime, which leaks in
+ // long-lived drivers.
+ //
+ // Initialisation order: the parent constructor runs validateTransportSupport
+ // (against the spec passed as a parameter, not subclass state) BEFORE this
+ // val is evaluated. If validation throws, this field is never assigned and
Review Comment:
Relying on the implicit initialization order of the (base) class here is not
that nice. Could we make this order implicit instead? E.g. have a function that
will be called for setup after the base class has finished validation? This
will make the semantics clearer, and this code much easier to reason about.
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/WorkerSession.scala:
##########
@@ -58,25 +58,29 @@ import org.apache.spark.udf.worker.Init
* See [[cancel]] for the full contract.
*
* The lifecycle is enforced here: [[init]] and [[process]] are `final`
- * and delegate to [[doInit]] / [[doProcess]] after AtomicBoolean guards.
- * Subclasses implement the protocol-specific work and do not re-check
- * the contract.
+ * and route through [[doInit]] / [[doProcess]] after AtomicBoolean
+ * guards. Subclasses provide the protocol-specific work and do not
+ * re-check the contract.
+ *
+ * @param workerHandle dispatcher-side handle used to release the worker
+ * on [[close]]. Package-private to `core`: the handle
+ * is dispatcher implementation detail.
+ * @param logger logger for lifecycle diagnostics.
*/
@Experimental
-abstract class WorkerSession extends AutoCloseable {
+abstract class WorkerSession(
+ private[core] val workerHandle: WorkerHandle,
+ protected val logger: WorkerLogger) extends AutoCloseable {
private val initialized = new AtomicBoolean(false)
private val processed = new AtomicBoolean(false)
+ private val released = new AtomicBoolean(false)
/**
* Initializes the UDF execution. Must be called exactly once before
* [[process]].
*
* Throws `IllegalStateException` if called more than once.
- *
Review Comment:
Why are we removing the descriptions for the method parameters in this class?
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/GrpcWorkerChannel.scala:
##########
@@ -0,0 +1,142 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.udf.worker.core.grpc
+
+import java.io.File
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.util.control.NonFatal
+
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.netty.NettyChannelBuilder
+import io.netty.channel.EventLoopGroup
+import io.netty.channel.unix.DomainSocketAddress
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.core.{WorkerConnection, WorkerLogger}
+
+/**
+ * :: Experimental ::
+ * A [[WorkerConnection]] backed by a gRPC `ManagedChannel` over a Netty
+ * Unix domain socket transport.
+ *
+ * Owns:
+ * - the [[ManagedChannel]] used by sessions to open Execute streams,
+ * - a dedicated Netty [[EventLoopGroup]] for that channel (one per worker so
+ * a slow worker cannot starve others sharing a global loop),
+ * - the socket file at `socketPath` (removed on [[close]]).
+ *
+ * Lifecycle:
+ * - [[isActive]] is true while the channel reports any non-`SHUTDOWN` state.
+ * Transient failures keep the connection usable -- gRPC will reconnect on
+ * the next RPC.
+ * - [[close]] shuts the channel down, drains the event loop, then removes
+ * the socket file. Idempotent.
+ *
+ * Construction does NOT block on the worker actually being reachable; the
+ * dispatcher proves liveness separately via its
+ * `waitForReady`/init-timeout machinery before handing the connection to a
+ * session.
+ *
+ * @param socketPath path of the UDS the worker server binds.
+ * @param logger logger for shutdown diagnostics; defaults to NoOp.
+ * @param shutdownPhaseTimeoutMs upper bound on each individual phase of
+ * close() -- graceful channel drain, forced
+ * channel drain, and event-loop drain are
+ * three separate phases that each get this
+ * budget. Total worst-case close time is up
+ * to three times this value. Defaults to 5s
+ * per phase.
+ */
+@Experimental
+class GrpcWorkerChannel(
+ val socketPath: String,
+ logger: WorkerLogger = WorkerLogger.NoOp,
+ shutdownPhaseTimeoutMs: Long =
GrpcWorkerChannel.DEFAULT_SHUTDOWN_PHASE_TIMEOUT_MS)
+ extends WorkerConnection {
+
+ private val transport: UnixDomainSocketTransport =
UnixDomainSocketTransport.detect()
+ private val eventLoopGroup: EventLoopGroup = transport.newEventLoopGroup()
+
+ /**
+ * Underlying gRPC channel. Sessions open one bidirectional Execute stream
+ * per session on this channel.
+ */
+ val channel: ManagedChannel = try {
+ NettyChannelBuilder
+ .forAddress(new DomainSocketAddress(socketPath))
+ .channelType(transport.channelType)
+ .eventLoopGroup(eventLoopGroup)
+ .usePlaintext()
+ .build()
+ } catch {
+ case NonFatal(e) =>
+ UnixDomainSocketTransport.shutdown(eventLoopGroup,
shutdownPhaseTimeoutMs)
+ throw e
+ }
+
+ private val closed = new AtomicBoolean(false)
+
+ override def isActive: Boolean = {
+ if (closed.get()) return false
+ channel.getState(false) != ConnectivityState.SHUTDOWN
+ }
+
+ /**
+ * Idempotently tears down the channel, event loop, and removes the socket
+ * file. Steps are guarded so a failure in one does not skip the others.
+ */
+ override def close(): Unit = {
+ if (!closed.compareAndSet(false, true)) return
+
+ try {
+ channel.shutdown()
+ if (!channel.awaitTermination(shutdownPhaseTimeoutMs,
TimeUnit.MILLISECONDS)) {
+ channel.shutdownNow()
+ channel.awaitTermination(shutdownPhaseTimeoutMs, TimeUnit.MILLISECONDS)
+ }
+ } catch {
+ case _: InterruptedException =>
+ Thread.currentThread().interrupt()
+ case NonFatal(e) =>
+ logger.warn(s"Error shutting down gRPC channel for $socketPath", e)
+ }
+
+ try {
+ UnixDomainSocketTransport.shutdown(eventLoopGroup,
shutdownPhaseTimeoutMs)
+ } catch {
+ case NonFatal(e) =>
+ logger.warn(s"Error shutting down event loop for $socketPath", e)
+ }
+
+ // Remove the UDS file last so the engine event loop has fully drained
+ // before the path disappears. delete() is best-effort: missing-file is
+ // not an error (already-removed sockets are common in test crashes).
+ val socketFile = new File(socketPath)
Review Comment:
The ownership of the socket seems a bit unclear. This socket is not created
in this channel and already cleaned-up in the Dispatcher. Why are we also
trying to remove it here? A single place of creation and deletion seems more
logical to me
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/grpc/GrpcWorkerSession.scala:
##########
@@ -0,0 +1,773 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.udf.worker.core.grpc
+
+import java.util.concurrent.{CountDownLatch, LinkedBlockingQueue,
TimeoutException, TimeUnit}
+import java.util.concurrent.atomic.{AtomicBoolean, AtomicReference}
+
+import scala.util.control.NonFatal
+
+import com.google.protobuf.ByteString
+import io.grpc.{ConnectivityState, ManagedChannel}
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.udf.worker.{Cancel, DataRequest, ExecutionError,
Finish, Init,
+ UdfControlRequest, UdfControlResponse, UdfRequest, UdfResponse,
UdfWorkerGrpc}
+import org.apache.spark.udf.worker.core.{WorkerHandle, WorkerLogger,
WorkerSession}
+import org.apache.spark.udf.worker.core.grpc.GrpcWorkerSession._
+
+/**
+ * :: Experimental ::
+ * gRPC implementation of [[WorkerSession]] for the `UdfWorker.Execute`
+ * bidirectional RPC.
+ *
+ * Drives one bidirectional `Execute` stream against the worker per the
+ * ordering invariants documented in `udf_protocol.proto`:
+ * {{{
+ * Engine -> Worker: Init -> (DataRequest)* -> Finish (Cancel)?
+ * | Cancel
+ * Worker -> Engine: InitResponse -> (DataResponse)* -> (ErrorResponse)? ->
(FinishResponse | CancelResponse)
+ * }}}
+ *
+ * Knows nothing about how the worker was provisioned (locally spawned,
+ * indirectly looked up, ...) -- the dispatcher constructs this with a
+ * [[WorkerHandle]] and channel; the base [[WorkerSession]] handles
+ * dispatcher-side cleanup on close.
+ *
+ * Threading:
+ * - [[doInit]] is synchronous: sends `Init` and blocks on `InitResponse`.
+ * - [[doProcess]] returns an iterator. Input batches are forwarded inline
+ * (the iterator's `next()` thread also sends `DataRequest`). Output
+ * batches arrive via the response observer (gRPC callback thread) and
+ * are consumed by the same iterator. A terminator (`FinishResponse`,
+ * `CancelResponse`, `ErrorResponse`, gRPC stream error) is published
+ * once.
+ * - [[cancel]] is thread-safe and idempotent.
+ * - [[doClose]] half-closes the request stream after the terminator arrives.
+ *
+ * NOTE: this class does not yet implement payload chunking; the entire
Review Comment:
nit: Add a TODO with the JIRA tracking ticket for this
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
