store lockfile in node-meta.json (#913)

* store lockfile in `node-meta.json`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test lockfile

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* split line

* typo

* test `mp_start_stage_lock` and `mp_join_stage_lock`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* mock `mp.Queue`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PythonFZ

tests/integration/test_node_meta.py

@@ -0,0 +1,68 @@
+import json
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+import zntrack
+from zntrack.utils.lockfile import get_stage_lock
+
+
+class ReadFileContent(zntrack.Node):
+    deps_file: Path = zntrack.deps_path()
+    params: str = zntrack.params()
+
+    def run(self):
+        pass
+
+
+@pytest.fixture()
+def lockfile_01():
+    return {
+        "cmd": "zntrack run test_node_meta.ReadFileContent --name ReadFileContent",
+        "deps": [
+            {
+                "hash": "md5",
+                "md5": "6dbd01b4309de2c22b027eb35a3ce18b",
+                "path": "data.txt",
+                "size": 11,
+            }
+        ],
+        "params": {"params.yaml": {"ReadFileContent": {"params": "test"}}},
+    }
+
+
+def test_node_meta_lock(proj_path, lockfile_01):
+    project = zntrack.Project()
+
+    file = Path("data.txt")
+    file.write_text("Lorem Ipsum")
+
+    with project:
+        _ = ReadFileContent(deps_file=file, params="test")
+
+    project.repro()
+    # TODO: do we want the node to require from_rev()?
+    #   to update lockfile and other node-meta states?
+    node = ReadFileContent.from_rev()
+    assert node.state.lockfile == lockfile_01
+
+
+def test_node_meta_lock_mp(proj_path, lockfile_01):
+    project = zntrack.Project()
+
+    file = Path("data.txt")
+    file.write_text("Lorem Ipsum")
+
+    with project:
+        node = ReadFileContent(deps_file=file, params="test")
+
+    project.repro()
+
+    queue = MagicMock()  # mock multiprocessing.Queue
+    get_stage_lock(node.name, queue)
+    assert queue.put.call_count == 1
+    lockfile = json.loads(json.dumps(queue.put.call_args[0][0]))
+    assert lockfile["cmd"] == lockfile_01["cmd"]
+    assert lockfile["deps"] == lockfile_01["deps"]
+    assert lockfile["params"] == lockfile_01["params"]

zntrack/cli/cli.py

@@ -14,6 +14,7 @@
 from zntrack import Node, utils
 from zntrack.state import PLUGIN_LIST
 from zntrack.utils.import_handler import import_handler
+from zntrack.utils.lockfile import mp_join_stage_lock, mp_start_stage_lock
 from zntrack.utils.misc import load_env_vars
 
 load_env_vars()
@@ -54,7 +55,11 @@ def main(
 
 @app.command()
 def run(
-    node_path: str, name: str = None, meta_only: bool = False, method: str = "run"
+    node_path: str,
+    name: str = None,
+    meta_only: bool = False,
+    method: str = "run",
+    save_lockfile: bool = True,
 ) -> None:
     """Execute a ZnTrack Node.
 
@@ -70,6 +75,8 @@ def run(
         Save only the metadata.
     method : str, default 'run'
         The method to run on the node.
+    save_lockfile : bool
+        Save the lockfile for the inputs into the node-meta.json file.
 
     """
     start_time = datetime.datetime.now()
@@ -78,13 +85,19 @@ def run(
 
     cls: Node = utils.import_handler.import_handler(node_path)
     node: Node = cls.from_rev(name=name, running=True)
+    if save_lockfile:
+        queue, proc = mp_start_stage_lock(name)
     node.state.increment_run_count()
     node.state.save_node_meta()
     # dynamic version of node.run()
     getattr(node, method)()
     node.save()
+    if save_lockfile:
+        stage_lock = mp_join_stage_lock(queue, proc)
+        node.state.set_lockfile(stage_lock)
     run_time = datetime.datetime.now() - start_time
     node.state.add_run_time(run_time)
+
     node.state.save_node_meta()
 
 

zntrack/node.py

@@ -242,12 +242,14 @@ def from_rev(
                 content = json.load(f)
                 run_count = content.get("run_count", 0)
                 run_time = content.get("run_time", 0)
+                lockfile = content.get("lockfile", None)
                 if node_uuid := content.get("uuid", None):
                     instance._uuid = uuid.UUID(node_uuid)
                 instance.__dict__["state"]["run_count"] = run_count
                 instance.__dict__["state"]["run_time"] = datetime.timedelta(
                     seconds=run_time
                 )
+                instance.__dict__["state"]["lockfile"] = lockfile
         if not instance.state.lazy_evaluation:
             for field in dataclasses.fields(cls):
                 _ = getattr(instance, field.name)

zntrack/state.py

@@ -80,6 +80,7 @@ class NodeStatus:
     group: Group | None = None
     run_time: datetime.timedelta | None = None
     path: pathlib.Path = dataclasses.field(default_factory=pathlib.Path)
+    lockfile: dict | None = None
     # TODO: move node name and nwd to here as well
 
     @property
@@ -258,6 +259,10 @@ def add_run_time(self, run_time: datetime.timedelta) -> None:
     def increment_run_count(self) -> None:
         self.node.__dict__["state"]["run_count"] = self.run_count + 1
 
+    def set_lockfile(self, lockfile: dict) -> None:
+        """Set the lockfile for the node."""
+        self.node.__dict__["state"]["lockfile"] = lockfile
+
     def save_node_meta(self) -> None:
         node_meta_content = {
             "uuid": str(self.node.uuid),
@@ -267,6 +272,8 @@ def save_node_meta(self) -> None:
 
         if self.run_time is not None:
             node_meta_content["run_time"] = self.run_time.total_seconds()
+        if self.lockfile is not None:
+            node_meta_content["lockfile"] = self.lockfile
 
         with contextlib.suppress(importlib.metadata.PackageNotFoundError):
             module = self.node.__module__.split(".")[0]

zntrack/utils/lockfile.py

@@ -0,0 +1,26 @@
+import typing as t
+from multiprocessing import Process, Queue
+
+from dvc.api import DVCFileSystem
+from dvc.stage.serialize import to_single_stage_lockfile
+
+
+def get_stage_lock(name: str, queue: Queue) -> None:
+    fs = DVCFileSystem()
+    stage = next(iter(fs.repo.stage.collect(name)))
+    stage.save_deps(allow_missing=False)
+    result = to_single_stage_lockfile(stage)
+    queue.put(result)  # Send the result back to the main process
+
+
+def mp_start_stage_lock(name: str) -> t.Tuple[Queue, Process]:
+    queue = Queue()
+    p = Process(target=get_stage_lock, args=(name, queue))
+    p.start()
+    return queue, p
+
+
+def mp_join_stage_lock(queue: Queue, p: Process) -> dict:
+    p.join()
+    stage_lock = queue.get()  # Receive the result
+    return {k: v for k, v in stage_lock.items() if k in ["cmd", "deps", "params"]}