[core] First implementation to kill submitted tasks

Author: Alxiice

meshroom/core/__init__.py

@@ -121,7 +121,7 @@ def loadClasses(folder: str, packageName: str, classType: type) -> list[type]:
                         classes.append(p)
             except Exception as e:
                 if classType == BaseSubmitter:
-                    logging.warning(f" Could not load submitter {pluginName} from package '{package.__name__}'")
+                    logging.warning(f" Could not load submitter {pluginName} from package '{package.__name__}'\n{e}")
                 else:
                     tb = traceback.extract_tb(e.__traceback__)
                     last_call = tb[-1]

meshroom/core/desc/node.py

@@ -7,6 +7,8 @@
 import shlex
 import shutil
 import sys
+import signal
+import subprocess
 
 from .computation import Level, StaticNodeSize
 from .attribute import StringParam, ColorParam, ChoiceParam
@@ -20,6 +22,29 @@
 _MESHROOM_COMPUTE_DEPS = ["psutil"]
 
 
+# Handle cleanup
+class ExitCleanup:
+    def __init__(self):
+        self._subprocesses = []
+        signal.signal(signal.SIGTERM, self.exit)
+    
+    def addSubprocess(self, process):
+        self._subprocesses.append(process)
+    
+    def exit(self, signum, frame):
+        for proc in self._subprocesses:
+            print(f"[ExitCleanup] (exit) kill subprocess {proc}")
+            try:
+                if proc.is_running():
+                    proc.terminate()
+                    proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+        raise RuntimeError("Process has been killed")
+
+exitCleanup = ExitCleanup()
+
+
 class MrNodeType(enum.Enum):
     NONE = enum.auto()
     BASENODE = enum.auto()
@@ -196,6 +221,7 @@ def executeChunkCommandLine(self, chunk, cmd, env=None):
                     env=env,
                     **platformArgs,
                 )
+                exitCleanup.addSubprocess(chunk.subprocess)
 
                 if hasattr(chunk, "statThread"):
                     # We only have a statThread if the node is running in the current process
@@ -281,6 +307,8 @@ def processChunkInEnvironment(self, chunk):
 
         if len(chunk.node.getChunks()) > 1:
             meshroomComputeCmd += f" --iteration {chunk.range.iteration}"
+        
+        print(f"(processChunkInEnvironment) meshroomComputeCmd={meshroomComputeCmd}")
 
         runtimeEnv = chunk.node.nodeDesc.plugin.runtimeEnv
         cmdPrefix = chunk.node.nodeDesc.plugin.commandPrefix

meshroom/core/node.py

@@ -720,9 +720,6 @@ def _processInIsolatedEnvironment(self):
         self.node.updateOutputAttr()
 
     def stopProcess(self):
-        if self.isExtern():
-            raise ValueError("Cannot stop process: node is computed externally (another instance of Meshroom)")
-
         # Ensure that we are up-to-date
         self.updateStatusFromCache()
 

meshroom/core/submitter.py

@@ -69,6 +69,12 @@ def __init__(self, jobId, submitter):
     def __repr__(self):
         return f"<{self.__class__.__name__} {self.jid}>"
 
+    def stopChunkTask(self, iteration):
+        if self.submitterOptions.includes(SubmitterOptionsEnum.INTERRUPT_JOB):
+            raise NotImplementedError("'stopChunkTask' method must be implemented in subclasses")
+        else:
+            raise RuntimeError(f"Submitter {self.__class__.__name__} cannot interrupt the job")
+
     def interruptJob(self):
         if self.submitterOptions.includes(SubmitterOptionsEnum.INTERRUPT_JOB):
             raise NotImplementedError("'interruptJob' method must be implemented in subclasses")
@@ -131,6 +137,11 @@ def retreiveJob(self, submitter, jid) -> Optional[BaseSubmittedJob]:
         job = submitter.retrieveJob(jid)
         return job
 
+    def stopChunkTask(self, chunk):
+        print(f"[JobManager] (stopChunkTask) {chunk}")
+        job = self.getNodeJob(chunk.node)
+        job.stopChunkTask(chunk.range.iteration)
+
 
 # Global instance that manages submitted jobs
 jobManager = JobManager()

meshroom/core/taskManager.py

@@ -98,6 +98,10 @@ def run(self):
             for cId, chunk in enumerate(node.chunks):
                 if chunk.isFinishedOrRunning() or not self.isRunning():
                     continue
+
+                if self._manager.isChunkCancelled(chunk):
+                    print(f"[TaskThread] Skip cancelled chunk {chunk}")
+                    continue
 
                 _nodeName, _node, _nbNodes = node.nodeType, nId+1, len(self._manager._nodesToProcess)
 
@@ -148,13 +152,18 @@ def __init__(self, parent: BaseObject = None):
         self._graph = None
         self._nodes = DictModel(keyAttrName='_name', parent=self)
         self._nodesToProcess = []
+        self._cancelledChunks = []
         self._nodesExtern = []
         # internal thread in which local tasks are executed
         self._thread = TaskThread(self)
 
         self._blockRestart = False
         self.restartRequested.connect(self.restart)
 
+    def join(self):
+        self._thread.wait()
+        self._cancelledChunks = []
+
     @Slot(BaseObject)
     def createChunks(self, node: Node):
         """ Create chunks on main process """
@@ -168,6 +177,13 @@ def createChunks(self, node: Node):
             logging.error(f"Failed to create chunks for {node.name}: {e}")
             self.chunksCreated.emit(node)  # Still emit to unblock waiting thread 
 
+    def isChunkCancelled(self, chunk):
+        for i, ch in enumerate(self._cancelledChunks):
+            if ch == chunk:
+                del self._cancelledChunks[i]
+                return True
+        return False
+
     def requestBlockRestart(self):  
         """
         Block computing.
@@ -188,6 +204,7 @@ def blockRestart(self):
 
         self._blockRestart = False
         self._nodesToProcess = []
+        self._cancelledChunks = []
         self._thread._state = State.DEAD
 
     @Slot()
@@ -198,7 +215,7 @@ def restart(self):
         """
         # Make sure to wait the end of the current thread
         if self._thread.isRunning():
-            self._thread.wait()
+            self.join()
 
         # Avoid restart if thread was globally stopped
         if self._blockRestart:

meshroom/ui/graph.py

@@ -28,6 +28,7 @@
 from meshroom.core.graphIO import GraphIO
 
 from meshroom.core.taskManager import TaskManager
+from meshroom.core.submitter import jobManager
 
 from meshroom.core.node import NodeChunk, Node, Status, ExecMode, CompatibilityNode, Position
 from meshroom.core import submitters, MrNodeType
@@ -377,6 +378,7 @@ def __init__(self, undoStack: commands.UndoStack, taskManager: TaskManager, pare
         self._sortedDFSChunks: QObjectListModel = QObjectListModel(parent=self)
         self._layout: GraphLayout = GraphLayout(self)
         self._selectedNode = None
+        self._selectedChunk = None
         self._nodeSelection: QItemSelectionModel = QItemSelectionModel(self._graph.nodes, parent=self)
         self._hoveredNode = None
 
@@ -561,7 +563,7 @@ def stopExecution(self):
             return
         self._taskManager.requestBlockRestart()
         self._graph.stopExecution()
-        self._taskManager._thread.wait()
+        self._taskManager.join()
 
     @Slot(Node)
     def stopNodeComputation(self, node):
@@ -572,7 +574,7 @@ def stopNodeComputation(self, node):
 
         # Stop the node and wait Task Manager
         node.stopComputation()
-        self._taskManager._thread.wait()
+        self._taskManager.join()
 
     @Slot(Node)
     def cancelNodeComputation(self, node):
@@ -608,13 +610,12 @@ def stopTask(self, chunk: NodeChunk):
         if self.isChunkComputingLocally(chunk):
             print(f"-> is local")
             chunk.stopProcess()
-            chunk.upgradeStatusTo(Status.STOPPED)
-            # TODO : remove the chunk from the thread process
-            self._taskManager._thread.wait()
+            # Remove the chunk from the thread process
+            self._taskManager._cancelledChunks.append(chunk)
         elif self.isChunkComputingExternally(chunk):
-            print("[UIGraph] (stopTask) Stop task is not implemented for ")
-            # TODO
-            pass
+            jobManager.stopChunkTask(chunk)
+            chunk.upgradeStatusTo(Status.STOPPED)
+            chunk.stopProcess()
 
     @Slot(NodeChunk)
     def pauseTask(self, chunk: NodeChunk):
@@ -1319,6 +1320,9 @@ def canSubmitNode(self, node: Node) -> bool:
     selectedNodeChanged = Signal()
     # Current main selected node
     selectedNode = makeProperty(QObject, "_selectedNode", selectedNodeChanged, resetOnDestroy=True)
+    # Current chunk selected (used to send signals from TaskManager to ChunksListView)
+    selectedChunkChanged = Signal()
+    selectedChunk = makeProperty(QObject, "_selectedChunk", selectedChunkChanged, resetOnDestroy=True)
 
     nodeSelection = makeProperty(QObject, "_nodeSelection")
 

meshroom/ui/qml/GraphEditor/ChunksListView.qml

@@ -5,11 +5,13 @@ import QtQuick.Layouts
 import Utils 1.0
 
 /**
- * ChunkListView
+ * ChunksListView
  */
 
 ColumnLayout {
     id: root
+
+    property var uigraph: null
     property variant chunks
     property int currentIndex: 0
     property variant currentChunk: (chunks && currentIndex >= 0) ? chunks.at(currentIndex) : undefined
@@ -90,4 +92,17 @@ ColumnLayout {
             }
         }
     }
+
+    Connections {
+        target: _reconstruction
+        function onSelectedChunkChanged() {
+            for (var i = 0; i < root.chunks.count; i++) {
+                if (_reconstruction.selectedChunk === root.chunks.at(i)) {
+                    root.currentIndex = i
+                    break;
+                }
+            }
+        }
+        ignoreUnknownSignals: true
+    }
 }

meshroom/ui/qml/GraphEditor/TaskManager.qml

@@ -28,6 +28,11 @@ Item {
     function selectNode(node) {
         uigraph.selectedNode = node
     }
+
+    function selectChunk(chunk) {
+        root.selectedChunk = chunk
+        uigraph.selectedChunk = chunk
+    }
 
     TextMetrics {
         id: nbMetrics
@@ -234,8 +239,21 @@ Item {
 
                     MouseArea {
                         anchors.fill: parent
-                        onPressed: {
-                            selectNode(object)
+                        acceptedButtons: Qt.LeftButton | Qt.RightButton
+                        onPressed: (mouse) => {
+                            if (mouse.button === Qt.LeftButton) {
+                                selectNode(object)
+                            } else if (mouse.button === Qt.RightButton) {
+                                contextMenu.popup()
+                            }
+                        }
+                        Menu {
+                            id: contextMenu
+                            MenuItem {
+                                text: "Open Folder"
+                                height: visible ? implicitHeight : 0
+                                onTriggered: Qt.openUrlExternally(Filepath.stringToUrl(object.internalFolder))
+                            }
                         }
                     }
                 }
@@ -306,10 +324,6 @@ Item {
                         model: object.chunks
                         property var node: object
 
-                        onModelChanged: {
-                            console.log("TaskManager model size :", chunkList.model.count)
-                        }
-
                         spacing: 3
 
                         delegate: Loader {
@@ -341,7 +355,7 @@ Item {
                                     anchors.fill: parent
                                     onPressed: {
                                         selectNode(chunkList.node)
-                                        selectedChunk = object
+                                        selectChunk(object)
                                     }
                                 }
                             }
@@ -363,7 +377,7 @@ Item {
                                 anchors.fill: parent
                                 onPressed: {
                                     selectNode(chunkList.node)
-                                    selectedChunk = null
+                                    selectChunk(null)
                                 }
                             }
                         }