import time
import pandas
from typing import List, Dict, Callable, Optional
from queue import Queue
from feldera.rest.errors import FelderaAPIError
from feldera.enums import PipelineStatus
from feldera.rest.pipeline import Pipeline as InnerPipeline
from feldera.rest.feldera_client import FelderaClient
from feldera._callback_runner import _CallbackRunnerInstruction, CallbackRunner
from feldera.output_handler import OutputHandler
from feldera._helpers import ensure_dataframe_has_columns, chunk_dataframe
[docs]
class Pipeline:
def __init__(self, name: str, client: FelderaClient):
self.name = name
self.client: FelderaClient = client
self._inner: InnerPipeline | None = None
self.views_tx: List[Dict[str, Queue]] = []
def __setup_output_listeners(self):
"""
Internal function used to set up the output listeners.
:meta private:
"""
for view_queue in self.views_tx:
for view_name, queue in view_queue.items():
# sends a message to the callback runner to start listening
queue.put(_CallbackRunnerInstruction.PipelineStarted)
# block until the callback runner is ready
queue.join()
[docs]
def status(self) -> PipelineStatus:
"""
Return the current status of the pipeline.
"""
try:
inner = self.client.get_pipeline(self.name)
self._inner = inner
return PipelineStatus.from_str(inner.deployment_status)
except FelderaAPIError as err:
if err.status_code == 404:
return PipelineStatus.NOT_FOUND
else:
raise err
[docs]
def listen(self, view_name: str) -> OutputHandler:
"""
Listen to the output of the provided view so that it is available in the notebook / python code.
:param view_name: The name of the view to listen to.
"""
queue: Optional[Queue] = None
if self.status() not in [PipelineStatus.PAUSED, PipelineStatus.RUNNING]:
queue = Queue(maxsize=1)
self.views_tx.append({view_name: queue})
handler = OutputHandler(self.client, self.name, view_name, queue)
handler.start()
return handler
[docs]
def foreach_chunk(self, view_name: str, callback: Callable[[pandas.DataFrame, int], None]):
"""
Run the given callback on each chunk of the output of the specified view.
:param view_name: The name of the view.
:param callback: The callback to run on each chunk. The callback should take two arguments:
- **chunk** -> The chunk as a pandas DataFrame
- **seq_no** -> The sequence number. The sequence number is a monotonically increasing integer that
starts from 0. Note that the sequence number is unique for each chunk, but not necessarily contiguous.
Please note that the callback is run in a separate thread, so it should be thread-safe.
Please note that the callback should not block for a long time, as by default, backpressure is enabled and
will block the pipeline.
.. note::
- The callback must be thread-safe as it will be run in a separate thread.
"""
queue: Optional[Queue] = None
if self.status() not in [PipelineStatus.RUNNING, PipelineStatus.PAUSED]:
queue = Queue(maxsize=1)
self.views_tx.append({view_name: queue})
handler = CallbackRunner(self.client, self.name, view_name, callback, queue)
handler.start()
[docs]
def wait_for_completion(self, shutdown: bool = False):
"""
Block until the pipeline has completed processing all input records.
This method blocks until (1) all input connectors attached to the pipeline
have finished reading their input data sources and issued end-of-input
notifications to the pipeline, and (2) all inputs received from these
connectors have been fully processed and corresponding outputs have been
sent out through the output connectors.
This method will block indefinitely if at least one of the input
connectors attached to the pipeline is a streaming connector, such as
Kafka, that does not issue the end-of-input notification.
:param shutdown: If True, the pipeline will be shutdown after completion. False by default.
:raises RuntimeError: If the pipeline returns unknown metrics.
"""
if self.status() not in [
PipelineStatus.RUNNING,
PipelineStatus.INITIALIZING,
PipelineStatus.PROVISIONING,
]:
raise RuntimeError("Pipeline must be running to wait for completion")
while True:
metrics: dict = self.client.get_pipeline_stats(self.name).get("global_metrics")
pipeline_complete: bool = metrics.get("pipeline_complete")
if pipeline_complete is None:
raise RuntimeError("received unknown metrics from the pipeline, pipeline_complete is None")
if pipeline_complete:
break
time.sleep(1)
if shutdown:
self.shutdown()
[docs]
def start(self):
"""
.. _start:
Starts this pipeline.
:raises RuntimeError: If the pipeline returns unknown metrics.
"""
status = self.status()
if status != PipelineStatus.SHUTDOWN:
raise RuntimeError(f"pipeline {self.name} in state: {str(status.name)} cannot be started")
self.pause()
self.__setup_output_listeners()
self.resume()
[docs]
def restart(self):
"""
Restarts the pipeline.
"""
self.shutdown()
self.start()
[docs]
def wait_for_idle(
self,
idle_interval_s: float = 5.0,
timeout_s: float = 600.0,
poll_interval_s: float = 0.2
):
"""
Wait for the pipeline to become idle and then returns.
Idle is defined as a sufficiently long interval in which the number of
input and processed records reported by the pipeline do not change, and
they equal each other (thus, all input records present at the pipeline
have been processed).
:param idle_interval_s: Idle interval duration (default is 5.0 seconds).
:param timeout_s: Timeout waiting for idle (default is 600.0 seconds).
:param poll_interval_s: Polling interval, should be set substantially
smaller than the idle interval (default is 0.2 seconds).
:raises ValueError: If idle interval is larger than timeout, poll interval
is larger than timeout, or poll interval is larger than idle interval.
:raises RuntimeError: If the metrics are missing or the timeout was
reached.
"""
if idle_interval_s > timeout_s:
raise ValueError(f"idle interval ({idle_interval_s}s) cannot be larger than timeout ({timeout_s}s)")
if poll_interval_s > timeout_s:
raise ValueError(f"poll interval ({poll_interval_s}s) cannot be larger than timeout ({timeout_s}s)")
if poll_interval_s > idle_interval_s:
raise ValueError(f"poll interval ({poll_interval_s}s) cannot be larger "
f"than idle interval ({idle_interval_s}s)")
start_time_s = time.monotonic()
idle_started_s = None
prev = (0, 0)
while True:
now_s = time.monotonic()
# Metrics retrieval
metrics: dict = self.client.get_pipeline_stats(self.name).get("global_metrics")
total_input_records: int | None = metrics.get("total_input_records")
total_processed_records: int | None = metrics.get("total_processed_records")
if total_input_records is None:
raise RuntimeError("total_input_records is missing from the pipeline metrics")
if total_processed_records is None:
raise RuntimeError("total_processed_records is missing from the pipeline metrics")
# Idle check
unchanged = prev[0] == total_input_records and prev[1] == total_processed_records
equal = total_input_records == total_processed_records
prev = (total_input_records, total_processed_records)
if unchanged and equal:
if idle_started_s is None:
idle_started_s = now_s
else:
idle_started_s = None
if idle_started_s is not None and now_s - idle_started_s >= idle_interval_s:
return
# Timeout
if now_s - start_time_s >= timeout_s:
raise RuntimeError(f"waiting for idle reached timeout ({timeout_s}s)")
time.sleep(poll_interval_s)
[docs]
def pause(self):
"""
Pause the pipeline.
"""
self.client.pause_pipeline(self.name)
[docs]
def shutdown(self):
"""
Shut down the pipeline.
"""
for view_queue in self.views_tx:
for view_name, queue in view_queue.items():
# sends a message to the callback runner to stop listening
queue.put(_CallbackRunnerInstruction.RanToCompletion)
# block until the callback runner has been stopped
queue.join()
self.client.shutdown_pipeline(self.name)
[docs]
def resume(self):
"""
Resumes the pipeline.
"""
self.client.start_pipeline(self.name)
[docs]
def delete(self):
"""
Deletes the pipeline.
"""
self.client.delete_pipeline(self.name)
[docs]
@staticmethod
def get(name: str, client: FelderaClient) -> 'Pipeline':
"""
Get the pipeline if it exists.
:param name: The name of the pipeline.
:param client: The FelderaClient instance.
"""
try:
inner = client.get_pipeline(name)
pipeline = Pipeline(inner.name, client)
pipeline.__inner = inner
return pipeline
except FelderaAPIError as err:
if err.status_code == 404:
raise RuntimeError(f"Pipeline with name {name} not found")