Skip to content

Commit b89b932

Browse files
authored
Added a timer to the KG creation pipeline (#107)
* Added a timer to the KG creation pipeline * Reverted TaskPipelineNode logger call to debug * Adds basic TQDM progress bar to KG creation pipeline * Removed progress bar from KG pipeline
1 parent af27578 commit b89b932

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed

src/neo4j_genai/experimental/pipeline/pipeline.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import enum
1919
import logging
2020
from datetime import datetime
21+
from timeit import default_timer
2122
from typing import Any, AsyncGenerator, Awaitable, Callable, Optional
2223

2324
from pydantic import BaseModel, Field
@@ -119,6 +120,7 @@ async def execute(self, **kwargs: Any) -> RunResult | None:
119120
was unsuccessful.
120121
"""
121122
logger.debug(f"Running component {self.name} with {kwargs}")
123+
start_time = default_timer()
122124
try:
123125
await self.set_status(RunStatus.RUNNING)
124126
except PipelineStatusUpdateError:
@@ -130,6 +132,8 @@ async def execute(self, **kwargs: Any) -> RunResult | None:
130132
status=self.status,
131133
result=component_result,
132134
)
135+
end_time = default_timer()
136+
logger.debug(f"Component {self.name} finished in {end_time - start_time}s")
133137
return run_result
134138

135139
def validate_inputs_config(self, input_data: dict[str, Any]) -> None:
@@ -467,8 +471,12 @@ def validate_inputs_config(self, data: dict[str, Any]) -> None:
467471
task.validate_inputs_config(data)
468472

469473
async def run(self, data: dict[str, Any]) -> dict[str, Any]:
474+
logger.debug("Starting pipeline")
475+
start_time = default_timer()
470476
self.validate_inputs_config(data)
471477
self.reinitialize()
472478
orchestrator = Orchestrator(self)
473479
await orchestrator.run(data)
480+
end_time = default_timer()
481+
logger.debug(f"Pipeline finished in {end_time - start_time}s")
474482
return self._final_results.all()

0 commit comments

Comments
 (0)