From 982802effcb3f0cb45d8eaf7907b6156ac9f360d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?=
 <peter@pmrowla.com>
Date: Thu, 4 Feb 2021 16:42:22 +0900
Subject: [PATCH] experiments: bug fixes (#5407)

* exp run: display metrics on -m/--metrics

* exp show: fix sort handling
---
 dvc/command/experiments.py                 | 65 ++++++++++++++++------
 tests/func/experiments/test_experiments.py | 12 ++++
 tests/func/experiments/test_show.py        | 19 +++++++
 3 files changed, 78 insertions(+), 18 deletions(-)

diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py
index 5c3c8b6be..c72c96ea6 100644
--- a/dvc/command/experiments.py
+++ b/dvc/command/experiments.py
@@ -130,14 +130,13 @@ def _collect_rows(
     from dvc.scm.git import Git
 
     if sort_by:
-        if sort_by in metric_names:
-            sort_type = "metrics"
-        elif sort_by in param_names:
-            sort_type = "params"
-        else:
-            raise InvalidArgumentError(f"Unknown sort column '{sort_by}'")
+        sort_path, sort_name, sort_type = _sort_column(
+            sort_by, metric_names, param_names
+        )
         reverse = sort_order == "desc"
-        experiments = _sort_exp(experiments, sort_by, sort_type, reverse)
+        experiments = _sort_exp(
+            experiments, sort_path, sort_name, sort_type, reverse
+        )
 
     new_checkpoint = True
     for i, (rev, exp) in enumerate(experiments.items()):
@@ -193,22 +192,45 @@ def _collect_rows(
         yield row, style
 
 
-def _sort_exp(experiments, sort_by, typ, reverse):
+def _sort_column(sort_by, metric_names, param_names):
+    path, _, sort_name = sort_by.rpartition(":")
+    matches = set()
+
+    if path:
+        if path in metric_names and sort_name in metric_names[path]:
+            matches.add((path, sort_name, "metrics"))
+        if path in param_names and sort_name in param_names[path]:
+            matches.add((path, sort_name, "params"))
+    else:
+        for path in metric_names:
+            if sort_name in metric_names[path]:
+                matches.add((path, sort_name, "metrics"))
+        for path in param_names:
+            if sort_name in param_names[path]:
+                matches.add((path, sort_name, "params"))
+
+    if len(matches) == 1:
+        return matches.pop()
+    if len(matches) > 1:
+        raise InvalidArgumentError(
+            "Ambiguous sort column '{}' matched '{}'".format(
+                sort_by,
+                ", ".join([f"{path}:{name}" for path, name, _ in matches]),
+            )
+        )
+    raise InvalidArgumentError(f"Unknown sort column '{sort_by}'")
+
+
+def _sort_exp(experiments, sort_path, sort_name, typ, reverse):
     def _sort(item):
         rev, exp = item
         tip = exp.get("checkpoint_tip")
         if tip and tip != rev:
             # Sort checkpoint experiments by tip commit
             return _sort((tip, experiments[tip]))
-        for fname, item in exp.get(typ, {}).items():
-            if isinstance(item, dict):
-                item = flatten(item)
-            else:
-                item = {fname: item}
-            if sort_by in item:
-                val = item[sort_by]
-                return (val is None, val)
-        return (True, None)
+        data = exp.get(typ, {}).get(sort_path, {})
+        val = flatten(data).get(sort_name)
+        return (val is None, val)
 
     ret = OrderedDict()
     if "baseline" in experiments:
@@ -472,10 +494,12 @@ class CmdExperimentsDiff(CmdBase):
 
 class CmdExperimentsRun(CmdRepro):
     def run(self):
+        from dvc.command.metrics import _show_metrics
+
         if self.args.reset:
             logger.info("Any existing checkpoints will be reset and re-run.")
 
-        self.repo.experiments.run(
+        results = self.repo.experiments.run(
             name=self.args.name,
             queue=self.args.queue,
             run_all=self.args.run_all,
@@ -487,6 +511,11 @@ class CmdExperimentsRun(CmdRepro):
             **self._repro_kwargs,
         )
 
+        if self.args.metrics and results:
+            metrics = self.repo.metrics.show(revs=list(results))
+            metrics.pop("workspace", None)
+            logger.info(_show_metrics(metrics))
+
         return 0
 
 
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index a10c7ca68..dcbaf30f1 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -581,3 +581,15 @@ def test_queue(tmp_dir, scm, dvc, exp_stage, mocker):
         with tree.open(tmp_dir / "metrics.yaml") as fobj:
             metrics.add(fobj.read().strip())
     assert expected == metrics
+
+
+def test_run_metrics(tmp_dir, scm, dvc, exp_stage, mocker):
+    from dvc.main import main
+
+    mocker.patch.object(
+        dvc.experiments, "run", return_value={"abc123": "abc123"}
+    )
+    show_mock = mocker.patch.object(dvc.metrics, "show", return_value={})
+
+    main(["exp", "run", "-m"])
+    assert show_mock.called_once()
diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py
index 9522aa4cf..8f91ca985 100644
--- a/tests/func/experiments/test_show.py
+++ b/tests/func/experiments/test_show.py
@@ -1,3 +1,4 @@
+import logging
 import os
 from datetime import datetime
 
@@ -214,3 +215,21 @@ def test_show_multiple_commits(tmp_dir, scm, dvc, exp_stage):
 
     results = dvc.experiments.show(num=100)
     assert set(results.keys()) == expected
+
+
+def test_show_sort(tmp_dir, scm, dvc, exp_stage, caplog):
+    with caplog.at_level(logging.ERROR):
+        assert main(["exp", "show", "--no-pager", "--sort-by=bar"]) != 0
+        assert "Unknown sort column" in caplog.text
+
+    with caplog.at_level(logging.ERROR):
+        assert main(["exp", "show", "--no-pager", "--sort-by=foo"]) != 0
+        assert "Ambiguous sort column" in caplog.text
+
+    assert (
+        main(["exp", "show", "--no-pager", "--sort-by=params.yaml:foo"]) == 0
+    )
+
+    assert (
+        main(["exp", "show", "--no-pager", "--sort-by=metrics.yaml:foo"]) == 0
+    )