Yicong-Huang commented on code in PR #7207: URL: https://github.com/apache/texera/pull/7207#discussion_r3742251585
########## common/workflow-operator/src/test/scala/org/apache/texera/amber/util/python/PythonWorkerPoolSpec.scala: ########## @@ -0,0 +1,172 @@ +/* + * 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.texera.amber.util.python + +import org.apache.texera.amber.util.JSONUtils.objectMapper +import org.scalatest.funsuite.AnyFunSuite + +import java.util.concurrent.{Executors, TimeUnit} +import scala.concurrent.duration._ +import scala.concurrent.{Await, ExecutionContext, ExecutionContextExecutorService, Future} +import scala.util.Try + +/** + * What the pool owes a caller when a worker misbehaves. An ordinary job failure + * is the other suites' business; this one is about a worker that stays alive and + * stops taking part, which is the case that does not end by itself: a crash + * closes the pipe and the pending read returns, while silence would hold the + * caller forever — neither a read nor a write on a process pipe answers an + * interrupt or a deadline, so no suite-level timeout can release one. + * + * Every wait here is bounded and runs on daemon threads, so a regression fails + * these tests instead of wedging the run: a lost non-daemon thread parked on a + * pipe would keep the JVM, and the build, alive. + * + * The fixture worker is stdlib-only and runs under `-I -S`, so this needs an + * interpreter but none of the operator packages. + */ +final class PythonWorkerPoolSpec extends AnyFunSuite { + + private val HangingWorker = "/python/hanging_worker.py" + private val CompileWorker = "/python/py_compile_worker.py" + + /** Short enough to keep the suite quick, far enough above process startup to + * not be mistaken for one: the fixture never answers, so it cannot race. + */ + private val ShortTimeouts: PythonWorkerPool.Timeouts = + PythonWorkerPool.Timeouts(responseMillis = 1500, startupMillis = 1500) + + /** Ceiling on a whole case, well above the deadlines under test. Reaching it + * means something never gave up. + */ + private val Bound: FiniteDuration = 25.seconds + + /** Any interpreter serves — the fixture imports only `json` and `time` — so + * this deliberately skips the configured `python.path` the suites that need + * pandas resolve. A machine without one cancels rather than fails. + */ + private def python(): String = { + def isRunnable(exe: String): Boolean = + Try(new ProcessBuilder(exe, "--version").redirectErrorStream(true).start()).toOption + .exists { p => + if (p.waitFor(5, TimeUnit.SECONDS)) p.exitValue() == 0 else { p.destroyForcibly(); false } + } + + List("python3", "python", "py").find(isRunnable).getOrElse(cancel("no runnable python")) + } + + private def onDaemonThreads[T](threads: Int)(body: ExecutionContext => T): T = { + val pool = Executors.newFixedThreadPool( + threads, + (r: Runnable) => { + val t = new Thread(r, "pool-spec-caller") + t.setDaemon(true) + t + } + ) + val ec: ExecutionContextExecutorService = ExecutionContext.fromExecutorService(pool) + try body(ec) + finally pool.shutdownNow() + } + + private def hangingCall( + launchArgs: Seq[String], + request: com.fasterxml.jackson.databind.node.ObjectNode = objectMapper.createObjectNode() + ): PythonWorkerPool.Outcome = + PythonWorkerPool.run( + resourcePath = HangingWorker, + launchArgs = launchArgs, + pythonExe = python(), + request = request, + interpreterArgs = Seq("-I", "-S"), + timeouts = ShortTimeouts + ) + + /** The call, on a daemon thread and under [[Bound]], expected to give up. */ + private def interceptBounded(call: => Any): PythonWorkerPool.WorkerDiedException = + intercept[PythonWorkerPool.WorkerDiedException] { + onDaemonThreads(1)(ec => Await.result(Future(call)(ec), Bound)) + } + + test("a worker that takes the job and stops answering is killed and reported") { + val startedAt = System.nanoTime() + val thrown = interceptBounded(hangingCall(Seq.empty)) + val elapsedMillis = (System.nanoTime() - startedAt) / 1000000 + + assert(thrown.getMessage.contains("did not answer")) + assert(thrown.getMessage.contains("killed it")) + // Well under the default response budget: what fired is the timeout passed in, + // not a wait that happened to end. + assert(elapsedMillis < PythonWorkerPool.Timeouts.Default.responseMillis / 2) + } + + test("a worker that never signals ready is killed and reported") { + val startedAt = System.nanoTime() + val thrown = interceptBounded(hangingCall(Seq("--hang-before-ready"))) + val elapsedMillis = (System.nanoTime() - startedAt) / 1000000 + + assert(thrown.getMessage.contains("did not signal ready")) + assert(elapsedMillis < PythonWorkerPool.Timeouts.Default.startupMillis / 2) + } + + test("a worker that never reads its request is killed and reported") { + val request = objectMapper.createObjectNode() + // Past any pipe buffer, so the write cannot simply be handed to the kernel and + // left there: it is the blocked write itself that has to be given up on. + request.put("source", "x" * (4 * 1024 * 1024)) + + val thrown = interceptBounded(hangingCall(Seq("--deaf"), request)) + + assert(thrown.getMessage.contains("did not read its request")) + assert(thrown.getMessage.contains("killed it")) + } + + test("a caller waiting at the worker cap is not stranded by a discarded worker") { + // One caller more than there are workers, all onto a worker that goes quiet: Review Comment: `borrow` creates a worker per caller up to the cap, so this starts `maxWorkers` hanging workers, not one. The singular hides the mechanism the case exists to test: the discards that free slots happen on different workers from the one the capped caller is waiting behind. ```suggestion // One caller more than there are workers, all onto workers that go quiet: ``` ########## common/workflow-operator/src/test/scala/org/apache/texera/amber/util/python/PythonWorkerPool.scala: ########## @@ -0,0 +1,445 @@ +/* + * 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.texera.amber.util.python + +import com.fasterxml.jackson.databind.node.ObjectNode +import com.typesafe.scalalogging.LazyLogging +import org.apache.texera.amber.util.JSONUtils.objectMapper + +import java.io.{BufferedReader, BufferedWriter, InputStreamReader, OutputStreamWriter} +import java.nio.charset.StandardCharsets +import java.nio.file.{Files, Path, StandardCopyOption} +import java.util.concurrent.{ + Callable, + ConcurrentHashMap, + ExecutionException, + ExecutorService, + Executors, + LinkedBlockingQueue, + TimeUnit, + TimeoutException +} +import java.util.concurrent.atomic.AtomicInteger +import scala.annotation.tailrec +import scala.collection.mutable +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +/** + * Pools of persistent Python "worker" processes that eliminate the per-call + * interpreter-boot + import cost a test otherwise pays on every subprocess + * spawn. Testing operators one at a time does not scale when each one costs a + * spawn: a bare `-I -S` interpreter boots in ~25 ms, and once pandas and plotly + * are imported a spawn costs ~260-310 ms — ~96% of a job whose real work is + * ~4 ms. A worker pays that once at startup, then serves many jobs over its + * lifetime, so N spawns become one. + * + * Lives in test scope here, rather than beside a single caller, because tests in + * several modules run generated operator code and would otherwise each + * hand-roll a driver, a stdout protocol and a timeout. Other modules reach it + * through a `test->test` dependency on this one. + * + * Generic over the worker script — the pool never interprets the payload — so + * one implementation serves a syntax check, template execution and DataFrame + * comparison alike. Each distinct (resource, interpreterArgs, launchArgs, + * python, env) combination gets its own sub-pool. + * + * Protocol (line-delimited JSON, shared by all worker scripts): + * startup worker -> pool: {"ready": true} + * request pool -> worker: <caller-supplied JSON object>\n + * response worker -> pool: {"exit": <int>, "stdout": "...", "stderr": "..."}\n + * + * Concurrency: callers submit from several threads at once — a spec run with + * ScalaTest's `-P4`, or one test fanning its cases out — so each sub-pool holds + * up to [[maxWorkers]] workers, each serving one job at a time (borrow -> run + * -> return). A worker script may chdir per job, so a worker must never run two + * jobs at once — the borrow/return discipline guarantees that. + * + * Robustness: an ordinary *job* failure comes back as an [[Outcome]] with + * `exit != 0` (worker keeps running). A hard interpreter crash ends a worker; + * the pool detects the EOF / broken pipe, discards it, and throws + * [[WorkerDiedException]] so the caller can fall back to a one-shot subprocess + * — behavior is then never worse than the pre-pool path. A worker that stays + * alive but stops answering ends the same way, on the [[Timeouts]] below. + */ +object PythonWorkerPool extends LazyLogging { + + /** Worker response: process-like exit code plus captured streams. */ + final case class Outcome(exit: Int, stdout: String, stderr: String) + + /** Thrown when a worker dies mid-job (hard crash / broken pipe). Callers + * catch this and fall back to a one-shot subprocess. + */ + final class WorkerDiedException(message: String, cause: Throwable = null) + extends RuntimeException(message, cause) + + /** Feature toggle. `TEXERA_TEST_PYTHON_WORKER=0` (or `false`/`off`) forces the + * one-subprocess-per-call paths everywhere — an escape hatch for debugging a + * suspected isolation leak. Default on. + */ + val enabled: Boolean = + !sys.env + .get("TEXERA_TEST_PYTHON_WORKER") + .map(_.trim.toLowerCase) + .exists(Set("0", "false", "off")) + + /** Max live workers per sub-pool. Defaults to 4 to match ScalaTest's `-P4`, so + * the two concurrency bounds agree on how many interpreters may be live; + * override via `TEXERA_TEST_PYTHON_WORKERS`. Public so a caller fanning out + * jobs within one test can size that fan-out to the workers it will get. + */ + val maxWorkers: Int = + sys.env + .get("TEXERA_TEST_PYTHON_WORKERS") + .flatMap(s => scala.util.Try(s.trim.toInt).toOption) + .filter(_ > 0) + .getOrElse(4) + + /** How long a caller waits on a worker before the pool kills and discards it. + * A read on a process pipe cannot be interrupted — a suite or executor timeout + * leaves the reading thread stuck on it — so a worker that stays alive without + * answering has to be bounded here. `response` keeps the 30 seconds the + * one-shot spawn this pool replaced allowed a job; `startup` is longer because + * a worker imports its libraries before it reports ready, and a loaded CI + * machine makes that slow. Override in seconds via + * `TEXERA_TEST_PYTHON_WORKER_TIMEOUT` / `TEXERA_TEST_PYTHON_WORKER_STARTUP_TIMEOUT`. + */ + final case class Timeouts(responseMillis: Long, startupMillis: Long) + + object Timeouts { + private def envSeconds(name: String, default: Long): Long = + sys.env + .get(name) + .flatMap(s => scala.util.Try(s.trim.toLong).toOption) + .filter(_ > 0) + .getOrElse(default) * 1000 + + val Default: Timeouts = Timeouts( + responseMillis = envSeconds("TEXERA_TEST_PYTHON_WORKER_TIMEOUT", 30), + startupMillis = envSeconds("TEXERA_TEST_PYTHON_WORKER_STARTUP_TIMEOUT", 60) + ) + } + + /** + * Run one job through a pooled worker for `resourcePath`, launched as + * `pythonExe <interpreterArgs> <script> <launchArgs>` with extra environment + * `env`. `request` is the worker-specific JSON payload (the pool does not + * interpret it). Throws [[WorkerDiedException]] on a hard worker crash. + * + * `interpreterArgs` are the flags that must precede the script — a syntax + * checker wants `-I -S` so it validates under the same isolation a one-shot + * `python -I -S -m py_compile` gave it. `launchArgs` are the script's own + * (e.g. `--serve`), and `env` carries what a flag cannot (e.g. PYTHONPATH). + * All three are part of a worker's identity: one started differently is not + * interchangeable, so it gets its own sub-pool. `timeouts` is not — it bounds + * this call, so a caller whose jobs are slower than most can raise it without + * splitting the pool. + */ + def run( + resourcePath: String, + launchArgs: Seq[String], + pythonExe: String, + request: ObjectNode, + env: Map[String, String] = Map.empty, + interpreterArgs: Seq[String] = Seq.empty, + timeouts: Timeouts = Timeouts.Default + ): Outcome = { + val pool = pools.computeIfAbsent( + Key(resourcePath, pythonExe, interpreterArgs.toList, launchArgs.toList, env.toList.sorted), + _ => new Pool(resourcePath, launchArgs, pythonExe, env, interpreterArgs) + ) + pool.run(request, timeouts) + } + + /** What makes two launches the same worker. Compared field by field rather than + * as one joined string, so a value carrying whatever the separator was — a + * python path with a space in it, a PYTHONPATH — cannot make two different + * launches share a pool. + */ + private final case class Key( + resourcePath: String, + pythonExe: String, + interpreterArgs: List[String], + launchArgs: List[String], + env: List[(String, String)] + ) + + /** How long a caller at the worker cap waits before re-examining it. Not a + * deadline — see [[Pool.borrow]]. + */ + private val CapRecheckMillis: Long = 250 + + private val pools = new ConcurrentHashMap[Key, Pool]() + + Runtime.getRuntime.addShutdownHook(new Thread(() => shutdownAll())) + + private def shutdownAll(): Unit = + pools.values().forEach(_.shutdown()) + + // A single sub-pool: up to `maxWorkers` live workers for one worker script. + private final class Pool( + resourcePath: String, + launchArgs: Seq[String], + pythonExe: String, + env: Map[String, String], + interpreterArgs: Seq[String] + ) { + private val idle = new LinkedBlockingQueue[Worker]() + private val liveCount = new AtomicInteger(0) + private val all = mutable.Set.empty[Worker] // guarded by `all` + @volatile private var script: Path = _ + + def run(request: ObjectNode, timeouts: Timeouts): Outcome = { + val w = borrow(timeouts) + try { + val outcome = w.run(request, timeouts.responseMillis) + idle.offer(w) // healthy — return to pool + outcome + } catch { + case e: WorkerDiedException => + discard(w) + throw e + } + } + + @tailrec + private def borrow(timeouts: Timeouts): Worker = { + val existing = idle.poll() + if (existing != null) existing + else if (liveCount.getAndIncrement() < maxWorkers) { + try create(timeouts) + catch { + case e: Throwable => + liveCount.decrementAndGet() + throw e + } + } else { + liveCount.decrementAndGet() + // At the cap. Waiting outright for a returned worker would strand this + // caller when the ones ahead are discarded instead: a discard frees a + // slot without putting anything back. So wait only briefly, then look at + // the cap again — the next pass starts a replacement. A long queue still + // waits as long as it takes; that is the caller's own backlog, not a hang. + val returned = idle.poll(CapRecheckMillis, TimeUnit.MILLISECONDS) + if (returned != null) returned else borrow(timeouts) + } + } + + private def create(timeouts: Timeouts): Worker = { + val cmd = + (((pythonExe +: interpreterArgs) :+ ensureScript().toString) ++ launchArgs).asJava + val pb = new ProcessBuilder(cmd).redirectErrorStream(false) + env.foreach { case (k, v) => pb.environment().put(k, v) } + val w = new Worker(pb.start(), s"$resourcePath ${launchArgs.mkString(" ")}".trim) + w.awaitReady(timeouts.startupMillis) + all.synchronized(all.add(w)) + logger.debug(s"Started python worker for $resourcePath (live=${liveCount.get}/$maxWorkers)") + w + } + + private def discard(w: Worker): Unit = { + all.synchronized(all.remove(w)) + liveCount.decrementAndGet() + w.destroy() + } + + private def ensureScript(): Path = { + if (script == null) synchronized { + if (script == null) { + val stream = getClass.getResourceAsStream(resourcePath) + require(stream != null, s"worker script not found on classpath at $resourcePath") + try { + val tmp = Files.createTempFile("py-worker-", ".py") + Files.copy(stream, tmp, StandardCopyOption.REPLACE_EXISTING) + tmp.toFile.deleteOnExit() + script = tmp + } finally stream.close() + } + } + script + } + + def shutdown(): Unit = + all.synchronized { + all.foreach(_.destroy()) + all.clear() + } + } + + // One live worker process plus its framed-JSON stdin and background drains of + // its stdout (the protocol) and stderr (only non-empty on a hard crash). + private final class Worker(process: Process, label: String) { + private val stdin: BufferedWriter = + new BufferedWriter(new OutputStreamWriter(process.getOutputStream, StandardCharsets.UTF_8)) + private val errBuf = new StringBuilder + + // Protocol lines the worker has written, `None` marking end of stream. A + // dedicated thread owns the blocking read so a caller can wait with a + // timeout: `readLine` on a process pipe answers neither an interrupt nor a + // deadline, and only closing the pipe — killing the process — releases it. + private val lines = new LinkedBlockingQueue[Option[String]]() + + // Owns the writing end for the same reason. + private val writer: ExecutorService = Executors.newSingleThreadExecutor { r => + val t = new Thread(r, "python-worker-stdin") + t.setDaemon(true) + t + } + + private val outThread: Thread = { + val t = new Thread(() => { + val r = + new BufferedReader(new InputStreamReader(process.getInputStream, StandardCharsets.UTF_8)) + try { + var line = r.readLine() + while (line != null) { + lines.put(Some(line)) + line = r.readLine() + } + } catch { case NonFatal(_) => () } + finally lines.put(None) + }) + t.setDaemon(true) + t.setName("python-worker-stdout") + t.start() + t + } + + private val errThread: Thread = { + val t = new Thread(() => { + val r = + new BufferedReader(new InputStreamReader(process.getErrorStream, StandardCharsets.UTF_8)) + try { + var line = r.readLine() + while (line != null) { + errBuf.synchronized(errBuf.append(line).append('\n')) + line = r.readLine() + } + } catch { case NonFatal(_) => () } + }) + t.setDaemon(true) + t.setName("python-worker-stderr") + t.start() + t + } + + /** Wait for the worker's startup `{"ready": true}`; if it dies first (e.g. an + * import failed) or never gets there, surface its stderr. A worker that + * fails here is killed: it is not in the pool's set yet, so nothing else + * will reap it. + */ + def awaitReady(timeoutMillis: Long): Unit = { + val line = nextLine(timeoutMillis, "signal ready") + if (!objectMapper.readTree(line).path("ready").asBoolean(false)) { Review Comment: `Worker.run` converts a failed protocol parse into `WorkerDiedException` (`:370`); `awaitReady` parses the same protocol here and does not. A non-JSON first line throws a raw `JsonParseException` past the caller's fallback. It also orphans the interpreter: `destroy()` here runs only on the `ready == false` branch, and the worker is not in `all` yet (`:253`), so the shutdown hook cannot reap it. Not reachable with the two scripts in this PR, since both write their JSON first under `-I`. But the class doc sells this seam as generic, and #7149 brings the second worker script. -- 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]
