| ################################################################################ |
| # 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. |
| ################################################################################ |
| import asyncio |
| |
| from pyflink.datastream import async_retry_strategies |
| from pyflink.fn_execution.datastream.operations import Operation, AsyncOperationMixin |
| from pyflink.fn_execution.datastream.process.async_function.operation import Emitter, \ |
| AsyncFunctionRunner, ResultHandler, RetryableResultHandler |
| from pyflink.fn_execution.datastream.process.async_function.queue import OrderedStreamElementQueue |
| from pyflink.fn_execution.metrics.process.metric_impl import GenericMetricGroup |
| from pyflink.fn_execution.utils import operation_utils |
| from pyflink.table import FunctionContext |
| |
| ASYNC_SCALAR_FUNCTION_URN = "flink:transform:async_scalar_function:v1" |
| |
| |
| class AsyncScalarFunctionOperation(Operation, AsyncOperationMixin): |
| """ |
| Operation for executing Python async scalar functions. |
| |
| This operation implements true asynchronous execution by leveraging the async |
| infrastructure from DataStream API's AsyncOperation: |
| - AsyncFunctionRunner: Manages asyncio event loop in a separate thread |
| - Queue: Maintains in-flight async operations with configurable capacity |
| - Emitter: Collects and emits results asynchronously |
| - Non-blocking: Multiple async operations can be in-flight simultaneously |
| |
| This provides high performance for I/O-bound async operations compared to |
| synchronous blocking execution. |
| """ |
| |
| def __init__(self, serialized_fn): |
| if serialized_fn.metric_enabled: |
| self.base_metric_group = GenericMetricGroup(None, None) |
| else: |
| self.base_metric_group = None |
| |
| self._capacity = serialized_fn.async_options.max_concurrent_operations |
| self._timeout = serialized_fn.async_options.timeout_ms / 1000.0 |
| self._retry_enabled = serialized_fn.async_options.retry_enabled |
| self._max_attempts = serialized_fn.async_options.retry_max_attempts |
| self._retry_delay = serialized_fn.async_options.retry_delay_ms / 1000.0 |
| |
| scalar_function, variable_dict, self.user_defined_funcs = \ |
| operation_utils.extract_user_defined_function( |
| serialized_fn.udfs[0], one_arg_optimization=False) |
| |
| # Create the eval function |
| self._eval_func = eval('lambda value: %s' % scalar_function, variable_dict) |
| |
| # Create ordered queue to maintain result order |
| self._queue = OrderedStreamElementQueue(self._capacity, self._raise_exception_if_exists) |
| |
| # Async execution components |
| self._async_function_runner = None |
| self._emitter = None |
| self._exception = None |
| self._output_processor = None |
| |
| # Job parameters |
| self._job_parameters = {p.key: p.value for p in serialized_fn.job_parameters} |
| if serialized_fn.HasField('runtime_context'): |
| rc = serialized_fn.runtime_context |
| self._runtime_context = { |
| 'task_name': rc.task_name, |
| 'task_name_with_subtasks': rc.task_name_with_subtasks, |
| 'number_of_parallel_subtasks': rc.number_of_parallel_subtasks, |
| 'max_number_of_parallel_subtasks': rc.max_number_of_parallel_subtasks, |
| 'index_of_this_subtask': rc.index_of_this_subtask, |
| 'attempt_number': rc.attempt_number, |
| } |
| else: |
| self._runtime_context = {} |
| |
| def set_output_processor(self, output_processor): |
| """Set the output processor for emitting results. |
| |
| This method is called by FunctionOperation for AsyncOperationMixin implementations. |
| """ |
| self._output_processor = output_processor |
| |
| def open(self): |
| # Open user defined functions |
| for user_defined_func in self.user_defined_funcs: |
| if hasattr(user_defined_func, 'open'): |
| user_defined_func.open(FunctionContext( |
| self.base_metric_group, self._job_parameters, |
| **self._runtime_context)) |
| |
| # Start emitter thread to collect async results |
| self._emitter = Emitter(self._mark_exception, self._output_processor, self._queue) |
| self._emitter.daemon = True |
| self._emitter.start() |
| |
| # Start async function runner with event loop |
| self._async_function_runner = AsyncFunctionRunner() |
| self._async_function_runner.daemon = True |
| self._async_function_runner.start() |
| self._async_function_runner.wait_ready() |
| |
| def close(self): |
| # Stop emitter |
| if self._emitter is not None: |
| self._emitter.stop() |
| self._emitter = None |
| |
| # Stop async function runner |
| if self._async_function_runner is not None: |
| self._async_function_runner.stop() |
| self._async_function_runner = None |
| |
| self._exception = None |
| |
| # Close user defined functions |
| for user_defined_func in self.user_defined_funcs: |
| if hasattr(user_defined_func, 'close'): |
| user_defined_func.close() |
| |
| def process_element(self, value): |
| """ |
| Process an input element asynchronously. |
| |
| This is non-blocking - it submits the async operation and returns immediately, |
| allowing multiple operations to be in-flight simultaneously. |
| """ |
| self._raise_exception_if_exists() |
| |
| entry = self._queue.put(None, 0, 0, value) |
| |
| async def execute_async(rh): |
| try: |
| # Call the eval function |
| result = self._eval_func(value) |
| |
| # Check if any result is a coroutine and await it |
| if asyncio.iscoroutine(result): |
| final_result = await result |
| else: |
| final_result = result |
| |
| # Complete with results (list format) |
| rh.complete([final_result]) |
| except Exception as e: |
| rh.complete_exceptionally(e) |
| |
| # Create result handler |
| result_handler = ResultHandler( |
| self.__class__.__name__, |
| self._timeout_func, |
| self._mark_exception, |
| value, |
| entry) |
| |
| if self._retry_enabled: |
| retry_strategy = async_retry_strategies.FixedDelayRetryStrategy( |
| max_attempts=self._max_attempts, |
| backoff_time_millis=int(self._retry_delay * 1000), |
| result_predicate=None, |
| exception_predicate=lambda ex: True) # Retry on all exceptions |
| |
| # Wrap with retryable handler |
| retryable_handler = RetryableResultHandler( |
| result_handler, |
| lambda r, h: self._async_function_runner.run_async( |
| lambda _: execute_async(h), r, h), |
| retry_strategy) |
| |
| # Register timeout |
| assert self._timeout > 0 |
| retryable_handler.register_timeout(self._timeout) |
| |
| # Submit to event loop asynchronously |
| self._async_function_runner.run_async( |
| lambda r: execute_async(retryable_handler), |
| value, |
| retryable_handler) |
| else: |
| # Register timeout if configured |
| if self._timeout > 0: |
| result_handler.register_timeout(self._timeout) |
| |
| # Submit to event loop asynchronously |
| self._async_function_runner.run_async( |
| lambda r: execute_async(result_handler), |
| value, |
| result_handler) |
| |
| def _timeout_func(self, record): |
| """Handle timeout for async operations.""" |
| raise TimeoutError("Async function call has timed out for input: " + str(record)) |
| |
| def finish(self): |
| """Wait for all in-flight async operations to complete.""" |
| self._wait_for_in_flight_inputs_finished() |
| self._update_gauge(self.base_metric_group) |
| |
| def _wait_for_in_flight_inputs_finished(self): |
| """Wait until all in-flight async operations are completed.""" |
| while not self._queue.is_empty(): |
| self._queue.wait_for_in_flight_elements_processed() |
| self._raise_exception_if_exists() |
| |
| def _mark_exception(self, exception): |
| """Mark an exception that occurred during async execution.""" |
| self._exception = exception |
| |
| def _raise_exception_if_exists(self): |
| """Raise exception if one occurred during async execution.""" |
| if self._exception is not None: |
| raise self._exception |
| |
| def _update_gauge(self, base_metric_group): |
| """Update metric gauges.""" |
| if base_metric_group is not None: |
| for name in base_metric_group._flink_gauge: |
| flink_gauge = base_metric_group._flink_gauge[name] |
| beam_gauge = base_metric_group._beam_gauge[name] |
| beam_gauge.set(flink_gauge()) |
| for sub_group in base_metric_group._sub_groups: |
| self._update_gauge(sub_group) |
