exp: init: Use persist: true if checkpoint.

hiroto7 · Jul 21, 2022 · ed1d8e2 · ed1d8e2
1 parent 0b07354
commit ed1d8e2
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 9 deletions.
diff --git a/dvc/repo/experiments/init.py b/dvc/repo/experiments/init.py
@@ -253,18 +253,25 @@ def init(
     live_metrics = f"{live_path}.json" if live_path else None
     live_plots = os.path.join(live_path, "scalars") if live_path else None
 
+    if type == "checkpoint":
+        outs_key = "checkpoints"
+        metrics_key = "metrics_persist_no_cache"
+        plots_key = "plots_persist_no_cache"
+    else:
+        outs_key = "outs"
+        metrics_key = "metrics_no_cache"
+        plots_key = "plots_no_cache"
+
     stage = repo.stage.create(
         name=name,
         cmd=context["cmd"],
         deps=compact([context.get("code"), context.get("data")]),
         params=[{params: None}] if params else None,
-        metrics_no_cache=compact([context.get("metrics"), live_metrics]),
-        plots_no_cache=compact([context.get("plots"), live_plots]),
         force=force,
         **{
-            "checkpoints"
-            if type == "checkpoint"
-            else "outs": compact([models])
+            outs_key: compact([models]),
+            metrics_key: compact([context.get("metrics"), live_metrics]),
+            plots_key: compact([context.get("plots"), live_plots]),
         },
     )
 

diff --git a/dvc/stage/utils.py b/dvc/stage/utils.py
@@ -50,10 +50,14 @@ def fill_stage_outputs(stage, **kwargs):
     keys = [
         "outs_persist",
         "outs_persist_no_cache",
-        "metrics_no_cache",
         "metrics",
-        "plots_no_cache",
+        "metrics_persist",
+        "metrics_no_cache",
+        "metrics_persist_no_cache",
         "plots",
+        "plots_persist",
+        "plots_no_cache",
+        "plots_persist_no_cache",
         "outs_no_cache",
         "outs",
         "checkpoints",

diff --git a/tests/func/experiments/test_init.py b/tests/func/experiments/test_init.py
@@ -396,12 +396,12 @@ def test_init_with_type_checkpoint_and_models_plots_provided(
                 "cmd": "cmd",
                 "deps": ["data", "src"],
                 "metrics": [
-                    {"m": {"cache": False}},
+                    {"m": {"cache": False, "persist": True}},
                 ],
                 "outs": [{"models": {"checkpoint": True}}],
                 "params": [{"params.yaml": None}],
                 "plots": [
-                    {"p": {"cache": False}},
+                    {"p": {"cache": False, "persist": True}},
                 ],
             }
         }