feat(parallel): add support for iMatrix generation tracking

- add support for iMatrix generation tracking
- don't adjust progress bar when indeterminate
This commit is contained in:
BuildTools 2024-09-09 19:26:56 -07:00
parent c96380cbf8
commit cee4294ecf
No known key found for this signature in database
GPG Key ID: 3270C066C15D530B
4 changed files with 27 additions and 6 deletions

View File

@ -1,6 +1,5 @@
import importlib
import json
import re
import shutil
import urllib.request
import urllib.error
@ -1547,6 +1546,7 @@ def task_finished(self, thread, task_item) -> None:
if thread in self.quant_threads:
self.quant_threads.remove(thread)
task_item.update_status(COMPLETED)
self.setAttribute(Qt.WA_WindowModified, True) # Set modified flag
def show_task_details(self, item) -> None:
self.logger.debug(SHOWING_TASK_DETAILS_FOR.format(item.text()))
@ -1649,7 +1649,7 @@ def generate_imatrix(self) -> None:
task_item = TaskListItem(
task_name,
log_file,
show_progress_bar=False,
show_progress_bar=True,
logger=self.logger,
quant_threads=self.quant_threads,
)
@ -1658,7 +1658,12 @@ def generate_imatrix(self) -> None:
self.task_list.addItem(list_item)
self.task_list.setItemWidget(list_item, task_item)
imatrix_chunks = None
thread.status_signal.connect(task_item.update_status)
thread.output_signal.connect(
lambda line, ti=task_item: self.parse_progress(line, ti, imatrix_chunks)
)
thread.finished_signal.connect(
lambda: self.task_finished(thread, task_item)
)

View File

@ -79,14 +79,31 @@ def parse_model_info(self, line) -> None:
f"{quant_type}: {tensors} tensors"
)
def parse_progress(self, line, task_item) -> None:
def parse_progress(self, line, task_item, imatrix_chunks=None) -> None:
# Parses the output line for progress information and updates the task item.
match = re.search(r"\[\s*(\d+)\s*/\s*(\d+)\s*].*", line)
if match:
current = int(match.group(1))
total = int(match.group(2))
progress = int((current / total) * 100)
task_item.update_progress(progress)
else:
imatrix_match = re.search(
r"compute_imatrix: computing over (\d+) chunks with batch_size \d+",
line,
)
if imatrix_match:
imatrix_chunks = int(imatrix_match.group(1))
elif imatrix_chunks is not None:
save_match = re.search(
r"save_imatrix: stored collected data after (\d+) chunks in .*",
line,
)
if save_match:
saved_chunks = int(save_match.group(1))
progress = int((saved_chunks / self.imatrix_chunks) * 100)
task_item.update_progress(progress)
def terminate(self) -> None:
# Terminate the subprocess if it's still running

View File

@ -162,8 +162,7 @@ def update_progress(self, value=None) -> None:
self.progress_value = value
self.progress_bar.setValue(self.progress_value)
else:
# Set progress bar to zero for indeterminate progress
self.progress_bar.setValue(0)
return
def restart_task(self, task_item) -> None:
self.logger.info(RESTARTING_TASK.format(task_item.task_name))

View File

@ -51,7 +51,7 @@ class Config:
class Task(BaseModel):
id: str = Field(..., description="Unique identifier for the task")
# id: str = Field(..., description="Unique identifier for the task")
status: str = Field(..., description="Current status of the task")
progress: float = Field(..., description="Progress of the task as a percentage")