aglinxinyuan commented on code in PR #6132: URL: https://github.com/apache/texera/pull/6132#discussion_r3524376840
########## amber/src/test/integration/org/apache/texera/amber/engine/e2e/MultiRegionWorkflowIntegrationSpec.scala: ########## @@ -0,0 +1,286 @@ +/* + * 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.e2e + +import com.twitter.util.{Await, Duration, Promise, Return} +import com.typesafe.scalalogging.Logger +import org.apache.pekko.actor.{ActorSystem, Props} +import org.apache.pekko.testkit.{ImplicitSender, TestKit} +import org.apache.pekko.util.Timeout +import org.apache.texera.amber.clustering.SingleNodeListener +import org.apache.texera.amber.core.workflow.PortIdentity +import org.apache.texera.amber.engine.architecture.controller.{ + ControllerConfig, + ExecutionStateUpdate +} +import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmptyRequest +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.COMPLETED +import org.apache.texera.amber.engine.architecture.scheduling.CostBasedScheduleGenerator +import org.apache.texera.amber.engine.common.AmberRuntime +import org.apache.texera.amber.engine.common.client.AmberClient +import org.apache.texera.amber.engine.common.virtualidentity.util.CONTROLLER +import org.apache.texera.amber.engine.e2e.TestUtils.{ + buildWorkflow, + cleanupWorkflowExecutionData, + initiateTexeraDBForTestCases, + runWorkflowAndReadTerminalResults, + setUpWorkflowExecutionData +} +import org.apache.texera.amber.operator.TestOperators +import org.apache.texera.amber.operator.source.scan.text.TextInputSourceOpDesc +import org.apache.texera.amber.operator.union.UnionOpDesc +import org.apache.texera.amber.tags.IntegrationTest +import org.apache.texera.workflow.LogicalLink +import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, Outcome, Retries} +import org.scalatest.flatspec.AnyFlatSpecLike + +import scala.concurrent.duration._ + +/** + * End-to-end coverage for a workflow that executes across MULTIPLE regions. + * + * Cross-region data used to be delivered by dedicated cache-source operators; + * #3425 replaced them with input-port materialization reader threads. Before + * removing the now-dead cache-source plumbing, this spec adds the missing + * end-to-end coverage for that multi-region path: until now it was exercised + * only by scheduling unit tests (region counting on a physical plan), so a + * regression in the materialization read/write path could pass CI unnoticed. + * + * This spec drives a big "X"-shaped workflow to guard that path: + * + * {{{ + * csvLeft ─┬─▶ join.build (port 0) ─┐ + * │ join ─▶ pythonSort ─▶ (terminal) + * csvRight ─┼─▶ join.probe (port 1) ─┘ + * │ + * csvLeft ─┼─▶ union ───────────────────────────────────▶ (terminal) + * csvRight ─┘ (two links fan into union's single port) + * }}} + * + * The hash join's probe-depends-on-build ordering forces the build input to be + * materialized, cutting the plan into >=2 regions whose boundary is crossed by + * the reader-thread path. A Python UDF (sort) makes it a real end-to-end run + * with worker subprocesses, so it is class-level `@IntegrationTest` tagged and + * routed to the `amber-integration` CI job (which provisions Python deps); the + * lighter `amber` job excludes this tag. + */ +@IntegrationTest +class MultiRegionWorkflowIntegrationSpec + extends TestKit(ActorSystem("MultiRegionWorkflowIntegrationSpec", AmberRuntime.pekkoConfig)) + with ImplicitSender + with AnyFlatSpecLike + with BeforeAndAfterAll + with BeforeAndAfterEach + with Retries { + + /** + * Retry each test once if it fails. Mirrors the other e2e integration specs: + * in CI there is a small chance the run does not observe "COMPLETED", so a + * single retry stabilizes the suite until that root cause is fixed. + */ + override def withFixture(test: NoArgTest): Outcome = + withRetry { super.withFixture(test) } + + implicit val timeout: Timeout = Timeout(5.seconds) + + private val logger = Logger("MultiRegionWorkflowIntegrationSpecLogger") + private val specId = 5 + private val ctx = TestUtils.workflowContext(specId) + + // country_sales_small.csv has 100 rows with a unique "Order ID" per row. + private val sourceRowCount = 100 + + // Buffer every input tuple, then emit them ordered by "Order ID" once the + // input port is exhausted. This makes the UDF a genuine sort and forces a real + // Python worker to process data that crossed a region boundary. + private val sortByOrderIdCode = + """ + |from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | def open(self) -> None: + | self.buffer = [] + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | self.buffer.append(tuple_) + | yield from [] + | + | @overrides + | def on_finish(self, port: int) -> Iterator[Optional[TupleLike]]: + | for row in sorted(self.buffer, key=lambda t: t["Order ID"]): + | yield row + |""".stripMargin + + override protected def beforeEach(): Unit = { + setUpWorkflowExecutionData(specId) + } + + override protected def afterEach(): Unit = { + cleanupWorkflowExecutionData(specId) + } + + override def beforeAll(): Unit = { + system.actorOf(Props[SingleNodeListener](), "cluster-info") + // These test cases access postgres in CI, but occasionally the jdbc driver cannot be found during CI run. + // Explicitly load the JDBC driver to avoid flaky CI failures. + Class.forName("org.postgresql.Driver") + initiateTexeraDBForTestCases() + warmupOnce() + } + + override def afterAll(): Unit = { + TestKit.shutdownActorSystem(system) + } + + /** + * Runs a TextInput -> Python UDF workflow once before the timed test so the + * Python worker cold-start is paid here, not inside the timed run. Capped and + * wrapped so warmup can never fail or hang the suite. + */ + private def warmupOnce(): Unit = { + val warmupCap = Duration.fromSeconds(60) + setUpWorkflowExecutionData(specId) + var client: AmberClient = null + try { + val src = new TextInputSourceOpDesc() + src.textInput = "warmup" + val udf = TestOperators.pythonOpDesc() + val warmupCtx = TestUtils.workflowContext(specId) + val workflow = buildWorkflow( + List(src, udf), + List( + LogicalLink( + src.operatorIdentifier, + PortIdentity(), + udf.operatorIdentifier, + PortIdentity() + ) + ), + warmupCtx + ) + client = new AmberClient( + system, + workflow.context, + workflow.physicalPlan, + ControllerConfig.default, + _ => {} + ) + val completion = Promise[Unit]() + client.registerCallback[ExecutionStateUpdate](evt => { + if (evt.state == COMPLETED) completion.updateIfEmpty(Return(())) + }) Review Comment: This is already handled on the current head: the AmberClient error callback is `e => completion.updateIfEmpty(Throw(e))` and a `FatalError` callback also fails the completion Promise, so a broken warmup returns immediately instead of waiting out the 60s cap (the surrounding try/catch then logs and lets tests run cold). Resolving. -- 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]
