Add logging for pipeline
This commit is contained in:
@@ -29,3 +29,6 @@ class FrameProcessor:
|
|||||||
async def interrupted(self) -> None:
|
async def interrupted(self) -> None:
|
||||||
"""Handle any cleanup if the pipeline was interrupted."""
|
"""Handle any cleanup if the pipeline was interrupted."""
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
def __str__(self):
|
||||||
|
return self.__class__.__name__
|
||||||
|
|||||||
@@ -1,8 +1,9 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import logging
|
||||||
from typing import AsyncGenerator, AsyncIterable, Iterable, List
|
from typing import AsyncGenerator, AsyncIterable, Iterable, List
|
||||||
from dailyai.pipeline.frame_processor import FrameProcessor
|
from dailyai.pipeline.frame_processor import FrameProcessor
|
||||||
|
|
||||||
from dailyai.pipeline.frames import EndPipeFrame, EndFrame, Frame
|
from dailyai.pipeline.frames import AudioFrame, EndPipeFrame, EndFrame, Frame
|
||||||
|
|
||||||
|
|
||||||
class Pipeline:
|
class Pipeline:
|
||||||
@@ -29,6 +30,8 @@ class Pipeline:
|
|||||||
self.source: asyncio.Queue[Frame] = source or asyncio.Queue()
|
self.source: asyncio.Queue[Frame] = source or asyncio.Queue()
|
||||||
self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
|
self.sink: asyncio.Queue[Frame] = sink or asyncio.Queue()
|
||||||
|
|
||||||
|
self._logger = logging.getLogger("dailyai.pipeline")
|
||||||
|
|
||||||
def set_source(self, source: asyncio.Queue[Frame]):
|
def set_source(self, source: asyncio.Queue[Frame]):
|
||||||
"""Set the source queue for this pipeline. Frames from this queue
|
"""Set the source queue for this pipeline. Frames from this queue
|
||||||
will be processed by each frame_processor in the pipeline, or order
|
will be processed by each frame_processor in the pipeline, or order
|
||||||
@@ -82,9 +85,11 @@ class Pipeline:
|
|||||||
try:
|
try:
|
||||||
while True:
|
while True:
|
||||||
initial_frame = await self.source.get()
|
initial_frame = await self.source.get()
|
||||||
|
self._log_frame(initial_frame, 0)
|
||||||
async for frame in self._run_pipeline_recursively(
|
async for frame in self._run_pipeline_recursively(
|
||||||
initial_frame, self._processors
|
initial_frame, self._processors
|
||||||
):
|
):
|
||||||
|
self._log_frame(frame, len(self._processors) + 1)
|
||||||
await self.sink.put(frame)
|
await self.sink.put(frame)
|
||||||
|
|
||||||
if isinstance(initial_frame, EndFrame) or isinstance(
|
if isinstance(initial_frame, EndFrame) or isinstance(
|
||||||
@@ -96,18 +101,32 @@ class Pipeline:
|
|||||||
# here.
|
# here.
|
||||||
for processor in self._processors:
|
for processor in self._processors:
|
||||||
await processor.interrupted()
|
await processor.interrupted()
|
||||||
pass
|
|
||||||
|
|
||||||
async def _run_pipeline_recursively(
|
async def _run_pipeline_recursively(
|
||||||
self, initial_frame: Frame, processors: List[FrameProcessor]
|
self, initial_frame: Frame, processors: List[FrameProcessor], depth=1
|
||||||
) -> AsyncGenerator[Frame, None]:
|
) -> AsyncGenerator[Frame, None]:
|
||||||
"""Internal function to add frames to the pipeline as they're yielded
|
"""Internal function to add frames to the pipeline as they're yielded
|
||||||
by each processor."""
|
by each processor."""
|
||||||
if processors:
|
if processors:
|
||||||
|
self._log_frame(initial_frame, depth)
|
||||||
async for frame in processors[0].process_frame(initial_frame):
|
async for frame in processors[0].process_frame(initial_frame):
|
||||||
async for final_frame in self._run_pipeline_recursively(
|
async for final_frame in self._run_pipeline_recursively(
|
||||||
frame, processors[1:]
|
frame, processors[1:], depth + 1
|
||||||
):
|
):
|
||||||
yield final_frame
|
yield final_frame
|
||||||
else:
|
else:
|
||||||
yield initial_frame
|
yield initial_frame
|
||||||
|
|
||||||
|
def _log_frame(self, frame: Frame, depth: int):
|
||||||
|
if not isinstance(frame, AudioFrame):
|
||||||
|
if depth == 0:
|
||||||
|
source = "source"
|
||||||
|
dest = str(self._processors[depth - 1])
|
||||||
|
elif depth == len(self._processors) + 1:
|
||||||
|
dest = "sink"
|
||||||
|
source = str(self._processors[depth - 2])
|
||||||
|
else:
|
||||||
|
dest = str(self._processors[depth - 1])
|
||||||
|
source = str(self._processors[depth - 2])
|
||||||
|
|
||||||
|
self._logger.debug(" " * (depth + 1) + " -> ".join([source, str(frame), dest]))
|
||||||
|
|||||||
Reference in New Issue
Block a user