Coverage for intelligence_toolkit/AI/base_batch_async.py: 100%

1# Copyright (c) 2024 Microsoft Corporation. All rights reserved. 

2# Licensed under the MIT license. See LICENSE file in the project. 

3# 

4import asyncio 

5import logging 

6 

7from intelligence_toolkit.helpers.progress_batch_callback import ProgressBatchCallback 

8 

9logger = logging.getLogger(__name__) 

10 

11 

12class BaseBatchAsync: 

13 total_tasks: int = 1 

14 completed_tasks: int = 0 

15 previous_completed_tasks: int = 0 

16 

17 async def track_progress( 

18 self, tasks: list[asyncio.Task], callbacks: list[ProgressBatchCallback] 

19 ): 

20 while not all(task.done() for task in tasks): 

21 await asyncio.sleep(0.1) 

22 if self.completed_tasks != self.previous_completed_tasks: 

23 for callback in callbacks: 

24 callback.on_batch_change(self.completed_tasks, self.total_tasks) 

25 self.previous_completed_tasks = self.completed_tasks 

26 

27 if self.completed_tasks != self.previous_completed_tasks: 

28 for callback in callbacks: 

29 callback.on_batch_change(self.completed_tasks, self.total_tasks) 

30 

31 if self.completed_tasks == self.total_tasks: 

32 for callback in callbacks: 

33 callback.on_batch_change(self.completed_tasks, self.total_tasks) 

34 

35 def progress_callback(self) -> None: 

36 self.completed_tasks += 1