You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@flink.apache.org by "codenohup (via GitHub)" <gi...@apache.org> on 2025/07/24 11:49:25 UTC

[PR] [WIP][API][Runtime] Support agent and action asynchronous execution [flink-agents]

codenohup opened a new pull request, #80:
URL: https://github.com/apache/flink-agents/pull/80

   <!--
   * Thank you very much for contributing to Flink Agents.
   * Please add the relevant components in the PR title. E.g., [api], [runtime], [java], [python], [hotfix], etc.
   -->
   
   <!-- Please link the PR to the relevant issue(s). Hotfix doesn't need this. -->
   Linked issue: #78 
   
   ### Purpose of change
   Enable the agent and action to support asynchronous and parallel processing for improved performance. Details can be found in the #78 .
   
   This feature is still under development and is being pushed now only to run CI. 
   **Please do not review it at this time.**
   
   ### Tests
   Added corresponding unit tests and e2e tests.
   
   ### API
   
   Yes, introduce runner_context#execute_async API
   
   ### Documentation
   
   No, I'll add the relevant documentation in the future.
   


-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong merged PR #80:
URL: https://github.com/apache/flink-agents/pull/80


-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2239278569


##########
python/flink_agents/runtime/local_runner.py:
##########
@@ -115,6 +115,23 @@ def get_short_term_memory(self) -> MemoryObject:
         """
         return self._short_term_memory
 
+    def execute_async(
+        self,
+        func: Callable[[Any], Any],
+        *args: Tuple[Any, ...],
+        **kwargs: Dict[str, Any],
+    ) -> Any:
+        """Asynchronously execute the provided function. Access to memory
+         is prohibited within the function.
+        """
+        logger.info(
+            "Local runner does not support asynchronous execution; falling back to synchronous execution."
+        )
+        func_result = func(*args, **kwargs)
+        yield func_result

Review Comment:
   > Why do we need this yeild here?
   
   Since the user is using _yield from_ to await the result of an asynchronous function in their _action_, we should also use yield to return the result of the function.



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2247225634


##########
python/flink_agents/plan/function.py:
##########
@@ -237,6 +237,27 @@ def check_signature(self, *args: Tuple[Any, ...]) -> None:
         """Check function signature is legal or not."""
 
 
+class PythonGeneratorWrapper:
+    """
+    A temporary wrapper class for Python generators to work around a
+    known issue in PEMJA, where the generator type is incorrectly handled.
+
+    TODO: This wrapper is intended to be a temporary solution. Once PEMJA
+    version 0.6.0 (or later) fixes the bug related to generator type conversion,
+    this wrapper should be removed.
+    """

Review Comment:
   Better link to the tracking issue.



##########
python/flink_agents/plan/function.py:
##########
@@ -305,3 +326,31 @@ def get_python_function_cache_keys() -> List[Tuple[str, str]]:
         List of (module, qualname) tuples representing cached functions.
     """
     return list(_PYTHON_FUNCTION_CACHE.keys())
+
+
+def call_python_generator(generator_wrapper: PythonGeneratorWrapper) -> (bool, Any):
+    """
+    Invokes the next step of a wrapped Python generator and returns whether
+    it is done, along with the yielded or returned value.
+
+    Args:
+        generator_wrapper (PythonGeneratorWrapper): A wrapper object that
+        contains a `generator` attribute. This attribute should be an instance
+        of a Python generator.
+
+    Returns:
+        Tuple[bool, Any]:
+            - The first element is a boolean flag indicating whether the generator
+            has finished:
+                * False: The generator has more values to yield.
+                * True: The generator has completed.
+            - The second element is either:
+                * The value yielded by the generator (when not exhausted), or
+                * The return value of the generator (when it has finished).
+    """
+    try:
+        result = next(generator_wrapper.generator)
+    except StopIteration as e:
+        return True, e.value
+    else:
+        return False, result

Review Comment:
   ```suggestion
   def call_python_generator(generator_wrapper: PythonGeneratorWrapper) -> (bool, Any):
       try:
           result = next(generator_wrapper.generator)
           return False, result
       except StopIteration as e:
           return True, e.value if hasattr(e, 'value') else None
       except Exception as e:
           # log and throw error
           logger.error(f"Error in generator execution: {e}")
           raise
   ```



##########
python/flink_agents/runtime/local_runner.py:
##########
@@ -203,7 +220,13 @@ def run(self, **data: Dict[str, Any]) -> Any:
                 logger.info(
                     "key: %s, performing action: %s", key, action.name
                 )
-                action.exec(event, context)
+                func_result = action.exec(event, context)
+                if isinstance(func_result, Generator):
+                    try:
+                        while True:
+                            next(func_result)
+                    except StopIteration:
+                        pass

Review Comment:
   ```suggestion
   func_result = action.exec(event, context)
   if isinstance(func_result, Generator):
       try:
           for _ in func_result:
               pass
       except Exception as e:
           logger.error(f"Error in async execution: {e}")
           raise
   ```



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java:
##########
@@ -40,63 +36,103 @@ public class PythonActionExecutor {
             "from flink_agents.plan import function\n"
                     + "from flink_agents.runtime import flink_runner_context\n"
                     + "from flink_agents.runtime import python_java_utils";
+
+    // =========== RUNNER CONTEXT ===========
     private static final String CREATE_FLINK_RUNNER_CONTEXT =
             "flink_runner_context.create_flink_runner_context";
+    private static final String FLINK_RUNNER_CONTEXT_REF_NAME_PREFIX = "flink_runner_context_";
+    private static final AtomicLong FLINK_RUNNER_CONTEXT_REF_ID = new AtomicLong(0);
+
+    // ========== ASYNC THREAD POOL ===========
+    private static volatile Object pythonAsyncThreadPoolRef;
+    private static final String CREATE_ASYNC_THREAD_POOL =
+            "flink_runner_context.create_async_thread_pool";
+    private static final String PYTHON_ASYNC_THREAD_POOL_REF_NAME = "python_async_thread_pool";
+    private static final AtomicLong PYTHON_ASYNC_THREAD_POOL_REF_ID = new AtomicLong(0);
+
+    // =========== PYTHON GENERATOR ===========
+    private static final String CALL_PYTHON_GENERATOR = "function.call_python_generator";
+    private static final String PYTHON_GENERATOR_VAR_NAME_PREFIX = "python_generator_";
+    private static final AtomicLong PYTHON_GENERATOR_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON AND JAVA OBJECT CONVERT ===========
     private static final String CONVERT_TO_PYTHON_OBJECT =
             "python_java_utils.convert_to_python_object";
     private static final String WRAP_TO_INPUT_EVENT = "python_java_utils.wrap_to_input_event";
     private static final String GET_OUTPUT_FROM_OUTPUT_EVENT =
             "python_java_utils.get_output_from_output_event";
-    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME = "flink_runner_context";
 
     private final PythonEnvironmentManager environmentManager;
     private final String agentPlanJson;
-    private PythonRunnerContextImpl runnerContext;
-
     private PythonInterpreter interpreter;
+    private String pythonAsyncThreadPoolObjectName;
 
     public PythonActionExecutor(PythonEnvironmentManager environmentManager, String agentPlanJson) {
         this.environmentManager = environmentManager;
         this.agentPlanJson = agentPlanJson;
     }
 
-    public void open(
-            MapState<String, MemoryObjectImpl.MemoryItem> shortTermMemState,
-            FlinkAgentsMetricGroupImpl metricGroup)
-            throws Exception {
+    public void open() throws Exception {
         environmentManager.open();
         EmbeddedPythonEnvironment env = environmentManager.createEnvironment();
 
         interpreter = env.getInterpreter();
         interpreter.exec(PYTHON_IMPORTS);
 
-        runnerContext = new PythonRunnerContextImpl(shortTermMemState, metricGroup);
-
-        // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject =
-                interpreter.invoke(CREATE_FLINK_RUNNER_CONTEXT, runnerContext, agentPlanJson);
-        interpreter.set(FLINK_RUNNER_CONTEXT_VAR_NAME, pythonRunnerContextObject);
+        // TODO: remove the set and get thread pool after updating pemja to version 0.5.3

Review Comment:
   same here for the issue link



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -112,78 +149,151 @@ public void open() throws Exception {
         metricGroup = new FlinkAgentsMetricGroupImpl(getMetricGroup());
         builtInMetrics = new BuiltInMetrics(metricGroup, agentPlan);
 
-        runnerContext = new RunnerContextImpl(shortTermMemState, metricGroup);
+        // init agent processing related state
+        actionTasksKState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "actionTasks", TypeInformation.of(ActionTask.class)));
+        pendingInputEventsKState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "pendingInputEvents", TypeInformation.of(Event.class)));
+        // We use UnionList here to ensure that the task can access all keys after parallelism
+        // modifications.
+        // Subsequent steps {@link #tryResumeProcessActionTasks} will then filter out keys that do
+        // not belong to the key range of current task.
+        currentProcessingKeysOpState =
+                getOperatorStateBackend()
+                        .getUnionListState(
+                                new ListStateDescriptor<>(
+                                        "currentProcessingKeys", TypeInformation.of(Object.class)));
 
         // init PythonActionExecutor
         initPythonActionExecutor();
+
+        mailboxProcessor = getMailboxProcessor();
+
+        // Since an operator restart may change the key range it manages due to changes in
+        // parallelism,
+        // and {@link tryProcessActionTaskForKey} mails might be lost,
+        // it is necessary to reprocess all keys to ensure correctness.
+        tryResumeProcessActionTasks();
     }
 
     @Override
     public void processElement(StreamRecord<IN> record) throws Exception {
         IN input = record.getValue();
         LOG.debug("Receive an element {}", input);
 
-        // 1. wrap to InputEvent first
+        // wrap to InputEvent first
         Event inputEvent = wrapToInputEvent(input);
 
-        // 2. execute action
-        LinkedList<Event> events = new LinkedList<>();
-        events.push(inputEvent);
-        while (!events.isEmpty()) {
-            Event event = events.pop();
-            builtInMetrics.markEventProcessed();
-            List<Action> actions = getActionsTriggeredBy(event);
-            if (actions != null && !actions.isEmpty()) {
-                for (Action action : actions) {
-                    // TODO: Support multi-action execution for a single event. Example: A Java
-                    // event
-                    // should be processable by both Java and Python actions.
-                    // TODO: Implement asynchronous action execution.
-
-                    // execute action and collect output events
-                    String actionName = action.getName();
-                    LOG.debug("Try execute action {} for event {}.", actionName, event);
-                    List<Event> actionOutputEvents;
-                    if (action.getExec() instanceof JavaFunction) {
-                        runnerContext.setActionName(actionName);
-                        action.getExec().call(event, runnerContext);
-                        actionOutputEvents = runnerContext.drainEvents();
-                    } else if (action.getExec() instanceof PythonFunction) {
-                        checkState(event instanceof PythonEvent);
-                        actionOutputEvents =
-                                pythonActionExecutor.executePythonFunction(
-                                        (PythonFunction) action.getExec(),
-                                        (PythonEvent) event,
-                                        actionName);
-                    } else {
-                        throw new RuntimeException("Unsupported action type: " + action.getClass());
-                    }
-                    builtInMetrics.markActionExecuted(actionName);
-
-                    for (Event actionOutputEvent : actionOutputEvents) {
-                        if (EventUtil.isOutputEvent(actionOutputEvent)) {
-                            builtInMetrics.markEventProcessed();
-                            OUT outputData = getOutputFromOutputEvent(actionOutputEvent);
-                            LOG.debug(
-                                    "Collect output data {} for input {} in action {}.",
-                                    outputData,
-                                    input,
-                                    action.getName());
-                            output.collect(reusedStreamRecord.replace(outputData));
-                        } else {
-                            LOG.debug(
-                                    "Collect event {} for event {} in action {}.",
-                                    actionOutputEvent,
-                                    event,
-                                    action.getName());
-                            events.add(actionOutputEvent);
-                        }
-                    }
+        if (currentKeyHasMoreActionTask()) {
+            // If there are already actions being processed for the current key, the newly incoming
+            // event should be queued and processed later. Therefore, we add it to
+            // pendingInputEventsState.
+            pendingInputEventsKState.add(inputEvent);
+        } else {
+            // Otherwise, the new event is processed immediately.
+            processEvent(getCurrentKey(), inputEvent, true);
+        }
+    }
+
+    private void processEvent(Object key, Event event, boolean shouldSubmitMail) throws Exception {

Review Comment:
   In terms of readability, I find this `shouldSubmitMail` a bit hard to understand. It's not straightforward that `processEvent` sometimes need to submit the mail to the mailbox and sometimes not.
   
   I think the rules for notifying the mailbox executor about pending action tasks in a key is quite simple. We should notify in two cases:
   1. When `actionTasksKState` changes from empty to non-empty. That is when the InputEvent is processed, because it is guaranteed that `actionTasksKState` is empty before processing the next InputEvent.
   2. When we processed an action task from `actionTasksKState` and there are more remaining. 
   
   A straightforward approach would be `processEvent` always submit to the mailbox when an input event is processed, and `processActionTaskForKey` always submit to the mailbox when there are remaining action tasks. In this way, we won't need the `shouldSubmitMail` argument here. 
   
   IIUC, the problem is when we finish all action tasks of the previous input event, and pick up a new pending input event. Currently, we check the remaining action tasks at the end of `processActionTaskForKey`. If we move that check to before processing the a new pending input event, the problem should be solved.



##########
python/flink_agents/runtime/flink_runner_context.py:
##########
@@ -105,7 +110,34 @@ def get_action_metric_group(self) -> FlinkMetricGroup:
         """
         return FlinkMetricGroup(self._j_runner_context.getActionMetricGroup())
 
-
-def create_flink_runner_context(j_runner_context: Any, agent_plan_json: str) -> FlinkRunnerContext:
+    @override
+    def execute_async(
+            self,
+            func: Callable[[Any], Any],
+            *args: Tuple[Any, ...],
+            **kwargs: Dict[str, Any],
+    ) -> Any:
+        """Asynchronously execute the provided function. Access to memory
+         is prohibited within the function.
+        """
+        future = self.executor.submit(func, *args, **kwargs)
+        while not future.done():
+            # TODO: Currently, we are using a polling mechanism to check whether
+            #  the future has completed. This approach should be optimized in the
+            #  future by switching to a notification-based model, where the Flink
+            #  operator is notified directly once the future is completed.
+            yield
+        return future.result()
+
+
+def create_flink_runner_context(
+    j_runner_context: Any, agent_plan_json: str, executor: ThreadPoolExecutor
+) -> FlinkRunnerContext:
     """Used to create a FlinkRunnerContext Python object in Pemja environment."""
-    return FlinkRunnerContext(j_runner_context, agent_plan_json)
+    return FlinkRunnerContext(j_runner_context, agent_plan_json, executor)
+
+
+def create_async_thread_pool() -> ThreadPoolExecutor:
+    """Used to create a thread pool to execute asynchronous
+    code block in action."""
+    return ThreadPoolExecutor(max_workers=os.cpu_count() * 2)

Review Comment:
   How do we clean-up the thread pool?



##########
plan/src/main/java/org/apache/flink/agents/plan/JavaFunction.java:
##########
@@ -71,7 +71,10 @@ public Class<?>[] getParameterTypes() {
         return parameterTypes;
     }
 
-    public Method getMethod() {
+    public Method getMethod() throws ClassNotFoundException, NoSuchMethodException {
+        if (method == null) {
+            this.method = Class.forName(qualName).getMethod(methodName, parameterTypes);
+        }

Review Comment:
   Why do we need this lazy initialization?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * 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;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.plan.Action;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+/**
+ * This class represents a task related to the execution of an action in {@link
+ * ActionExecutionOperator}.
+ *
+ * <p>An action is split into multiple code blocks, and each code block is represented by an {@code
+ * ActionTask}. You can call {@link #invoke()} to execute a code block and obtain invoke result
+ * {@link ActionTaskResult}. If the action contains additional code blocks, you can obtain the next
+ * {@code ActionTask} via {@link ActionTaskResult#getGeneratedActionTask()} and continue executing
+ * it.
+ */
+public abstract class ActionTask {
+
+    protected static final Logger LOG = LoggerFactory.getLogger(ActionTask.class);
+
+    protected final Object key;
+    protected final Event event;
+    protected final Action action;
+    /**
+     * Since RunnerContextImpl contains references to the Operator and state, it should not be
+     * serialized and included in the state with ActionTask. Instead, we should check if a valid
+     * RunnerContext exists before each ActionTask invocation and create a new one if necessary.
+     */
+    protected transient RunnerContextImpl runnerContext;
+
+    public ActionTask(Object key, Event event, Action action) {
+        this.key = key;
+        this.event = event;
+        this.action = action;
+    }
+
+    public RunnerContextImpl getRunnerContext() {
+        return runnerContext;
+    }
+
+    public void setRunnerContext(RunnerContextImpl runnerContext) {
+        this.runnerContext = runnerContext;
+    }
+
+    public Object getKey() {
+        return key;
+    }
+
+    /** Invokes the action task. */
+    public abstract ActionTaskResult invoke() throws Exception;
+
+    public class ActionTaskResult {
+        private final boolean finished;
+        private final List<Event> outputEvents;
+        private final ActionTask generatedActionTask;
+
+        public ActionTaskResult(
+                boolean finished, List<Event> outputEvents, ActionTask generatedActionTask) {
+            this.finished = finished;
+            this.outputEvents = outputEvents;
+            this.generatedActionTask = generatedActionTask;
+        }
+
+        public boolean isFinished() {
+            return finished;
+        }
+
+        public List<Event> getOutputEvents() {
+            return outputEvents;
+        }
+
+        public ActionTask getGeneratedActionTask() {

Review Comment:
   Maybe return an `Optional<ActionTask>` and force the checking and null-handling at the caller side?



##########
python/flink_agents/runtime/local_runner.py:
##########
@@ -128,6 +128,23 @@ def get_action_metric_group(self) -> MetricGroup:
         err_msg = "Metric mechanism is not supported for local agent execution yet."
         raise NotImplementedError(err_msg)
 
+    def execute_async(
+        self,
+        func: Callable[[Any], Any],
+        *args: Tuple[Any, ...],
+        **kwargs: Dict[str, Any],
+    ) -> Any:
+        """Asynchronously execute the provided function. Access to memory
+         is prohibited within the function.
+        """
+        logger.warning(
+            "Local runner does not support asynchronous execution; falling back to synchronous execution."
+        )
+        func_result = func(*args, **kwargs)
+        yield func_result
+        return func_result

Review Comment:
   Should not need this `return` statement.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -240,4 +362,70 @@ private List<Action> getActionsTriggeredBy(Event event) {
             return agentPlan.getActionsTriggeredBy(event.getClass().getName());
         }
     }
+
+    private MailboxProcessor getMailboxProcessor() throws Exception {
+        Field field = MailboxExecutorImpl.class.getDeclaredField("mailboxProcessor");
+        field.setAccessible(true);
+        return (MailboxProcessor) field.get(mailboxExecutor);
+    }
+
+    private void checkMailboxThread() {
+        checkState(
+                mailboxProcessor.isMailboxThread(),
+                "Expected to be running on the task mailbox thread, but was not.");
+    }
+
+    private ActionTask createActionTask(Object key, Action action, Event event) {
+        if (action.getExec() instanceof JavaFunction) {
+            return new JavaActionTask(key, event, action);
+        } else if (action.getExec() instanceof PythonFunction) {
+            return new PythonActionTask(key, event, action, pythonActionExecutor);
+        } else {
+            throw new IllegalStateException(
+                    "Unsupported action type: " + action.getExec().getClass());
+        }
+    }
+
+    private void createAndSetRunnerContext(ActionTask actionTask) {
+        if (actionTask.getRunnerContext() != null) {
+            return;
+        }
+
+        RunnerContextImpl runnerContext;
+        if (actionTask.action.getExec() instanceof JavaFunction) {
+            runnerContext =
+                    new RunnerContextImpl(shortTermMemState, metricGroup, this::checkMailboxThread);
+        } else if (actionTask.action.getExec() instanceof PythonFunction) {
+            runnerContext =
+                    new PythonRunnerContextImpl(
+                            shortTermMemState, metricGroup, this::checkMailboxThread);
+        } else {
+            throw new IllegalStateException(
+                    "Unsupported action type: " + actionTask.action.getExec().getClass());
+        }
+
+        runnerContext.setActionName(actionTask.action.getName());
+        actionTask.setRunnerContext(runnerContext);
+    }
+
+    private boolean currentKeyHasMoreActionTask() throws Exception {
+        return listStateNotEmpty(actionTasksKState);
+    }
+
+    private void tryResumeProcessActionTasks() throws Exception {
+        Iterable<Object> keys = currentProcessingKeysOpState.get();
+        if (keys != null) {
+            for (Object key : keys) {
+                mailboxExecutor.submit(
+                        () -> tryProcessActionTaskForKey(key), "process action task");
+            }
+        }
+    }

Review Comment:
   This is smart!



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2240270394


##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java:
##########
@@ -39,62 +37,103 @@ public class PythonActionExecutor {
             "from flink_agents.plan import function\n"
                     + "from flink_agents.runtime import flink_runner_context\n"
                     + "from flink_agents.runtime import python_java_utils";
+
+    // =========== RUNNER CONTEXT ===========
     private static final String CREATE_FLINK_RUNNER_CONTEXT =
             "flink_runner_context.create_flink_runner_context";
+    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME_PREFIX = "flink_runner_context_";
+    private static final AtomicLong FLINK_RUNNER_CONTEXT_VAR_ID = new AtomicLong(0);
+
+    // ========== ASYNC THREAD POOL ===========
+    private static final String CREATE_ASYNC_THREAD_POOL =
+            "flink_runner_context.create_async_thread_pool";
+    private static final String PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX =
+            "python_async_thread_pool";
+    private static final AtomicLong PYTHON_ASYNC_THREAD_POOL_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON GENERATOR ===========
+    private static final String CALL_PYTHON_GENERATOR = "function.call_python_generator";
+    private static final String PYTHON_GENERATOR_VAR_NAME_PREFIX = "python_generator_";
+    private static final AtomicLong PYTHON_GENERATOR_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON AND JAVA OBJECT CONVERT ===========
     private static final String CONVERT_TO_PYTHON_OBJECT =
             "python_java_utils.convert_to_python_object";
     private static final String WRAP_TO_INPUT_EVENT = "python_java_utils.wrap_to_input_event";
     private static final String GET_OUTPUT_FROM_OUTPUT_EVENT =
             "python_java_utils.get_output_from_output_event";
-    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME = "flink_runner_context";
 
     private final PythonEnvironmentManager environmentManager;
     private final String agentPlanJson;
-    private PythonRunnerContextImpl runnerContext;
-
     private PythonInterpreter interpreter;
+    private String pythonAsyncThreadPoolObjectName;
 
     public PythonActionExecutor(PythonEnvironmentManager environmentManager, String agentPlanJson) {
         this.environmentManager = environmentManager;
         this.agentPlanJson = agentPlanJson;
     }
 
-    public void open(
-            MapState<String, MemoryObjectImpl.MemoryItem> shortTermMemState,
-            Runnable mailboxThreadChecker)
-            throws Exception {
+    public void open() throws Exception {
         environmentManager.open();
         EmbeddedPythonEnvironment env = environmentManager.createEnvironment();
 
         interpreter = env.getInterpreter();
         interpreter.exec(PYTHON_IMPORTS);
 
-        runnerContext = new PythonRunnerContextImpl(shortTermMemState, mailboxThreadChecker);
-
-        // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject =
-                interpreter.invoke(CREATE_FLINK_RUNNER_CONTEXT, runnerContext, agentPlanJson);
-        interpreter.set(FLINK_RUNNER_CONTEXT_VAR_NAME, pythonRunnerContextObject);
+        // TODO: remove the set and get thread pool after updating pemja to version 0.5.3
+        Object pythonAsyncThreadPool = interpreter.invoke(CREATE_ASYNC_THREAD_POOL);
+        this.pythonAsyncThreadPoolObjectName =
+                PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX
+                        + PYTHON_ASYNC_THREAD_POOL_VAR_ID.incrementAndGet();
+        interpreter.set(pythonAsyncThreadPoolObjectName, pythonAsyncThreadPool);
     }
 
-    public List<Event> executePythonFunction(PythonFunction function, PythonEvent event)
+    /**
+     * Execute the Python function, which may return a Python generator that needs to be processed
+     * in the future. Due to an issue in Pemja regarding incorrect object reference counting, this
+     * may lead to garbage collection of the object. To prevent this, we use the set and get methods
+     * to manually increment the object's reference count, then return the name of the Python
+     * generator variable.
+     *
+     * @return The name of the Python generator variable. It may be null if the Python function does
+     *     not return a generator.
+     */
+    public String executePythonFunction(
+            PythonFunction function, PythonEvent event, RunnerContextImpl runnerContext)
             throws Exception {
         runnerContext.checkNoPendingEvents();
         function.setInterpreter(interpreter);
 
         // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject = interpreter.get(FLINK_RUNNER_CONTEXT_VAR_NAME);
+        Object pythonRunnerContextObject =
+                interpreter.invoke(
+                        CREATE_FLINK_RUNNER_CONTEXT,
+                        runnerContext,
+                        agentPlanJson,
+                        interpreter.get(pythonAsyncThreadPoolObjectName));
+        String pythonRunnerContextObjectName =
+                FLINK_RUNNER_CONTEXT_VAR_NAME_PREFIX
+                        + FLINK_RUNNER_CONTEXT_VAR_ID.incrementAndGet();

Review Comment:
   > Why do we need to create a new object for each time calling the function?
   
   Since actions may be executed in parallel, to avoid interference between different action instances, we create a new instance for each run of an action.
   
   



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "wenjin272 (via GitHub)" <gi...@apache.org>.
wenjin272 commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2232367244


##########
python/flink_agents/plan/function.py:
##########
@@ -146,7 +146,59 @@ def check_signature(self, *args: Tuple[Any, ...]) -> None:
         """Check function signature is legal or not."""
 
 
+class PythonGeneratorWrapper:
+    """
+    A temporary wrapper class for Python generators to work around a
+    known issue in PEMJA, where the generator type is incorrectly handled.
+
+    Note: This wrapper is intended to be a temporary solution. Once PEMJA
+    version 0.6.0 (or later) fixes the bug related to generator type conversion,
+    this wrapper should be removed.
+    """

Review Comment:
   Maybe need add TODO and create correspond issue to track 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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2250498638


##########
plan/src/main/java/org/apache/flink/agents/plan/JavaFunction.java:
##########
@@ -71,7 +71,10 @@ public Class<?>[] getParameterTypes() {
         return parameterTypes;
     }
 
-    public Method getMethod() {
+    public Method getMethod() throws ClassNotFoundException, NoSuchMethodException {
+        if (method == null) {
+            this.method = Class.forName(qualName).getMethod(methodName, parameterTypes);
+        }

Review Comment:
   > Why do we need this lazy initialization?
   
   Since we need to store the `ActionTask` in the state, and it indirectly contains the `method` from the `JavaFunction`, which is not serializable in Java.



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2250795038


##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -112,75 +150,167 @@ public void open() throws Exception {
         metricGroup = new FlinkAgentsMetricGroupImpl(getMetricGroup());
         builtInMetrics = new BuiltInMetrics(metricGroup, agentPlan);
 
-        runnerContext = new RunnerContextImpl(shortTermMemState, metricGroup);
+        // init agent processing related state
+        actionTasksKState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "actionTasks", TypeInformation.of(ActionTask.class)));
+        pendingInputEventsKState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "pendingInputEvents", TypeInformation.of(Event.class)));
+        // We use UnionList here to ensure that the task can access all keys after parallelism
+        // modifications.
+        // Subsequent steps {@link #tryResumeProcessActionTasks} will then filter out keys that do
+        // not belong to the key range of current task.
+        currentProcessingKeysOpState =
+                getOperatorStateBackend()
+                        .getUnionListState(
+                                new ListStateDescriptor<>(
+                                        "currentProcessingKeys", TypeInformation.of(Object.class)));
 
         // init PythonActionExecutor
         initPythonActionExecutor();
+
+        mailboxProcessor = getMailboxProcessor();
+
+        // Since an operator restart may change the key range it manages due to changes in
+        // parallelism,
+        // and {@link tryProcessActionTaskForKey} mails might be lost,
+        // it is necessary to reprocess all keys to ensure correctness.
+        tryResumeProcessActionTasks();
     }
 
     @Override
     public void processElement(StreamRecord<IN> record) throws Exception {
         IN input = record.getValue();
         LOG.debug("Receive an element {}", input);
 
-        // 1. wrap to InputEvent first
+        // wrap to InputEvent first
         Event inputEvent = wrapToInputEvent(input);
 
-        // 2. execute action
-        LinkedList<Event> events = new LinkedList<>();
-        events.push(inputEvent);
-        while (!events.isEmpty()) {
-            Event event = events.pop();
-            builtInMetrics.markEventProcessed();
-            List<Action> actions = getActionsTriggeredBy(event);
-            if (actions != null && !actions.isEmpty()) {
-                for (Action action : actions) {
-                    // TODO: Support multi-action execution for a single event. Example: A Java
-                    // event
-                    // should be processable by both Java and Python actions.
-                    // TODO: Implement asynchronous action execution.
-
-                    // execute action and collect output events
-                    String actionName = action.getName();
-                    LOG.debug("Try execute action {} for event {}.", actionName, event);
-                    List<Event> actionOutputEvents;
-                    if (action.getExec() instanceof JavaFunction) {
-                        runnerContext.setActionName(actionName);
-                        action.getExec().call(event, runnerContext);
-                        actionOutputEvents = runnerContext.drainEvents();
-                    } else if (action.getExec() instanceof PythonFunction) {
-                        checkState(event instanceof PythonEvent);
-                        actionOutputEvents =
-                                pythonActionExecutor.executePythonFunction(
-                                        (PythonFunction) action.getExec(),
-                                        (PythonEvent) event,
-                                        actionName);
-                    } else {
-                        throw new RuntimeException("Unsupported action type: " + action.getClass());
-                    }
-                    builtInMetrics.markActionExecuted(actionName);
-
-                    for (Event actionOutputEvent : actionOutputEvents) {
-                        if (EventUtil.isOutputEvent(actionOutputEvent)) {
-                            builtInMetrics.markEventProcessed();
-                            OUT outputData = getOutputFromOutputEvent(actionOutputEvent);
-                            LOG.debug(
-                                    "Collect output data {} for input {} in action {}.",
-                                    outputData,
-                                    input,
-                                    action.getName());
-                            output.collect(reusedStreamRecord.replace(outputData));
-                        } else {
-                            LOG.debug(
-                                    "Collect event {} for event {} in action {}.",
-                                    actionOutputEvent,
-                                    event,
-                                    action.getName());
-                            events.add(actionOutputEvent);
-                        }
-                    }
+        if (currentKeyHasMoreActionTask()) {
+            // If there are already actions being processed for the current key, the newly incoming
+            // event should be queued and processed later. Therefore, we add it to
+            // pendingInputEventsState.
+            pendingInputEventsKState.add(inputEvent);
+        } else {
+            // Otherwise, the new event is processed immediately.
+            Object key = getCurrentKey();
+            boolean shouldSubmitMail = processEvent(key, inputEvent);
+            if (shouldSubmitMail) {
+                mailboxExecutor.submit(
+                        () -> tryProcessActionTaskForKey(key), "process action task");
+            }
+        }
+    }
+
+    /**
+     * Processes an incoming event for the given key and determines whether a new mail
+     * `tryProcessActionTaskForKey` should be submitted to continue processing.
+     *
+     * <p>The return value `shouldSubmitMail` is set to true only if the action tasks list
+     * (actionTasksKState) was previously empty and is now non-empty.
+     *
+     * @param key The key associated with the event.
+     * @param event The event to process.
+     * @return True if a new mail should be submitted, false otherwise.
+     * @throws Exception If any error occurs during processing.
+     */
+    private boolean processEvent(Object key, Event event) throws Exception {
+        boolean shouldSubmitMail = false;
+        builtInMetrics.markEventProcessed();
+        if (EventUtil.isOutputEvent(event)) {
+            // If the event is an OutputEvent, we send it downstream.
+            OUT outputData = getOutputFromOutputEvent(event);
+            output.collect(reusedStreamRecord.replace(outputData));
+        } else {
+            if (EventUtil.isInputEvent(event)) {
+                // If the event is an InputEvent, we mark that the key is currently being processed.
+                currentProcessingKeysOpState.add(key);
+            }
+            // We then obtain the triggered action and add ActionTasks to the waiting processing
+            // queue.
+            List<Action> triggerActions = getActionsTriggeredBy(event);
+            if (triggerActions != null && !triggerActions.isEmpty()) {
+                shouldSubmitMail = !listStateNotEmpty(actionTasksKState);
+                for (Action triggerAction : triggerActions) {
+                    actionTasksKState.add(createActionTask(key, triggerAction, event));
+                }
+            }
+        }
+
+        return shouldSubmitMail;

Review Comment:
   Why not submit the mail here?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTask.java:
##########
@@ -0,0 +1,112 @@
+/*
+ * 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;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.plan.Action;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+import java.util.Optional;
+
+/**
+ * This class represents a task related to the execution of an action in {@link
+ * ActionExecutionOperator}.
+ *
+ * <p>An action is split into multiple code blocks, and each code block is represented by an {@code
+ * ActionTask}. You can call {@link #invoke()} to execute a code block and obtain invoke result
+ * {@link ActionTaskResult}. If the action contains additional code blocks, you can obtain the next
+ * {@code ActionTask} via {@link ActionTaskResult#getGeneratedActionTask()} and continue executing
+ * it.
+ */
+public abstract class ActionTask {
+
+    protected static final Logger LOG = LoggerFactory.getLogger(ActionTask.class);
+
+    protected final Object key;
+    protected final Event event;
+    protected final Action action;
+    /**
+     * Since RunnerContextImpl contains references to the Operator and state, it should not be
+     * serialized and included in the state with ActionTask. Instead, we should check if a valid
+     * RunnerContext exists before each ActionTask invocation and create a new one if necessary.
+     */
+    protected transient RunnerContextImpl runnerContext;
+
+    public ActionTask(Object key, Event event, Action action) {
+        this.key = key;
+        this.event = event;
+        this.action = action;
+    }
+
+    public RunnerContextImpl getRunnerContext() {
+        return runnerContext;
+    }
+
+    public void setRunnerContext(RunnerContextImpl runnerContext) {
+        this.runnerContext = runnerContext;
+    }
+
+    public Object getKey() {
+        return key;
+    }
+
+    /** Invokes the action task. */
+    public abstract ActionTaskResult invoke() throws Exception;
+
+    public class ActionTaskResult {
+        private final boolean finished;
+        private final List<Event> outputEvents;
+        private final Optional<ActionTask> generatedActionTaskOpt;
+
+        public ActionTaskResult(
+                boolean finished,
+                List<Event> outputEvents,
+                Optional<ActionTask> generatedActionTaskOpt) {
+            this.finished = finished;
+            this.outputEvents = outputEvents;
+            this.generatedActionTaskOpt = generatedActionTaskOpt;
+        }

Review Comment:
   ```suggestion
           public ActionTaskResult(
                   boolean finished,
                   List<Event> outputEvents,
                   @Nullable ActionTask generatedActionTaskOpt) {
               this.finished = finished;
               this.outputEvents = outputEvents;
               this.generatedActionTaskOpt = Optional.ofNullable(generatedActionTaskOpt);
           }
   ```



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2250359039


##########
python/flink_agents/runtime/local_runner.py:
##########
@@ -128,6 +128,23 @@ def get_action_metric_group(self) -> MetricGroup:
         err_msg = "Metric mechanism is not supported for local agent execution yet."
         raise NotImplementedError(err_msg)
 
+    def execute_async(
+        self,
+        func: Callable[[Any], Any],
+        *args: Tuple[Any, ...],
+        **kwargs: Dict[str, Any],
+    ) -> Any:
+        """Asynchronously execute the provided function. Access to memory
+         is prohibited within the function.
+        """
+        logger.warning(
+            "Local runner does not support asynchronous execution; falling back to synchronous execution."
+        )
+        func_result = func(*args, **kwargs)
+        yield func_result
+        return func_result

Review Comment:
   > Should not need this `return` statement.
   
   Since the action needs to obtain the result of the asynchronous function via `yield from runner_context.execute_async(...)`, it is necessary to return the value here. It should be modified as follows:
   
   ```python
   func_result = func(*args, **kwargs)
   yield
   return func_result
   ```
   



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on PR #80:
URL: https://github.com/apache/flink-agents/pull/80#issuecomment-3139013008

   @xintongsong 
   Thanks for your review! I've addressed all the comments in the PR and rebased onto master. Could you PTAL when you have a moment?


-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "xintongsong (via GitHub)" <gi...@apache.org>.
xintongsong commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2238354910


##########
python/flink_agents/runtime/local_runner.py:
##########
@@ -115,6 +115,23 @@ def get_short_term_memory(self) -> MemoryObject:
         """
         return self._short_term_memory
 
+    def execute_async(
+        self,
+        func: Callable[[Any], Any],
+        *args: Tuple[Any, ...],
+        **kwargs: Dict[str, Any],
+    ) -> Any:
+        """Asynchronously execute the provided function. Access to memory
+         is prohibited within the function.
+        """
+        logger.info(
+            "Local runner does not support asynchronous execution; falling back to synchronous execution."
+        )
+        func_result = func(*args, **kwargs)
+        yield func_result

Review Comment:
   Why do we need this yeild here?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -240,6 +308,90 @@ private List<Action> getActionsTriggeredBy(Event event) {
         }
     }
 
+    private <T> T pollFromListState(ListState<T> listState) throws Exception {
+        Iterator<T> listStateIterator = listState.get().iterator();
+        if (!listStateIterator.hasNext()) {
+            return null;
+        }
+
+        T polled = listStateIterator.next();
+        List<T> remaining = new ArrayList<>();
+        while (listStateIterator.hasNext()) {
+            remaining.add(listStateIterator.next());
+        }
+        listState.clear();
+        listState.update(remaining);
+        return polled;
+    }
+
+    private ActionTask createActionTask(Object key, Action action, Event event) {
+        if (action.getExec() instanceof JavaFunction) {
+            return new JavaActionTask(
+                    key,
+                    event,
+                    action,
+                    new RunnerContextImpl(shortTermMemState, this::checkMailboxThread));
+        } else if (action.getExec() instanceof PythonFunction) {
+            return new PythonActionTask(
+                    key,
+                    event,
+                    action,
+                    pythonActionExecutor,
+                    new PythonRunnerContextImpl(shortTermMemState, this::checkMailboxThread));
+        } else {
+            throw new IllegalStateException(
+                    "Unsupported action type: " + action.getExec().getClass());
+        }
+    }
+
+    private void addActionTask(ActionTask actionTask) throws Exception {
+        actionTasksState.add(actionTask);
+        mailboxExecutor.submit(
+                () -> processActionTaskWrapper(actionTask.getKey()), "process action task");

Review Comment:
   processActionTaskWrapper -> tryProcessActionFromKey



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -116,7 +137,23 @@ public void open() throws Exception {
                         TypeInformation.of(String.class),
                         TypeInformation.of(MemoryObjectImpl.MemoryItem.class));
         shortTermMemState = getRuntimeContext().getMapState(shortTermMemStateDescriptor);
-        runnerContext = new RunnerContextImpl(shortTermMemState, this::checkMailboxThread);
+
+        // init agent processing related state
+        actionTasksState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "actionTasks", TypeInformation.of(ActionTask.class)));
+        pendingInputEventsState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "pendingInputEvents", TypeInformation.of(Event.class)));
+        currentProcessingKeysOperatorState =
+                getOperatorStateBackend()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "currentProcessingKeys", TypeInformation.of(Object.class)));

Review Comment:
   1. The fact, that `actionTasksState` and `pendingInputEventsState` are keyed states and `currentProcessingKeysOperatorState` is operator states, can be more explicit.
   - `actionTasksKState`
   - `pendingInputEventsKState`
   - `currentProcessingKeysOpState`



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -223,4 +239,16 @@ private List<Action> getActionsTriggeredBy(Event event) {
             return agentPlan.getActionsTriggeredBy(event.getClass().getName());
         }
     }
+
+    private MailboxProcessor getMailboxProcessor() throws Exception {
+        Field field = MailboxExecutorImpl.class.getDeclaredField("mailboxProcessor");
+        field.setAccessible(true);
+        return (MailboxProcessor) field.get(mailboxExecutor);
+    }

Review Comment:
   I guess the elegant way is to add an interface on MailboxExecutor for getting the mailbox main thread, or verifying whether we are in the main thread.
   
   I'm okay with the workaround here, but let's add a TODO to fix this latter.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -126,60 +163,91 @@ public void open() throws Exception {
 
     @Override
     public void processElement(StreamRecord<IN> record) throws Exception {
+        checkError();
+
         IN input = record.getValue();
         LOG.debug("Receive an element {}", input);
 
-        // 1. wrap to InputEvent first
+        // wrap to InputEvent first
         Event inputEvent = wrapToInputEvent(input);
 
-        // 2. execute action
-        LinkedList<Event> events = new LinkedList<>();
-        events.push(inputEvent);
-        while (!events.isEmpty()) {
-            Event event = events.pop();
-            List<Action> actions = getActionsTriggeredBy(event);
-            if (actions != null && !actions.isEmpty()) {
-                for (Action action : actions) {
-                    // TODO: Support multi-action execution for a single event. Example: A Java
-                    // event
-                    // should be processable by both Java and Python actions.
-                    // TODO: Implement asynchronous action execution.
-
-                    // execute action and collect output events
-                    LOG.debug("Try execute action {} for event {}.", action.getName(), event);
-                    List<Event> actionOutputEvents;
-                    if (action.getExec() instanceof JavaFunction) {
-                        action.getExec().call(event, runnerContext);
-                        actionOutputEvents = runnerContext.drainEvents();
-                    } else if (action.getExec() instanceof PythonFunction) {
-                        checkState(event instanceof PythonEvent);
-                        actionOutputEvents =
-                                pythonActionExecutor.executePythonFunction(
-                                        (PythonFunction) action.getExec(), (PythonEvent) event);
-                    } else {
-                        throw new RuntimeException("Unsupported action type: " + action.getClass());
-                    }
-
-                    for (Event actionOutputEvent : actionOutputEvents) {
-                        if (EventUtil.isOutputEvent(actionOutputEvent)) {
-                            OUT outputData = getOutputFromOutputEvent(actionOutputEvent);
-                            LOG.debug(
-                                    "Collect output data {} for input {} in action {}.",
-                                    outputData,
-                                    input,
-                                    action.getName());
-                            output.collect(reusedStreamRecord.replace(outputData));
-                        } else {
-                            LOG.debug(
-                                    "Collect event {} for event {} in action {}.",
-                                    actionOutputEvent,
-                                    event,
-                                    action.getName());
-                            events.add(actionOutputEvent);
-                        }
-                    }
+        if (currentKeyHasMoreActionTask()) {
+            // If there are already actions being processed for the current key, the newly incoming
+            // event should be queued and processed later. Therefore, we add it to
+            // pendingInputEventsState.
+            pendingInputEventsState.add(inputEvent);
+        } else {
+            // Otherwise, the new event is processed immediately.
+            processEvent(getCurrentKey(), inputEvent);
+        }
+    }
+
+    private void processEvent(Object key, Event event) throws Exception {
+        if (EventUtil.isOutputEvent(event)) {
+            // If the event is an OutputEvent, we send it downstream.
+            OUT outputData = getOutputFromOutputEvent(event);
+            output.collect(reusedStreamRecord.replace(outputData));
+        } else {
+            if (EventUtil.isInputEvent(event)) {
+                // If the event is an InputEvent, we mark that the key is currently being processed.
+                currentProcessingKeysOperatorState.add(key);
+            }
+            // We then obtain the triggered action and add ActionTasks to the waiting processing
+            // queue.
+            List<Action> triggerActions = getActionsTriggeredBy(event);
+            for (Action triggerAction : triggerActions) {
+                addActionTask(createActionTask(key, triggerAction, event));
+            }
+        }
+    }
+
+    private void processActionTask(Object key) throws Exception {
+        setCurrentKey(key);
+        ActionTask actionTask = pollFromListState(actionTasksState);
+        if (actionTask == null) {
+            throw new RuntimeException(
+                    "A null value was encountered while trying to process the ActionTask. This likely indicates a bug in the internal logic.");
+        }
+        boolean finished = actionTask.invoke();
+        List<Event> outputEvents = actionTask.getOutputEvents();

Review Comment:
   It seems `getOutputEvents()` and `getGeneratedActionTask()` are only expected to be called after `invoke()`. Why not combine these three into a single method?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -251,4 +403,17 @@ private void checkMailboxThread() {
                 mailboxProcessor.isMailboxThread(),
                 "Expected to be running on the task mailbox thread, but was not.");
     }
+
+    private void checkError() {
+        if (errorCause != null) {

Review Comment:
   1. Is `errorCause` always accessed in the mailbox thread? Do we need to consider the thread safety?
   2. Is it guaranteed there's only one error?
   3. In which case will the error not be handled immediately?



##########
python/flink_agents/runtime/local_runner.py:
##########
@@ -115,6 +115,23 @@ def get_short_term_memory(self) -> MemoryObject:
         """
         return self._short_term_memory
 
+    def execute_async(
+        self,
+        func: Callable[[Any], Any],
+        *args: Tuple[Any, ...],
+        **kwargs: Dict[str, Any],
+    ) -> Any:
+        """Asynchronously execute the provided function. Access to memory
+         is prohibited within the function.
+        """
+        logger.info(

Review Comment:
   warn



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java:
##########
@@ -39,62 +37,103 @@ public class PythonActionExecutor {
             "from flink_agents.plan import function\n"
                     + "from flink_agents.runtime import flink_runner_context\n"
                     + "from flink_agents.runtime import python_java_utils";
+
+    // =========== RUNNER CONTEXT ===========
     private static final String CREATE_FLINK_RUNNER_CONTEXT =
             "flink_runner_context.create_flink_runner_context";
+    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME_PREFIX = "flink_runner_context_";
+    private static final AtomicLong FLINK_RUNNER_CONTEXT_VAR_ID = new AtomicLong(0);
+
+    // ========== ASYNC THREAD POOL ===========
+    private static final String CREATE_ASYNC_THREAD_POOL =
+            "flink_runner_context.create_async_thread_pool";
+    private static final String PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX =
+            "python_async_thread_pool";
+    private static final AtomicLong PYTHON_ASYNC_THREAD_POOL_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON GENERATOR ===========
+    private static final String CALL_PYTHON_GENERATOR = "function.call_python_generator";
+    private static final String PYTHON_GENERATOR_VAR_NAME_PREFIX = "python_generator_";
+    private static final AtomicLong PYTHON_GENERATOR_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON AND JAVA OBJECT CONVERT ===========
     private static final String CONVERT_TO_PYTHON_OBJECT =
             "python_java_utils.convert_to_python_object";
     private static final String WRAP_TO_INPUT_EVENT = "python_java_utils.wrap_to_input_event";
     private static final String GET_OUTPUT_FROM_OUTPUT_EVENT =
             "python_java_utils.get_output_from_output_event";
-    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME = "flink_runner_context";
 
     private final PythonEnvironmentManager environmentManager;
     private final String agentPlanJson;
-    private PythonRunnerContextImpl runnerContext;
-
     private PythonInterpreter interpreter;
+    private String pythonAsyncThreadPoolObjectName;
 
     public PythonActionExecutor(PythonEnvironmentManager environmentManager, String agentPlanJson) {
         this.environmentManager = environmentManager;
         this.agentPlanJson = agentPlanJson;
     }
 
-    public void open(
-            MapState<String, MemoryObjectImpl.MemoryItem> shortTermMemState,
-            Runnable mailboxThreadChecker)
-            throws Exception {
+    public void open() throws Exception {
         environmentManager.open();
         EmbeddedPythonEnvironment env = environmentManager.createEnvironment();
 
         interpreter = env.getInterpreter();
         interpreter.exec(PYTHON_IMPORTS);
 
-        runnerContext = new PythonRunnerContextImpl(shortTermMemState, mailboxThreadChecker);
-
-        // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject =
-                interpreter.invoke(CREATE_FLINK_RUNNER_CONTEXT, runnerContext, agentPlanJson);
-        interpreter.set(FLINK_RUNNER_CONTEXT_VAR_NAME, pythonRunnerContextObject);
+        // TODO: remove the set and get thread pool after updating pemja to version 0.5.3
+        Object pythonAsyncThreadPool = interpreter.invoke(CREATE_ASYNC_THREAD_POOL);
+        this.pythonAsyncThreadPoolObjectName =
+                PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX
+                        + PYTHON_ASYNC_THREAD_POOL_VAR_ID.incrementAndGet();
+        interpreter.set(pythonAsyncThreadPoolObjectName, pythonAsyncThreadPool);
     }
 
-    public List<Event> executePythonFunction(PythonFunction function, PythonEvent event)
+    /**
+     * Execute the Python function, which may return a Python generator that needs to be processed
+     * in the future. Due to an issue in Pemja regarding incorrect object reference counting, this
+     * may lead to garbage collection of the object. To prevent this, we use the set and get methods
+     * to manually increment the object's reference count, then return the name of the Python
+     * generator variable.
+     *
+     * @return The name of the Python generator variable. It may be null if the Python function does
+     *     not return a generator.
+     */
+    public String executePythonFunction(
+            PythonFunction function, PythonEvent event, RunnerContextImpl runnerContext)
             throws Exception {
         runnerContext.checkNoPendingEvents();
         function.setInterpreter(interpreter);
 
         // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject = interpreter.get(FLINK_RUNNER_CONTEXT_VAR_NAME);
+        Object pythonRunnerContextObject =
+                interpreter.invoke(
+                        CREATE_FLINK_RUNNER_CONTEXT,
+                        runnerContext,
+                        agentPlanJson,
+                        interpreter.get(pythonAsyncThreadPoolObjectName));
+        String pythonRunnerContextObjectName =
+                FLINK_RUNNER_CONTEXT_VAR_NAME_PREFIX
+                        + FLINK_RUNNER_CONTEXT_VAR_ID.incrementAndGet();

Review Comment:
   Why do we need to create a new object for each time calling the function?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -240,6 +308,90 @@ private List<Action> getActionsTriggeredBy(Event event) {
         }
     }
 
+    private <T> T pollFromListState(ListState<T> listState) throws Exception {
+        Iterator<T> listStateIterator = listState.get().iterator();
+        if (!listStateIterator.hasNext()) {
+            return null;
+        }
+
+        T polled = listStateIterator.next();
+        List<T> remaining = new ArrayList<>();
+        while (listStateIterator.hasNext()) {
+            remaining.add(listStateIterator.next());
+        }
+        listState.clear();
+        listState.update(remaining);
+        return polled;
+    }
+
+    private ActionTask createActionTask(Object key, Action action, Event event) {
+        if (action.getExec() instanceof JavaFunction) {
+            return new JavaActionTask(
+                    key,
+                    event,
+                    action,
+                    new RunnerContextImpl(shortTermMemState, this::checkMailboxThread));
+        } else if (action.getExec() instanceof PythonFunction) {
+            return new PythonActionTask(
+                    key,
+                    event,
+                    action,
+                    pythonActionExecutor,
+                    new PythonRunnerContextImpl(shortTermMemState, this::checkMailboxThread));
+        } else {
+            throw new IllegalStateException(
+                    "Unsupported action type: " + action.getExec().getClass());
+        }
+    }
+
+    private void addActionTask(ActionTask actionTask) throws Exception {
+        actionTasksState.add(actionTask);
+        mailboxExecutor.submit(
+                () -> processActionTaskWrapper(actionTask.getKey()), "process action task");

Review Comment:
   It seems we are assuming the actions in `actionTasksState` are consistent with `processActionTaskWrapper` in the mailbox queue. This seems fragile. We should control the order of actions at one place.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -116,7 +137,23 @@ public void open() throws Exception {
                         TypeInformation.of(String.class),
                         TypeInformation.of(MemoryObjectImpl.MemoryItem.class));
         shortTermMemState = getRuntimeContext().getMapState(shortTermMemStateDescriptor);
-        runnerContext = new RunnerContextImpl(shortTermMemState, this::checkMailboxThread);
+
+        // init agent processing related state
+        actionTasksState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "actionTasks", TypeInformation.of(ActionTask.class)));
+        pendingInputEventsState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "pendingInputEvents", TypeInformation.of(Event.class)));
+        currentProcessingKeysOperatorState =
+                getOperatorStateBackend()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "currentProcessingKeys", TypeInformation.of(Object.class)));

Review Comment:
   2. How do we handle the redistribution of the operator state?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.python.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.plan.Action;
+import org.apache.flink.agents.plan.PythonFunction;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.agents.runtime.python.event.PythonEvent;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * A special {@link ActionTask} designed to execute a Python action task.
+ *
+ * <p>During asynchronous execution in Python, the {@link PythonActionTask} can produce a {@link
+ * PythonGeneratorActionTask} to represent the subsequent code block when needed.
+ */
+public class PythonActionTask extends ActionTask {
+
+    protected final PythonActionExecutor pythonActionExecutor;
+
+    public PythonActionTask(
+            Object key,
+            Event event,
+            Action action,
+            PythonActionExecutor pythonActionExecutor,
+            RunnerContextImpl runnerContext) {
+        super(key, event, action, runnerContext);
+        checkState(action.getExec() instanceof PythonFunction);
+        checkState(
+                event instanceof PythonEvent,
+                "Python action only accept python event, but got " + event);
+        this.pythonActionExecutor = pythonActionExecutor;
+    }
+
+    public boolean invoke() throws Exception {
+        LOG.debug(
+                "Try execute python action {} for event {} with key {}.",
+                action.getName(),
+                event,
+                key);
+        runnerContext.checkNoPendingEvents();
+
+        String pythonGeneratorVarName =
+                pythonActionExecutor.executePythonFunction(
+                        (PythonFunction) action.getExec(), (PythonEvent) event, runnerContext);
+        if (pythonGeneratorVarName != null) {
+            // The Python action generates a generator. We need to execute it once, which will
+            // submit an asynchronous task and return whether the action has been completed.

Review Comment:
   I think the key information here, which should be explained, is that we don't want to convert actions which don't call `async_execute()` into generators.



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.python.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.plan.Action;
+import org.apache.flink.agents.plan.PythonFunction;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.agents.runtime.python.event.PythonEvent;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * A special {@link ActionTask} designed to execute a Python action task.
+ *
+ * <p>During asynchronous execution in Python, the {@link PythonActionTask} can produce a {@link
+ * PythonGeneratorActionTask} to represent the subsequent code block when needed.
+ */
+public class PythonActionTask extends ActionTask {
+
+    protected final PythonActionExecutor pythonActionExecutor;
+
+    public PythonActionTask(
+            Object key,
+            Event event,
+            Action action,
+            PythonActionExecutor pythonActionExecutor,
+            RunnerContextImpl runnerContext) {
+        super(key, event, action, runnerContext);
+        checkState(action.getExec() instanceof PythonFunction);
+        checkState(
+                event instanceof PythonEvent,
+                "Python action only accept python event, but got " + event);
+        this.pythonActionExecutor = pythonActionExecutor;
+    }
+
+    public boolean invoke() throws Exception {
+        LOG.debug(
+                "Try execute python action {} for event {} with key {}.",
+                action.getName(),
+                event,
+                key);
+        runnerContext.checkNoPendingEvents();
+
+        String pythonGeneratorVarName =

Review Comment:
   Or create a Java class for python generator?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonGeneratorActionTask.java:
##########
@@ -0,0 +1,55 @@
+/*
+ * 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.python.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.plan.Action;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+
+/** An {@link ActionTask} wrapper a Python Generator to represent a code block in Python action. */
+public class PythonGeneratorActionTask extends PythonActionTask {
+    private final String pythonGeneratorVarName;
+
+    public PythonGeneratorActionTask(
+            Object key,
+            Event event,
+            Action action,
+            PythonActionExecutor pythonActionExecutor,
+            RunnerContextImpl runnerContext,
+            String pythonGeneratorVarName) {
+        super(key, event, action, pythonActionExecutor, runnerContext);
+        this.pythonGeneratorVarName = pythonGeneratorVarName;
+    }
+
+    @Override
+    public boolean invoke() throws Exception {
+        LOG.debug(
+                "Try execute python generator action {} for event {} with key {}.",
+                action.getName(),
+                event,
+                key);
+        boolean finished = pythonActionExecutor.callPythonGenerator(pythonGeneratorVarName).f0;

Review Comment:
   `f1` is never used?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java:
##########
@@ -39,62 +37,103 @@ public class PythonActionExecutor {
             "from flink_agents.plan import function\n"
                     + "from flink_agents.runtime import flink_runner_context\n"
                     + "from flink_agents.runtime import python_java_utils";
+
+    // =========== RUNNER CONTEXT ===========
     private static final String CREATE_FLINK_RUNNER_CONTEXT =
             "flink_runner_context.create_flink_runner_context";
+    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME_PREFIX = "flink_runner_context_";
+    private static final AtomicLong FLINK_RUNNER_CONTEXT_VAR_ID = new AtomicLong(0);
+
+    // ========== ASYNC THREAD POOL ===========
+    private static final String CREATE_ASYNC_THREAD_POOL =
+            "flink_runner_context.create_async_thread_pool";
+    private static final String PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX =
+            "python_async_thread_pool";
+    private static final AtomicLong PYTHON_ASYNC_THREAD_POOL_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON GENERATOR ===========
+    private static final String CALL_PYTHON_GENERATOR = "function.call_python_generator";
+    private static final String PYTHON_GENERATOR_VAR_NAME_PREFIX = "python_generator_";
+    private static final AtomicLong PYTHON_GENERATOR_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON AND JAVA OBJECT CONVERT ===========
     private static final String CONVERT_TO_PYTHON_OBJECT =
             "python_java_utils.convert_to_python_object";
     private static final String WRAP_TO_INPUT_EVENT = "python_java_utils.wrap_to_input_event";
     private static final String GET_OUTPUT_FROM_OUTPUT_EVENT =
             "python_java_utils.get_output_from_output_event";
-    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME = "flink_runner_context";
 
     private final PythonEnvironmentManager environmentManager;
     private final String agentPlanJson;
-    private PythonRunnerContextImpl runnerContext;
-
     private PythonInterpreter interpreter;
+    private String pythonAsyncThreadPoolObjectName;
 
     public PythonActionExecutor(PythonEnvironmentManager environmentManager, String agentPlanJson) {
         this.environmentManager = environmentManager;
         this.agentPlanJson = agentPlanJson;
     }
 
-    public void open(
-            MapState<String, MemoryObjectImpl.MemoryItem> shortTermMemState,
-            Runnable mailboxThreadChecker)
-            throws Exception {
+    public void open() throws Exception {
         environmentManager.open();
         EmbeddedPythonEnvironment env = environmentManager.createEnvironment();
 
         interpreter = env.getInterpreter();
         interpreter.exec(PYTHON_IMPORTS);
 
-        runnerContext = new PythonRunnerContextImpl(shortTermMemState, mailboxThreadChecker);
-
-        // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject =
-                interpreter.invoke(CREATE_FLINK_RUNNER_CONTEXT, runnerContext, agentPlanJson);
-        interpreter.set(FLINK_RUNNER_CONTEXT_VAR_NAME, pythonRunnerContextObject);
+        // TODO: remove the set and get thread pool after updating pemja to version 0.5.3
+        Object pythonAsyncThreadPool = interpreter.invoke(CREATE_ASYNC_THREAD_POOL);
+        this.pythonAsyncThreadPoolObjectName =
+                PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX
+                        + PYTHON_ASYNC_THREAD_POOL_VAR_ID.incrementAndGet();
+        interpreter.set(pythonAsyncThreadPoolObjectName, pythonAsyncThreadPool);

Review Comment:
   1. What does this todo mean?
   2. Is it necessary to have separate thread pools for each operator?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/operator/PythonActionTask.java:
##########
@@ -0,0 +1,84 @@
+/*
+ * 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.python.operator;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.plan.Action;
+import org.apache.flink.agents.plan.PythonFunction;
+import org.apache.flink.agents.runtime.context.RunnerContextImpl;
+import org.apache.flink.agents.runtime.operator.ActionTask;
+import org.apache.flink.agents.runtime.python.event.PythonEvent;
+import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * A special {@link ActionTask} designed to execute a Python action task.
+ *
+ * <p>During asynchronous execution in Python, the {@link PythonActionTask} can produce a {@link
+ * PythonGeneratorActionTask} to represent the subsequent code block when needed.
+ */
+public class PythonActionTask extends ActionTask {
+
+    protected final PythonActionExecutor pythonActionExecutor;
+
+    public PythonActionTask(
+            Object key,
+            Event event,
+            Action action,
+            PythonActionExecutor pythonActionExecutor,
+            RunnerContextImpl runnerContext) {
+        super(key, event, action, runnerContext);
+        checkState(action.getExec() instanceof PythonFunction);
+        checkState(
+                event instanceof PythonEvent,
+                "Python action only accept python event, but got " + event);
+        this.pythonActionExecutor = pythonActionExecutor;
+    }
+
+    public boolean invoke() throws Exception {
+        LOG.debug(
+                "Try execute python action {} for event {} with key {}.",
+                action.getName(),
+                event,
+                key);
+        runnerContext.checkNoPendingEvents();
+
+        String pythonGeneratorVarName =

Review Comment:
   The name is confusing. Maybe `pythonGeneratorRef`?



##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -240,6 +308,90 @@ private List<Action> getActionsTriggeredBy(Event event) {
         }
     }
 
+    private <T> T pollFromListState(ListState<T> listState) throws Exception {

Review Comment:
   1. These list state operations seems heavy. They should be avoid as much as possible.
   2. We might move them into a util.



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2241416406


##########
runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java:
##########
@@ -39,62 +37,103 @@ public class PythonActionExecutor {
             "from flink_agents.plan import function\n"
                     + "from flink_agents.runtime import flink_runner_context\n"
                     + "from flink_agents.runtime import python_java_utils";
+
+    // =========== RUNNER CONTEXT ===========
     private static final String CREATE_FLINK_RUNNER_CONTEXT =
             "flink_runner_context.create_flink_runner_context";
+    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME_PREFIX = "flink_runner_context_";
+    private static final AtomicLong FLINK_RUNNER_CONTEXT_VAR_ID = new AtomicLong(0);
+
+    // ========== ASYNC THREAD POOL ===========
+    private static final String CREATE_ASYNC_THREAD_POOL =
+            "flink_runner_context.create_async_thread_pool";
+    private static final String PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX =
+            "python_async_thread_pool";
+    private static final AtomicLong PYTHON_ASYNC_THREAD_POOL_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON GENERATOR ===========
+    private static final String CALL_PYTHON_GENERATOR = "function.call_python_generator";
+    private static final String PYTHON_GENERATOR_VAR_NAME_PREFIX = "python_generator_";
+    private static final AtomicLong PYTHON_GENERATOR_VAR_ID = new AtomicLong(0);
+
+    // =========== PYTHON AND JAVA OBJECT CONVERT ===========
     private static final String CONVERT_TO_PYTHON_OBJECT =
             "python_java_utils.convert_to_python_object";
     private static final String WRAP_TO_INPUT_EVENT = "python_java_utils.wrap_to_input_event";
     private static final String GET_OUTPUT_FROM_OUTPUT_EVENT =
             "python_java_utils.get_output_from_output_event";
-    private static final String FLINK_RUNNER_CONTEXT_VAR_NAME = "flink_runner_context";
 
     private final PythonEnvironmentManager environmentManager;
     private final String agentPlanJson;
-    private PythonRunnerContextImpl runnerContext;
-
     private PythonInterpreter interpreter;
+    private String pythonAsyncThreadPoolObjectName;
 
     public PythonActionExecutor(PythonEnvironmentManager environmentManager, String agentPlanJson) {
         this.environmentManager = environmentManager;
         this.agentPlanJson = agentPlanJson;
     }
 
-    public void open(
-            MapState<String, MemoryObjectImpl.MemoryItem> shortTermMemState,
-            Runnable mailboxThreadChecker)
-            throws Exception {
+    public void open() throws Exception {
         environmentManager.open();
         EmbeddedPythonEnvironment env = environmentManager.createEnvironment();
 
         interpreter = env.getInterpreter();
         interpreter.exec(PYTHON_IMPORTS);
 
-        runnerContext = new PythonRunnerContextImpl(shortTermMemState, mailboxThreadChecker);
-
-        // TODO: remove the set and get runner context after updating pemja to version 0.5.3
-        Object pythonRunnerContextObject =
-                interpreter.invoke(CREATE_FLINK_RUNNER_CONTEXT, runnerContext, agentPlanJson);
-        interpreter.set(FLINK_RUNNER_CONTEXT_VAR_NAME, pythonRunnerContextObject);
+        // TODO: remove the set and get thread pool after updating pemja to version 0.5.3
+        Object pythonAsyncThreadPool = interpreter.invoke(CREATE_ASYNC_THREAD_POOL);
+        this.pythonAsyncThreadPoolObjectName =
+                PYTHON_ASYNC_THREAD_POOL_VAR_NAME_PREFIX
+                        + PYTHON_ASYNC_THREAD_POOL_VAR_ID.incrementAndGet();
+        interpreter.set(pythonAsyncThreadPoolObjectName, pythonAsyncThreadPool);

Review Comment:
   > 1. What does this todo mean?
   Due to a bug in Pemja that may prevent proper cleanup of Python objects, we are using get and set methods as a temporary workaround. For detailed context, please refer to: #83 .
   
   > 2. Is it necessary to have separate thread pools for each operator?
   Since each operator has its own Python interpreter, and the thread pool is created and managed within the interpreter, it is necessary to allocate a separate thread pool for each operator.



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


Re: [PR] [API][Runtime] Support agent and action asynchronous execution [flink-agents]

Posted by "codenohup (via GitHub)" <gi...@apache.org>.
codenohup commented on code in PR #80:
URL: https://github.com/apache/flink-agents/pull/80#discussion_r2241466900


##########
runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java:
##########
@@ -116,7 +137,23 @@ public void open() throws Exception {
                         TypeInformation.of(String.class),
                         TypeInformation.of(MemoryObjectImpl.MemoryItem.class));
         shortTermMemState = getRuntimeContext().getMapState(shortTermMemStateDescriptor);
-        runnerContext = new RunnerContextImpl(shortTermMemState, this::checkMailboxThread);
+
+        // init agent processing related state
+        actionTasksState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "actionTasks", TypeInformation.of(ActionTask.class)));
+        pendingInputEventsState =
+                getRuntimeContext()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "pendingInputEvents", TypeInformation.of(Event.class)));
+        currentProcessingKeysOperatorState =
+                getOperatorStateBackend()
+                        .getListState(
+                                new ListStateDescriptor<>(
+                                        "currentProcessingKeys", TypeInformation.of(Object.class)));

Review Comment:
   > 2. How do we handle the redistribution of the operator state?
   
   I suggest replacing _currentProcessingKeysOpState_ to _UnionList_. This change would ensure that all keys are retained even after a parallelism modification or restart. After that, we can simply iterate through the list and clean up any keys that fall outside the current key range.
   
   



-- 
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: issues-unsubscribe@flink.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org