wenjin272 commented on code in PR #1147: URL: https://github.com/apache/flink-agents/pull/1147#discussion_r4068158267
########## runtime/src/main/java/org/apache/flink/agents/runtime/operator/parallel/ParallelExecutionCoordinator.java: ########## @@ -0,0 +1,299 @@ +/* + * 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.flink.agents.runtime.operator.parallel; + +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.util.function.ThrowingRunnable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Supplier; + +import static org.apache.flink.util.Preconditions.checkNotNull; + +/** + * Coordinates worker execution of {@link ParallelExecutionTask}: callers {@link #addTask(Object) + * add} one node per queued task, and the coordinator submits one anonymous <em>dispatch permit</em> + * per node to an elastic {@link ThreadPoolExecutor}. A pool thread redeems a permit under the lock + * (node poll and FIFO task pull share the lock hold, keeping their binding correct — permits never + * name a task); commits always run on the mailbox thread in per-key taskIndex order. + */ +public final class ParallelExecutionCoordinator implements AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(ParallelExecutionCoordinator.class); + + private static final long WORKER_TERMINATION_TIMEOUT_SECONDS = 180L; + + @FunctionalInterface + public interface MailboxDispatcher { + void execute(ThrowingRunnable<? extends Exception> command, String description); + } + + private final ParallelExecutionTaskQueue schedule; + private final ParallelExecutionLock parallelExecutionLock; + private final ThreadPoolExecutor workerExecutor; + private final MailboxDispatcher mailboxDispatcher; + + /** Supplies a fresh work; the work pulls and prepares its own task inside {@code execute()}. */ + private final Supplier<ParallelExecutionTask> workFactory; + + /** Tasks pulled but not yet committed; read by {@link #isQuiesced()} for checkpoint quiesce. */ + private final AtomicInteger inFlightExecuting = new AtomicInteger(); + + /** + * When {@code true}, workers finish in-flight tasks but pull no new nodes. Set before a + * checkpoint, cleared once the snapshot is taken. + */ + private volatile boolean draining; + + private final Object drainMonitor = new Object(); + + /** + * Permits submitted and not yet fully redeemed; compared against the pool size by {@link + * DispatchQueue#offer} to decide when the pool must grow. + */ + private final AtomicInteger submittedPermits = new AtomicInteger(); + + /** + * Creates a coordinator with an elastic worker pool: one resident thread, on-demand growth up + * to {@code maxWorkers}, and idle retirement after {@code workerIdleTimeoutMillis}. Admission + * is not gated here; the operator bounds outstanding work via input-record backpressure. + */ + public ParallelExecutionCoordinator( + ParallelExecutionLock parallelExecutionLock, + MailboxDispatcher mailboxDispatcher, + Supplier<ParallelExecutionTask> workFactory, + int maxWorkers, + long workerIdleTimeoutMillis) { + this.parallelExecutionLock = checkNotNull(parallelExecutionLock); + this.schedule = new ParallelExecutionTaskQueue(parallelExecutionLock); + this.mailboxDispatcher = checkNotNull(mailboxDispatcher); + this.workFactory = checkNotNull(workFactory); + DispatchQueue dispatchQueue = new DispatchQueue(); + this.workerExecutor = + new ThreadPoolExecutor( Review Comment: These workers need to participate in the existing managed-worker interpreter lifecycle from #1092. With the default thread factory, `PythonInterpreterManager` treats them as unmanaged callers and forwards Python invocations to a separate callback worker. This introduces two regressions for Java actions using Python resources: - **Cross-thread Python object access:** `PythonVectorStore` receives `PyObject`s created on the callback worker, but reads their attributes and closes them on the original action worker. - **Resource-initialization deadlock:** the action worker holds the `ResourceCache.getResource()` monitor while initializing a Python chat model setup. It waits for the callback worker to run Python `open()`. That `open()` resolves the model connection through the same resource cache, so the callback worker waits for the monitor held by the action worker. Neither can proceed. Previously, this lookup ran on the same thread and could re-enter the monitor. I reproduced both with the production Java components and mocked native interpreter calls. Could we integrate the new pool with managed-worker interpreter ownership and cleanup, and add regression coverage for these paths? ########## runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java: ########## @@ -1049,17 +1398,256 @@ private void tryResumeProcessActionTasks() throws Exception { }); long[] pendingInputEvents = {0L}; + // Recovered pending input records will each be dequeued, processed, and decremented when + // they finish, so they must occupy in-flight budget too — count one unit per pending event. stateManager.forEachPendingInputEventKey( getKeyedStateBackend(), - (key, state) -> { - for (Event ignored : state.get()) { - eventRouter.getKeySegmentQueue().addKeyToLastSegment(key); - pendingInputEvents[0]++; - } - }); + (key, state) -> + state.get() + .forEach( + event -> { + inFlightInputRecords++; + pendingInputEvents[0]++; + eventRouter + .getKeySegmentQueue() + .addKeyToLastSegment(key); + })); builtInMetrics.restorePendingInputEvents(pendingInputEvents[0]); } + private void runWithParallelExecutionLock(ThrowingRunnable<? extends Exception> action) + throws Exception { + boolean acquired = false; + try { + parallelExecutionLock.acquireByMain(); + acquired = true; + action.run(); + } finally { + if (acquired) { + parallelExecutionLock.release(); + } + } + } + + /** + * A single action task's full worker lifecycle on the parallel engine: a pool thread {@code + * setup}s and {@code execute}s it under the lock; the mailbox thread later restores and commits + * it. All task-specific processing lives here; the operator keeps only mailbox flow. + */ + private final class Work implements ParallelExecutionTask { + private Object key; + private long recordIndex; + private long taskIndex; + private String contextKey; + + private ActionTask actionTask; + private long sequenceNumber; + private boolean replayCompletedAction; + + @Nullable private ActionTask.ActionTaskResult result; + @Nullable private Throwable failure; + + @Override + public void setup(Object key, long recordIndex, long taskIndex) { + this.key = key; + this.recordIndex = recordIndex; + this.taskIndex = taskIndex; + this.contextKey = resolveContextKey(key); + } + + @Override + public void execute() { + try { + // Pull + prepare under the lock hold: pairing the FIFO poll with the under-lock + // taskIndex keeps the (taskIndex, task) binding correct across concurrent workers. + // createAndSetRunnerContext already binds the context on this thread. + setCurrentKey(key); + ActionTask task = stateManager.pollNextActionTask(); Review Comment: This direct call to `stateManager.pollNextActionTask()` bypasses the operator's existing dequeue helper, which calls `builtInMetrics.markActionTaskDequeued()`. Enqueueing an action still stores its execution ID and timestamp in `initialTaskEnqueueNanos`. The skipped dequeue notification is the only place that removes this entry; action completion does not clean it up. Consequently, the parallel path retains an entry for every completed action indefinitely, causing unbounded memory growth as well as an incorrect pending-action gauge. Using the actual operator harness, 10 fully completed input records producing 20 actions left **20 retained entries and a pending count of 20** in parallel mode; serial mode left **zero**. Could we preserve the existing dequeue bookkeeping here, accounting for internal noop tasks? The direct `pollNextPendingInputEvent()` call in `maybeFinishCurrentInput()` also skips the corresponding pending-input metric decrement. -- 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]
