from typing import List from PySide6.QtCore import * from PySide6.QtGui import QAction from PySide6.QtWidgets import * from Localizations import ( DELETING_TASK, CANCELLING_TASK, CONFIRM_DELETION_TITLE, CONFIRM_DELETION, SHOWING_TASK_CONTEXT_MENU, CANCELED, CANCEL, PROPERTIES, COMPLETED, SHOWING_PROPERTIES_FOR_TASK, DELETE, RESTART, IN_PROGRESS, ERROR, RESTARTING_TASK, ) from ModelInfoDialog import ModelInfoDialog from QuantizationThread import QuantizationThread from Logger import Logger from error_handling import handle_error class TaskListItem(QWidget): def __init__( self, task_name, log_file, show_progress_bar=True, parent=None, show_properties=False, logger=Logger, quant_threads=List[QuantizationThread], ) -> None: super().__init__(parent) self.quant_threads = quant_threads self.task_name = task_name self.log_file = log_file self.logger = logger self.show_properties = show_properties self.status = "Pending" layout = QHBoxLayout(self) self.task_label = QLabel(task_name) self.progress_bar = QProgressBar() self.progress_bar.setRange(0, 100) self.status_label = QLabel(self.status) layout.addWidget(self.task_label) layout.addWidget(self.progress_bar) layout.addWidget(self.status_label) # Hide progress bar if show_progress_bar is False self.progress_bar.setVisible(show_progress_bar) # Use indeterminate progress bar if not showing percentage if not show_progress_bar: self.progress_bar.setRange(0, 0) self.progress_timer = QTimer(self) self.progress_timer.timeout.connect(self.update_progress) self.progress_value = 0 def show_task_context_menu(self, position) -> None: self.logger.debug(SHOWING_TASK_CONTEXT_MENU) item = self.task_list.itemAt(position) if item is not None: context_menu = QMenu(self) properties_action = QAction(PROPERTIES, self) properties_action.triggered.connect(lambda: self.show_task_properties(item)) context_menu.addAction(properties_action) task_item = self.task_list.itemWidget(item) if task_item.status != COMPLETED: cancel_action = QAction(CANCEL, self) cancel_action.triggered.connect(lambda: self.cancel_task(item)) context_menu.addAction(cancel_action) if task_item.status == CANCELED: restart_action = QAction(RESTART, self) restart_action.triggered.connect(lambda: self.restart_task(task_item)) context_menu.addAction(restart_action) delete_action = QAction(DELETE, self) delete_action.triggered.connect(lambda: self.delete_task(item)) context_menu.addAction(delete_action) context_menu.exec(self.task_list.viewport().mapToGlobal(position)) def show_task_properties(self, item) -> None: self.logger.debug(SHOWING_PROPERTIES_FOR_TASK.format(item.text())) for thread in self.quant_threads: model_info_dialog = ModelInfoDialog(thread.model_info, self) model_info_dialog.exec() break def cancel_task(self, item) -> None: # TODO: fix possibly buggy signal behavior task_item = self.task_list.itemWidget(item) if task_item: task_name = task_item.task_name # Store the name before any changes self.logger.info(CANCELLING_TASK.format(task_name)) # Find the thread and disconnect signals before terminating for thread in self.quant_threads: if thread.log_file == task_item.log_file: # Disconnect all signals from this thread first try: thread.error_signal.disconnect() # Disconnect all error signal connections thread.output_signal.disconnect() # Disconnect all output signal connections except TypeError: # No connections to disconnect pass # Now terminate the thread thread.terminate() self.quant_threads.remove(thread) break def delete_task(self, item) -> None: task_item = self.task_list.itemWidget(item) if not task_item: return task_name = task_item.task_name # Store task_name before deletion self.logger.info(DELETING_TASK.format(task_name)) reply = QMessageBox.question( self, CONFIRM_DELETION_TITLE, CONFIRM_DELETION, QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No, QMessageBox.StandardButton.No, ) if reply == QMessageBox.StandardButton.Yes: # Cancel the task first (which disconnects signals) self.cancel_task(item) # Now remove from list and delete row = self.task_list.row(item) self.task_list.takeItem(row) # Delete the widget after removing from list task_item.deleteLater() def update_status(self, status) -> None: self.status = status self.status_label.setText(status) if status == IN_PROGRESS: # Only start timer if showing percentage progress if self.progress_bar.isVisible(): self.progress_bar.setRange(0, 100) self.progress_timer.start(100) elif status == COMPLETED: self.progress_timer.stop() self.progress_bar.setValue(100) elif status == CANCELED: self.progress_timer.stop() self.progress_bar.setValue(0) def set_error(self) -> None: self.status = ERROR self.status_label.setText(ERROR) self.status_label.setStyleSheet("color: red;") self.progress_bar.setRange(0, 100) self.progress_timer.stop() def update_progress(self, value=None) -> None: if value is not None: # Update progress bar with specific value self.progress_value = value self.progress_bar.setValue(self.progress_value) else: return def restart_task(self, task_item) -> None: self.logger.info(RESTARTING_TASK.format(task_item.task_name)) for thread in self.quant_threads: if thread.log_file == task_item.log_file: new_thread = QuantizationThread( thread.command, thread.cwd, thread.log_file ) self.quant_threads.append(new_thread) new_thread.status_signal.connect(task_item.update_status) new_thread.finished_signal.connect( lambda: self.task_finished(new_thread, task_item) ) new_thread.error_signal.connect( lambda err: handle_error(self.logger, err, task_item) ) new_thread.model_info_signal.connect(self.update_model_info) new_thread.start() task_item.update_status(IN_PROGRESS) break