Compute and store node cache folder size in node status; display in Node Editor header

meshroom/core/node.py

@@ -66,7 +66,8 @@ class ExecMode(Enum):
 
 class NodeStatusData(BaseObject):
     __slots__ = ("nodeName", "nodeType", "status", "execMode", "packageName", "mrNodeType",
-                 "submitterSessionUid", "chunksBlockSize", "chunksFullSize", "chunksNbBlocks", "jobInfo")
+                 "submitterSessionUid", "chunksBlockSize", "chunksFullSize", "chunksNbBlocks", "jobInfo",
+                 "storageSizeKB")
 
     def __init__(self, nodeName='', nodeType='', packageName='',
                  mrNodeType: MrNodeType = MrNodeType.NONE, parent: BaseObject = None):
@@ -92,6 +93,7 @@ def resetDynamicValues(self):
         self.status: Status = Status.NONE
         self.execMode: ExecMode = ExecMode.NONE
         self.jobInfo: dict = {}
+        self.storageSizeKB: int = 0
 
     def setNodeType(self, node):
         """
@@ -676,6 +678,12 @@ def process(self, forceCompute=False, inCurrentEnv=False):
 
             if executionStatus:
                 self.upgradeStatusTo(executionStatus)
+            # After the chunk status is finalized, compute and persist the storage size
+            # if the node has reached a terminal state
+            if self.node.getGlobalStatus() in (Status.SUCCESS, Status.ERROR, Status.STOPPED, Status.KILLED):
+                self.node._nodeStatus.storageSizeKB = self.node._computeCacheFolderSize()
+                self.node.upgradeStatusFile()
+                self.node.nodeStatusChanged.emit()
             logging.info(f"[Process chunk] elapsed time: {self._status.elapsedTimeStr}")
             # Ask and wait for the stats thread to stop
             self.statThread.stopRequest()
@@ -1623,6 +1631,20 @@ def upgradeStatusFile(self):
             json.dump(data, jsonFile, indent=4)
         renameWritingToFinalPath(statusFilepathWriting, statusFilepath)
 
+    def _computeCacheFolderSize(self):
+        """ Compute the total size of all files in the node's internal cache folder in KB. """
+        totalSizeBytes = 0
+        folder = self.internalFolder
+        if folder and os.path.exists(folder):
+            for dirpath, dirnames, filenames in os.walk(folder):
+                for filename in filenames:
+                    filepath = os.path.join(dirpath, filename)
+                    try:
+                        totalSizeBytes += os.path.getsize(filepath)
+                    except OSError:
+                        logging.warning(f"Could not get size of file: {filepath}")
+        return totalSizeBytes // 1024
+
     def setJobId(self, jid, submitterName):
         self._nodeStatus.setJob(jid, submitterName)
         self.upgradeStatusFile()
@@ -2132,6 +2154,7 @@ def _hasDisplayableShape(self):
     elapsedTime = Property(float, lambda self: self.getFusedStatus().elapsedTime, notify=globalStatusChanged)
     recursiveElapsedTime = Property(float, lambda self: self.getRecursiveFusedStatus().elapsedTime,
                                     notify=globalStatusChanged)
+    storageSizeKB = Property(int, lambda self: self._nodeStatus.storageSizeKB, notify=nodeStatusChanged)
     isCompatibilityNode = Property(bool, lambda self: self._isCompatibilityNode(), constant=True)
     isInputNode = Property(bool, lambda self: self._isInputNode(), constant=True)
     isBackdropNode = Property(bool, lambda self: self._isBackdropNode(), constant=True)

meshroom/ui/qml/GraphEditor/NodeEditor.qml

@@ -270,6 +270,15 @@ Panel {
             }
         }
 
+        Label {
+            id: storageSizeInfo
+            color: node && node.isComputableType ? Colors.statusColors[node.globalStatus] : palette.text
+            padding: 2
+            font.italic: true
+            visible: node !== null && node.isComputableType && node.storageSizeKB > 0
+            text: node !== null ? Format.KB2SizeStr(node.storageSizeKB) : ""
+        }
+
         SearchBar {
             id: searchBar
             toggle: true  // Enable toggling the actual text field by the search button

meshroom/ui/qml/Utils/format.js

@@ -102,3 +102,17 @@ function GB2SizeStr(GB) {
     }
     return sizeStr
 }
+
+function KB2SizeStr(KB) {
+    // Convert KB to a human-readable size string
+    // e.g. 1.5G, 45M, 234K
+    var MB = KB / 1024
+    var GB = MB / 1024
+    if (GB >= 1) {
+        return parseFloat(GB.toFixed(1)) + "G"
+    } else if (MB >= 1) {
+        return parseFloat(MB.toFixed(1)) + "M"
+    } else {
+        return Math.round(KB) + "K"
+    }
+}

tests/test_compute.py

@@ -413,3 +413,26 @@ def test_noUpstreamLockAfterParameterChange(self, graphSavedOnDisk):
         downstreamNode.input.disconnectEdge()
         self.checkNodeStatusAndLock(node, Status.SUCCESS, False)
         self.checkNodeStatusAndLock(downstreamNode, Status.NONE, False)
+
+    def test_storageSizeKBAfterComputation(self, graphSavedOnDisk):
+        """
+        Test that the storageSizeKB field is populated in the node status after computation.
+        """
+        import json
+
+        graph: Graph = graphSavedOnDisk
+        node = graph.addNewNode("PluginANodeA")
+        graph.save()
+
+        assert node._nodeStatus.storageSizeKB == 0.0
+
+        node.process(inCurrentEnv=True)
+        self.checkNodeStatusAndLock(node, Status.SUCCESS, False)
+
+        # The storageSizeKB should be greater than 0 after computation (status files were written)
+        assert node._nodeStatus.storageSizeKB > 0.0
+
+        # The storageSizeKB should also be persisted in the node status file
+        with open(node.nodeStatusFile) as f:
+            data = json.load(f)
+        assert data.get("storageSizeKB", 0.0) > 0.0