From 480b7c173b4a3f432937c2f7ddfcf2fbcc567476 Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Sun, 9 Aug 2026 17:57:45 -0700 Subject: [PATCH] test(amber): replace commented-out PythonWorkflowWorkerSpec with pythonworker proxy unit tests PythonWorkflowWorkerSpec.scala has been fully commented out for years and no longer compiles against today's APIs; the actor-level path it targeted is covered by the e2e tests. Replace it with real unit tests for the JVM side of the JVM<->Python Arrow Flight bridge, which had none: - PythonProxyServerSpec: handshake port promise + ok reply; control actions route ControlInvocation/ReturnInvocation to the output gateway on the control channel and ack with a little-endian credit; Data/State/ ECM puts are reassembled into DataFrame/StateFrame (loop envelope preserved)/EmbeddedControlMessage and acked with credits. - PythonProxyClientSpec: heartbeat-before-drain connection order; queued control/actor commands and Data/State/ECM payloads arrive as the right Flight actions/puts with intact payloads; queue-size acks drive getQueuedCredit; retry exhaustion and non-ack heartbeats abort with WorkflowRuntimeException; close() before connecting does not throw. Both specs stand in for the Python worker with plain Arrow Flight components, so no Python process is involved. --- .../pythonworker/PythonProxyClientSpec.scala | 379 ++++++++++++++++++ .../pythonworker/PythonProxyServerSpec.scala | 306 ++++++++++++++ .../PythonWorkflowWorkerSpec.scala | 201 ---------- 3 files changed, 685 insertions(+), 201 deletions(-) create mode 100644 amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala create mode 100644 amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala delete mode 100644 amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala new file mode 100644 index 00000000000..f536b61d1c3 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClientSpec.scala @@ -0,0 +1,379 @@ +/* + * 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.engine.architecture.pythonworker + +import com.twitter.util.Promise +import org.apache.arrow.flight.{ + Action, + FlightProducer, + FlightServer, + FlightStream, + Location, + NoOpFlightProducer, + PutResult, + Result +} +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.VarBinaryVector +import org.apache.texera.amber.core.WorkflowRuntimeException +import org.apache.texera.amber.core.state.State +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple} +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + ChannelIdentity, + EmbeddedControlMessageIdentity +} +import org.apache.texera.amber.engine.architecture.pythonworker.WorkerBatchInternalQueue.{ + DataElement, + EmbeddedControlMessageElement +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + ControlInvocation, + EmbeddedControlMessage, + EmbeddedControlMessageType, + EmptyRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} +import org.apache.texera.amber.engine.common.actormessage.{CreditUpdate, PythonActorMessage} +import org.apache.texera.amber.engine.common.ambermessage.{ + DataFrame, + DirectControlMessagePayloadV2, + PythonControlMessage, + PythonDataHeader, + StateFrame +} +import org.apache.texera.amber.util.ArrowUtils +import org.scalatest.flatspec.AnyFlatSpec + +import java.net.ServerSocket +import java.nio.charset.StandardCharsets +import java.util.concurrent.ConcurrentLinkedQueue +import scala.jdk.CollectionConverters._ + +/** + * Exercises PythonProxyClient without a Python process: a Scala + * NoOpFlightProducer stands in for the Python worker's Flight server + * (`network_receiver.py`), speaking the same action/put protocol — heartbeat + * acks, per-action queue-size replies, and per-put credit metadata. + */ +class PythonProxyClientSpec extends AnyFlatSpec { + + private val clientWorkerId = ActorVirtualIdentity("python-proxy-client") + private val controllerId = ActorVirtualIdentity("CONTROLLER") + private val upstreamChannel = + ChannelIdentity(controllerId, clientWorkerId, isControl = true) + + private def freePort: Int = { + val socket = new ServerSocket(0) + try socket.getLocalPort + finally socket.close() + } + + private def quietly(f: => Unit): Unit = + try f + catch { case _: Exception => () } + + private def awaitTrue(timeoutMs: Long = 10000)(cond: => Boolean): Unit = { + val deadline = System.currentTimeMillis() + timeoutMs + while (!cond && System.currentTimeMillis() < deadline) Thread.sleep(20) + assert(cond, s"condition not met within ${timeoutMs}ms") + } + + // --------------------------------------------------------------------------- + // Stand-in for the Python worker's Flight server: records every action and + // put it receives and acks them the way network_receiver.py does. + // --------------------------------------------------------------------------- + + private case class RecordedPut( + header: PythonDataHeader, + tuples: Vector[Tuple], + ecmBytes: Option[Array[Byte]] + ) + + private class FakePythonWorkerServer(heartbeatBody: String) extends NoOpFlightProducer { + val actionsReceived = new ConcurrentLinkedQueue[(String, Array[Byte])]() + val putsReceived = new ConcurrentLinkedQueue[RecordedPut]() + @volatile var reportedQueueSize: Long = 0L + + override def doAction( + context: FlightProducer.CallContext, + action: Action, + listener: FlightProducer.StreamListener[Result] + ): Unit = { + actionsReceived.add((action.getType, action.getBody)) + action.getType match { + case "heartbeat" => + listener.onNext(new Result(heartbeatBody.getBytes(StandardCharsets.UTF_8))) + listener.onCompleted() + case "control" | "actor" => + // The ack body is the Python worker's queue size, sent as a decimal + // string; the client parses it for flow control. + listener.onNext( + new Result(reportedQueueSize.toString.getBytes(StandardCharsets.UTF_8)) + ) + listener.onCompleted() + case _ => // "shutdown": the client does not consume a reply + listener.onCompleted() + } + } + + override def acceptPut( + context: FlightProducer.CallContext, + flightStream: FlightStream, + ackStream: FlightProducer.StreamListener[PutResult] + ): Runnable = { () => + { + val header = PythonDataHeader.parseFrom(flightStream.getDescriptor.getCommand) + val tuples = Vector.newBuilder[Tuple] + var ecmBytes: Option[Array[Byte]] = None + while (flightStream.next()) { + val root = flightStream.getRoot + if (header.payloadType == "ECM") { + ecmBytes = Some(root.getVector("payload").asInstanceOf[VarBinaryVector].get(0)) + } else { + (0 until root.getRowCount).foreach(i => tuples += ArrowUtils.getTexeraTuple(i, root)) + } + } + putsReceived.add(RecordedPut(header, tuples.result(), ecmBytes)) + // Ack with the queue size as put metadata, the way the Python side does. + val ackAllocator = new RootAllocator(1024) + try { + val buf = ackAllocator.buffer(java.lang.Long.BYTES) + buf.writeLong(reportedQueueSize) + ackStream.onNext(PutResult.metadata(buf)) + buf.close() + } finally { + ackAllocator.close() + } + } + } + } + + // --------------------------------------------------------------------------- + // Test harness — fake server + a PythonProxyClient main loop on its own thread + // --------------------------------------------------------------------------- + + private class ClientFixture(heartbeatBody: String = "ack") { + val producer = new FakePythonWorkerServer(heartbeatBody) + val serverAllocator = new RootAllocator() + val port: Int = freePort + val server: FlightServer = FlightServer + .builder(serverAllocator, Location.forGrpcInsecure("localhost", port), producer) + .build() + server.start() + val portPromise: Promise[Int] = Promise[Int]() + portPromise.setValue(port) + val client = new PythonProxyClient(portPromise, clientWorkerId) + private val loopThread = new Thread(() => client.run(), "python-proxy-client-test-loop") + loopThread.setDaemon(true) + // The tear-down interrupt below surfaces as an InterruptedException out of + // the blocking queue take; keep it off stderr. + loopThread.setUncaughtExceptionHandler((_, _) => ()) + + def start(): Unit = loopThread.start() + + def close(): Unit = { + quietly(client.close()) + loopThread.interrupt() // unblock the main loop's queue take + loopThread.join(5000) + quietly(server.close()) + quietly(serverAllocator.close()) + } + } + + private def withClient(heartbeatBody: String = "ack")(test: ClientFixture => Unit): Unit = { + val fixture = new ClientFixture(heartbeatBody) + try test(fixture) + finally fixture.close() + } + + // --------------------------------------------------------------------------- + // connection + control path + // --------------------------------------------------------------------------- + + "PythonProxyClient" should + "connect via heartbeat and forward an enqueued ControlInvocation as a control action" in { + withClient() { fixture => + fixture.start() + val invocation = ControlInvocation( + "OpenExecutor", + EmptyRequest(), + AsyncRPCContext(controllerId, clientWorkerId), + 5L + ) + fixture.client.enqueueCommand(invocation, upstreamChannel) + + awaitTrue()(fixture.producer.actionsReceived.asScala.exists(_._1 == "control")) + assert( + fixture.producer.actionsReceived.asScala.head._1 == "heartbeat", + "the client must handshake before draining its queue" + ) + val body = fixture.producer.actionsReceived.asScala.find(_._1 == "control").get._2 + val parsed = PythonControlMessage.parseFrom(body) + assert(parsed.tag == upstreamChannel) + assert( + parsed.payload == + DirectControlMessagePayloadV2.defaultInstance.withControlInvocation(invocation) + ) + } + } + + it should "forward an enqueued ReturnInvocation as a control action" in { + withClient() { fixture => + fixture.start() + val reply = ReturnInvocation(42L, EmptyReturn()) + fixture.client.enqueueCommand(reply, upstreamChannel) + + awaitTrue()(fixture.producer.actionsReceived.asScala.exists(_._1 == "control")) + val body = fixture.producer.actionsReceived.asScala.find(_._1 == "control").get._2 + val parsed = PythonControlMessage.parseFrom(body) + assert(parsed.tag == upstreamChannel) + assert( + parsed.payload == + DirectControlMessagePayloadV2.defaultInstance.withReturnInvocation(reply) + ) + } + } + + it should "forward an enqueued actor command as an actor action" in { + withClient() { fixture => + fixture.start() + fixture.client.enqueueActorCommand(CreditUpdate()) + + awaitTrue()(fixture.producer.actionsReceived.asScala.exists(_._1 == "actor")) + val body = fixture.producer.actionsReceived.asScala.find(_._1 == "actor").get._2 + assert(PythonActorMessage.parseFrom(body) == PythonActorMessage(CreditUpdate())) + } + } + + it should "update its Python-queue credit from the action ack" in { + withClient() { fixture => + fixture.producer.reportedQueueSize = 7L + fixture.start() + assert(fixture.client.getQueuedCredit == 0L) + + fixture.client.enqueueActorCommand(CreditUpdate()) + awaitTrue()(fixture.client.getQueuedCredit == 7L) + } + } + + // --------------------------------------------------------------------------- + // data path — Data / State / ECM puts + // --------------------------------------------------------------------------- + + it should "stream an enqueued DataFrame to the server under a Data header" in { + withClient() { fixture => + fixture.start() + val schema = Schema() + .add(new Attribute("v", AttributeType.INTEGER)) + .add(new Attribute("s", AttributeType.STRING)) + def tuple(v: Int, s: String): Tuple = + Tuple.builder(schema).addSequentially(Array(Int.box(v), s)).build() + val tuples = Array(tuple(1, "a"), tuple(2, "b")) + val dataChannel = ChannelIdentity(clientWorkerId, controllerId, isControl = false) + + fixture.client.enqueueData(DataElement(DataFrame(tuples), dataChannel)) + + awaitTrue()(!fixture.producer.putsReceived.isEmpty) + val put = fixture.producer.putsReceived.peek() + assert(put.header == PythonDataHeader(dataChannel, "Data")) + assert(put.tuples == tuples.toVector) + } + } + + it should "stream an enqueued StateFrame as a single State row carrying the loop envelope" in { + withClient() { fixture => + fixture.start() + val state = State(Map("count" -> "5")) + val dataChannel = ChannelIdentity(clientWorkerId, controllerId, isControl = false) + + fixture.client.enqueueData( + DataElement(StateFrame(state, loopCounter = 3L, loopStartId = "loop-start-1"), dataChannel) + ) + + awaitTrue()(!fixture.producer.putsReceived.isEmpty) + val put = fixture.producer.putsReceived.peek() + assert(put.header == PythonDataHeader(dataChannel, "State")) + assert(put.tuples.size == 1) + val row = put.tuples.head + assert(State.fromTuple(row) == state) + assert(State.loopCounterFrom(row) == 3L) + assert(State.loopStartIdFrom(row) == "loop-start-1") + } + } + + it should "stream an enqueued EmbeddedControlMessage as a serialized ECM put" in { + withClient() { fixture => + fixture.start() + val ecm = EmbeddedControlMessage( + EmbeddedControlMessageIdentity("ecm-1"), + EmbeddedControlMessageType.NO_ALIGNMENT, + Seq.empty, + Map.empty + ) + val dataChannel = ChannelIdentity(clientWorkerId, controllerId, isControl = false) + + fixture.client.enqueueData(EmbeddedControlMessageElement(ecm, dataChannel)) + + awaitTrue()(!fixture.producer.putsReceived.isEmpty) + val put = fixture.producer.putsReceived.peek() + assert(put.header == PythonDataHeader(dataChannel, "ECM")) + assert(put.ecmBytes.isDefined) + assert(EmbeddedControlMessage.parseFrom(put.ecmBytes.get) == ecm) + } + } + + // --------------------------------------------------------------------------- + // failure paths + // --------------------------------------------------------------------------- + + it should "abort with WorkflowRuntimeException once connection retries are exhausted" in { + // A free port with nothing listening: every connection attempt is refused. + val portPromise = Promise[Int]() + portPromise.setValue(freePort) + val client = new PythonProxyClient(portPromise, ActorVirtualIdentity("no-server")) + assertThrows[WorkflowRuntimeException] { + client.run() + } + } + + it should "abort when the server heartbeat does not reply ack" in { + withClient(heartbeatBody = "nak") { fixture => + // run() on the test thread: it must give up after the retry budget + // because the heartbeat body check fails on every attempt. + assertThrows[WorkflowRuntimeException] { + fixture.client.run() + } + } + } + + it should "tolerate close() before any connection is established" in { + val portPromise = Promise[Int]() + portPromise.setValue(freePort) + val client = new PythonProxyClient(portPromise, ActorVirtualIdentity("never-connected")) + // Must not throw: the internal null flight client is handled. + client.close() + succeed + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala new file mode 100644 index 00000000000..1e4c531a0b5 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServerSpec.scala @@ -0,0 +1,306 @@ +/* + * 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.engine.architecture.pythonworker + +import com.twitter.util.{Await, Duration, Promise} +import org.apache.arrow.flight.{Action, FlightClient, FlightDescriptor, Location, SyncPutListener} +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema => ArrowSchema} +import org.apache.arrow.vector.{VarBinaryVector, VectorSchemaRoot} +import org.apache.texera.amber.core.state.State +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple} +import org.apache.texera.amber.core.virtualidentity.{ + ActorVirtualIdentity, + ChannelIdentity, + EmbeddedControlMessageIdentity +} +import org.apache.texera.amber.engine.architecture.messaginglayer.NetworkOutputGateway +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ + AsyncRPCContext, + ControlInvocation, + EmbeddedControlMessage, + EmbeddedControlMessageType, + EmptyRequest +} +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} +import org.apache.texera.amber.engine.common.ambermessage.{ + DataFrame, + DirectControlMessagePayloadV2, + PythonControlMessage, + PythonDataHeader, + StateFrame, + WorkflowFIFOMessage +} +import org.apache.texera.amber.util.ArrowUtils +import org.scalatest.flatspec.AnyFlatSpec + +import java.nio.charset.StandardCharsets +import java.nio.{ByteBuffer, ByteOrder} +import java.util.concurrent.{ConcurrentLinkedQueue, TimeUnit} +import scala.jdk.CollectionConverters._ + +/** + * Exercises the JVM-side Flight server of the Python worker bridge without a + * Python process: the test plays the Python worker's part with a plain Arrow + * FlightClient and asserts on what the server forwards to its + * NetworkOutputGateway. This is the JVM half of the protocol implemented by + * `network_sender.py` / `network_receiver.py`. + */ +class PythonProxyServerSpec extends AnyFlatSpec { + + private val jvmWorkerId = ActorVirtualIdentity("jvm-proxy-server") + private val pythonWorkerId = ActorVirtualIdentity("python-worker-1") + private val controllerId = ActorVirtualIdentity("CONTROLLER") + + // --------------------------------------------------------------------------- + // Test harness — a running PythonProxyServer plus a FlightClient connected to + // it, with every message the server emits captured from the output gateway. + // --------------------------------------------------------------------------- + + private class ServerFixture { + val sentMessages = new ConcurrentLinkedQueue[WorkflowFIFOMessage]() + val gateway = new NetworkOutputGateway(jvmWorkerId, msg => { sentMessages.add(msg); () }) + val portPromise: Promise[Int] = Promise[Int]() + val server = new PythonProxyServer(gateway, jvmWorkerId, portPromise) + server.run() // FlightServer.start() binds and returns; gRPC threads serve + val clientAllocator = new RootAllocator() + val flightClient: FlightClient = FlightClient + .builder( + clientAllocator, + Location.forGrpcInsecure("localhost", server.getPortNumber.get()) + ) + .build() + + def messages: List[WorkflowFIFOMessage] = sentMessages.asScala.toList + + def close(): Unit = { + quietly(flightClient.close()) + quietly(server.close()) + quietly(clientAllocator.close()) + } + } + + private def quietly(f: => Unit): Unit = + try f + catch { case _: Exception => () } + + private def withServer(test: ServerFixture => Unit): Unit = { + val fixture = new ServerFixture + try test(fixture) + finally fixture.close() + } + + private def awaitTrue(timeoutMs: Long = 10000)(cond: => Boolean): Unit = { + val deadline = System.currentTimeMillis() + timeoutMs + while (!cond && System.currentTimeMillis() < deadline) Thread.sleep(20) + assert(cond, s"condition not met within ${timeoutMs}ms") + } + + /** Mirror of PythonProxyClient.writeArrowStream: put one batch of rows. */ + private def putBatch( + fixture: ServerFixture, + header: PythonDataHeader, + arrowSchema: ArrowSchema, + fillRoot: VectorSchemaRoot => Unit + ): Long = { + val descriptor = FlightDescriptor.command(header.toByteArray) + val listener = new SyncPutListener + val root = VectorSchemaRoot.create(arrowSchema, fixture.clientAllocator) + try { + val writer = fixture.flightClient.startPut(descriptor, root, listener) + root.allocateNew() + fillRoot(root) + writer.putNext() + root.clear() + writer.completed() + val ack = listener.poll(5, TimeUnit.SECONDS) + assert(ack != null, "no credit ack received for the put within 5s") + val ackBuf = ack.getApplicationMetadata + val credit = ackBuf.getLong(0) + ackBuf.close() + listener.close() + credit + } finally { + root.close() + } + } + + // --------------------------------------------------------------------------- + // Port + handshake + // --------------------------------------------------------------------------- + + "PythonProxyServer" should "bind a free local port exposed via getPortNumber" in { + withServer { fixture => + assert(fixture.server.getPortNumber.get() > 0) + } + } + + it should "complete the port-number promise and reply ok on a handshake action" in { + withServer { fixture => + val result = fixture.flightClient + .doAction(new Action("handshake", "45678".getBytes(StandardCharsets.UTF_8))) + .next() + assert(new String(result.getBody, StandardCharsets.UTF_8) == "ok") + assert(Await.result(fixture.portPromise, Duration.fromSeconds(5)) == 45678) + } + } + + // --------------------------------------------------------------------------- + // control actions — routed to the gateway on the control channel, acked with + // a little-endian credit value + // --------------------------------------------------------------------------- + + it should "route a control action carrying a ControlInvocation to the tag's receiver" in { + withServer { fixture => + val tag = ChannelIdentity(pythonWorkerId, controllerId, isControl = true) + val invocation = ControlInvocation( + "StartWorker", + EmptyRequest(), + AsyncRPCContext(pythonWorkerId, controllerId), + 3L + ) + val controlMessage = PythonControlMessage( + tag, + DirectControlMessagePayloadV2.defaultInstance.withControlInvocation(invocation) + ) + + val results = fixture.flightClient.doAction(new Action("control", controlMessage.toByteArray)) + val credit = ByteBuffer.wrap(results.next().getBody).order(ByteOrder.LITTLE_ENDIAN).getLong + assert(credit == 30L, "the action ack must carry the credit value little-endian") + assert(!results.hasNext, "a control action must produce exactly one result") + + val out = fixture.messages + assert(out.size == 1) + // The gateway sends on ITS control channel to the tag's receiver. + assert(out.head.channelId == ChannelIdentity(jvmWorkerId, controllerId, isControl = true)) + assert(out.head.payload == invocation) + } + } + + it should "route a control action carrying a ReturnInvocation to the tag's receiver" in { + withServer { fixture => + val tag = ChannelIdentity(pythonWorkerId, controllerId, isControl = true) + val reply = ReturnInvocation(42L, EmptyReturn()) + val controlMessage = PythonControlMessage( + tag, + DirectControlMessagePayloadV2.defaultInstance.withReturnInvocation(reply) + ) + + val results = fixture.flightClient.doAction(new Action("control", controlMessage.toByteArray)) + results.next() + + val out = fixture.messages + assert(out.size == 1) + assert(out.head.channelId == ChannelIdentity(jvmWorkerId, controllerId, isControl = true)) + assert(out.head.payload == reply) + } + } + + // --------------------------------------------------------------------------- + // puts — data / state / ECM payload types, acked with a credit value + // --------------------------------------------------------------------------- + + it should "deliver a Data put as a DataFrame on the header's channel and ack with credits" in { + withServer { fixture => + val to = + ChannelIdentity(pythonWorkerId, ActorVirtualIdentity("downstream"), isControl = false) + val schema = Schema() + .add(new Attribute("v", AttributeType.INTEGER)) + .add(new Attribute("s", AttributeType.STRING)) + def tuple(v: Int, s: String): Tuple = + Tuple.builder(schema).addSequentially(Array(Int.box(v), s)).build() + val tuples = Array(tuple(1, "a"), tuple(2, "b")) + + val credit = putBatch( + fixture, + PythonDataHeader(to, "Data"), + ArrowUtils.fromTexeraSchema(schema), + root => tuples.foreach(t => ArrowUtils.appendTexeraTuple(t, root)) + ) + assert(credit == 31L, "the put ack must carry the credit value") + + awaitTrue()(fixture.messages.nonEmpty) + val out = fixture.messages + assert(out.size == 1) + assert(out.head.channelId == to) + assert(out.head.payload == DataFrame(tuples)) + } + } + + it should "reassemble a single-row State put into a StateFrame carrying the loop envelope" in { + withServer { fixture => + val to = + ChannelIdentity(pythonWorkerId, ActorVirtualIdentity("downstream"), isControl = false) + val state = State(Map("count" -> "5")) + val row = state.toTuple(loopCounter = 3L, loopStartId = "loop-start-1") + + putBatch( + fixture, + PythonDataHeader(to, "State"), + ArrowUtils.fromTexeraSchema(State.schema), + root => ArrowUtils.appendTexeraTuple(row, root) + ) + + awaitTrue()(fixture.messages.nonEmpty) + val out = fixture.messages + assert(out.size == 1) + assert(out.head.channelId == to) + assert(out.head.payload == StateFrame(state, 3L, "loop-start-1")) + } + } + + it should "parse an ECM put back into the original EmbeddedControlMessage" in { + withServer { fixture => + val to = + ChannelIdentity(pythonWorkerId, ActorVirtualIdentity("downstream"), isControl = false) + val ecm = EmbeddedControlMessage( + EmbeddedControlMessageIdentity("ecm-1"), + EmbeddedControlMessageType.NO_ALIGNMENT, + Seq.empty, + Map.empty + ) + // Mirror of PythonProxyClient.sendECM: one VarBinary "payload" column, + // one row holding the serialized ECM. + val field = new Field("payload", FieldType.nullable(new ArrowType.Binary), null) + val arrowSchema = new ArrowSchema(List(field).asJava) + + putBatch( + fixture, + PythonDataHeader(to, "ECM"), + arrowSchema, + root => { + val vector = root.getVector("payload").asInstanceOf[VarBinaryVector] + vector.setSafe(0, ecm.toByteArray) + vector.setValueCount(1) + root.setRowCount(1) + } + ) + + awaitTrue()(fixture.messages.nonEmpty) + val out = fixture.messages + assert(out.size == 1) + assert(out.head.channelId == to) + assert(out.head.payload == ecm) + } + } +} diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala deleted file mode 100644 index 9741dc4e453..00000000000 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonWorkflowWorkerSpec.scala +++ /dev/null @@ -1,201 +0,0 @@ -/* - * 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.engine.architecture.pythonworker -// -//import org.apache.pekko.actor.{ActorRef, ActorSystem, Props} -//import org.apache.pekko.testkit.{ImplicitSender, TestActorRef, TestKit} -//import org.apache.texera.amber.clustering.SingleNodeListener -//import org.apache.texera.amber.engine.architecture.common.WorkflowActor.{NetworkAck, NetworkMessage} -//import org.apache.texera.amber.engine.architecture.pythonworker.promisehandlers.InitializeOperatorLogicHandler.InitializeOperatorLogic -//import org.apache.texera.amber.engine.architecture.sendsemantics.partitionings.OneToOnePartitioning -//import org.apache.texera.amber.engine.architecture.worker.controlcommands.LinkOrdinal -//import org.apache.texera.amber.engine.architecture.worker.promisehandlers.AddPartitioningHandler.AddPartitioning -//import org.apache.texera.amber.engine.architecture.worker.promisehandlers.OpenOperatorHandler.OpenOperator -//import org.apache.texera.amber.engine.architecture.worker.promisehandlers.UpdateInputLinkingHandler.UpdateInputLinking -//import org.apache.texera.amber.engine.common.Constants -//import org.apache.texera.amber.engine.common.ambermessage.{ -// ChannelIdentity, -// ControlPayload, -// DataFrame, -// DataPayload, -// EndOfUpstream, -// WorkflowFIFOMessage -//} -//import org.apache.texera.amber.engine.common.rpc.AsyncRPCClient.{ControlInvocation, ReturnInvocation} -//import org.apache.texera.amber.engine.common.virtualidentity.util.COORDINATOR -//import org.apache.texera.amber.core.virtualidentity.{ -// ActorVirtualIdentity, -// PhysicalLink, -// PhysicalLink, -// OperatorIdentity -//} -//import org.apache.texera.amber.engine.e2e.TestOperators -//import org.apache.texera.workflow.common.tuple.Tuple -//import org.apache.texera.workflow.common.tuple.schema.{Attribute, AttributeType, Schema} -//import org.scalamock.scalatest.MockFactory -//import org.scalatest.BeforeAndAfterAll -//import org.scalatest.flatspec.AnyFlatSpecLike -// -//import scala.concurrent.duration.DurationInt -// -//class PythonWorkflowWorkerSpec -// extends TestKit(ActorSystem("PythonWorkerSpec")) -// with ImplicitSender -// with AnyFlatSpecLike -// with BeforeAndAfterAll -// with MockFactory { -// -// override def beforeAll: Unit = { -// system.actorOf(Props[SingleNodeListener], "cluster-info") -// } -// override def afterAll: Unit = { -// TestKit.shutdownActorSystem(system) -// } -// private val identifier1 = ActorVirtualIdentity("worker-1") -// private val identifier2 = ActorVirtualIdentity("worker-2") -// private val operatorIdentity = OperatorIdentity("testWorkflow", "testOperator") -// private val layerId1 = -// PhysicalLink(operatorIdentity.workflow, operatorIdentity.operator, "1st-layer") -// private val layerId2 = -// PhysicalLink(operatorIdentity.workflow, operatorIdentity.operator, "2nd-layer") -// private val pythonOp = TestOperators.pythonOpDesc() -// private val link = PhysicalLink(layerId1, 0, layerId2, 0) -// private val schema = Schema -// .newBuilder() -// .add(new Attribute("text", AttributeType.STRING)) -// .build() -// private val initialization = InitializeOperatorLogic( -// pythonOp.code, -// isSource = false, -// Seq(LinkOrdinal(link, 0)), -// Seq(LinkOrdinal(link, 0)), -// schema -// ) -// -// def sendControlToWorker( -// worker: ActorRef, -// controls: Array[ControlInvocation], -// beginSeqNum: Long = 0 -// ): Unit = { -// var seq = beginSeqNum -// controls.foreach { ctrl => -// worker ! NetworkMessage( -// seq, -// WorkflowFIFOMessage(ChannelIdentity(COORDINATOR, identifier1, true), seq, ctrl) -// ) -// val received = receiveWhile(3.seconds) { -// case NetworkAck(id, credits) => -// // pass -// case NetworkMessage(id, fifoPayload) => -// fifoPayload.payload.asInstanceOf[ControlPayload] match { -// case ControlInvocation(commandID, command) => assert(commandID == seq) -// case ReturnInvocation(originalCommandID, controlReturn) => -// assert(originalCommandID == seq) -// case _ => ??? -// } -// worker ! NetworkAck(id, Constants.unprocessedBatchesSizeLimitInBytesPerWorkerPair) -// } -// seq += 1 -// } -// } -// -// def mkWorker: ActorRef = TestActorRef(new PythonWorkflowWorker(identifier1)) -// -// "python worker" should "start" in { -// val worker = mkWorker -// sendControlToWorker(worker, Array(ControlInvocation(0, initialization))) -// } -// -// "python worker" should "process data" in { -// val worker = mkWorker -// sendControlToWorker(worker, Array(ControlInvocation(0, initialization))) -// val mockPolicy = OneToOnePartitioning(1, Array(identifier2)) -// val openControl = ControlInvocation(1, OpenOperator()) -// val invocation = ControlInvocation(2, AddPartitioning(link, mockPolicy)) -// val updateInputLinking = ControlInvocation(3, UpdateInputLinking(identifier2, link)) -// sendControlToWorker(worker, Array(openControl, invocation, updateInputLinking), 1) -// worker ! NetworkMessage( -// 4, -// WorkflowFIFOMessage( -// ChannelIdentity(identifier2, identifier1, false), -// 0, -// DataFrame( -// Array( -// Tuple -// .newBuilder(schema) -// .add("text", AttributeType.STRING, "123") -// .build() -// ) -// ) -// ) -// ) -// expectMsgClass(classOf[NetworkAck]) -// val data = receiveOne(30.seconds) -// assert(data.asInstanceOf[NetworkMessage].internalMessage.payload.isInstanceOf[DataFrame]) -// } -// -// "python worker" should "process data and receive end marker" in { -// val worker = mkWorker -// sendControlToWorker(worker, Array(ControlInvocation(0, initialization))) -// val mockPolicy = OneToOnePartitioning(100, Array(identifier2)) -// val openControl = ControlInvocation(1, OpenOperator()) -// val invocation = ControlInvocation(2, AddPartitioning(link, mockPolicy)) -// val updateInputLinking = ControlInvocation(3, UpdateInputLinking(identifier2, link)) -// sendControlToWorker(worker, Array(openControl, invocation, updateInputLinking), 1) -// worker ! NetworkMessage( -// 4, -// WorkflowFIFOMessage( -// ChannelIdentity(identifier2, identifier1, false), -// 0, -// DataFrame( -// (0 until 100) -// .map(_ => -// Tuple -// .newBuilder(schema) -// .add("text", AttributeType.STRING, "123") -// .build() -// ) -// .toArray -// ) -// ) -// ) -// expectMsgClass(classOf[NetworkAck]) -// val data = receiveOne(30.seconds) -// assert(data.asInstanceOf[NetworkMessage].internalMessage.payload.isInstanceOf[DataFrame]) -// worker ! NetworkMessage( -// 5, -// WorkflowFIFOMessage( -// ChannelIdentity(identifier2, identifier1, false), -// 1, -// EndOfUpstream() -// ) -// ) -// expectMsgClass(classOf[NetworkAck]) -// receiveWhile(10.seconds) { -// case NetworkMessage(id, fifoPayload) => -// fifoPayload.payload match { -// case payload: ControlPayload => //skip -// case payload: DataPayload => assert(payload.isInstanceOf[EndOfUpstream]) -// case _ => ??? -// } -// } -// } -// -//}