From e93e8f05a38efcea4cea56da14aa19ede4c72711 Mon Sep 17 00:00:00 2001
From: skshetry <18718008+skshetry@users.noreply.github.com>
Date: Tue, 2 May 2023 08:46:10 +0545
Subject: [PATCH] exp push: notify studio even if cache push fails (#9380)

* exp push: notify Studio even when cache push fails

* exp push: refactor and extract exp_ref resolving logic as the function is getting complex
---
 dvc/repo/experiments/push.py | 68 +++++++++++++++++++++---------------
 1 file changed, 40 insertions(+), 28 deletions(-)

diff --git a/dvc/repo/experiments/push.py b/dvc/repo/experiments/push.py
index 6e6c59fda..7e0913a68 100644
--- a/dvc/repo/experiments/push.py
+++ b/dvc/repo/experiments/push.py
@@ -20,6 +20,7 @@ from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
 from dvc.scm import Git, TqdmGit, iter_revs
 from dvc.utils import env2bool
+from dvc.utils.collections import ensure_list
 
 from .exceptions import UnresolvedExpNamesError
 from .refs import ExpRefInfo
@@ -66,12 +67,37 @@ def notify_refs_to_studio(
     return d.get("url")
 
 
+def exp_refs_from_names(scm: "Git", exp_names: List[str]) -> Set["ExpRefInfo"]:
+    exp_ref_set = set()
+    exp_ref_dict = resolve_name(scm, exp_names)
+    unresolved_exp_names = []
+    for exp_name, exp_ref in exp_ref_dict.items():
+        if exp_ref is None:
+            unresolved_exp_names.append(exp_name)
+        else:
+            exp_ref_set.add(exp_ref)
+
+    if unresolved_exp_names:
+        raise UnresolvedExpNamesError(unresolved_exp_names)
+    return exp_ref_set
+
+
+def exp_refs_from_rev(scm: "Git", rev: str, num: int = 1) -> Set["ExpRefInfo"]:
+    exp_ref_set = set()
+    rev_dict = iter_revs(scm, [rev], num)
+    rev_set = set(rev_dict.keys())
+    ref_info_dict = exp_refs_by_baseline(scm, rev_set)
+    for _, ref_info_list in ref_info_dict.items():
+        exp_ref_set.update(ref_info_list)
+    return exp_ref_set
+
+
 @locked
 @scm_context
-def push(  # noqa: C901
+def push(
     repo: "Repo",
     git_remote: str,
-    exp_names: Union[Iterable[str], str],
+    exp_names: Union[List[str], str],
     all_commits: bool = False,
     rev: Optional[str] = None,
     num: int = 1,
@@ -83,29 +109,10 @@ def push(  # noqa: C901
     assert isinstance(repo.scm, Git)
     if all_commits:
         exp_ref_set.update(exp_refs(repo.scm))
-
-    else:
-        if exp_names:
-            if isinstance(exp_names, str):
-                exp_names = [exp_names]
-            exp_ref_dict = resolve_name(repo.scm, exp_names)
-
-            unresolved_exp_names = []
-            for exp_name, exp_ref in exp_ref_dict.items():
-                if exp_ref is None:
-                    unresolved_exp_names.append(exp_name)
-                else:
-                    exp_ref_set.add(exp_ref)
-
-            if unresolved_exp_names:
-                raise UnresolvedExpNamesError(unresolved_exp_names)
-
-        if rev:
-            rev_dict = iter_revs(repo.scm, [rev], num)
-            rev_set = set(rev_dict.keys())
-            ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set)
-            for _, ref_info_list in ref_info_dict.items():
-                exp_ref_set.update(ref_info_list)
+    if exp_names:
+        exp_ref_set.update(exp_refs_from_names(repo.scm, ensure_list(exp_names)))
+    if rev:
+        exp_ref_set.update(exp_refs_from_rev(repo.scm, rev, num=num))
 
     push_result = _push(repo, git_remote, exp_ref_set, force)
 
@@ -118,15 +125,20 @@ def push(  # noqa: C901
     pushed_refs_info = (
         push_result[SyncStatus.UP_TO_DATE] + push_result[SyncStatus.SUCCESS]
     )
+
+    e = None
     if push_cache:
         try:
             result["uploaded"] = _push_cache(repo, pushed_refs_info, **kwargs)
         except Exception as exc:  # noqa: BLE001, pylint: disable=broad-except
-            raise UploadError("failed to push cache", result) from exc
+            e = exc
 
     pushed_refs = [str(r) for r in pushed_refs_info]
-    url = notify_refs_to_studio(repo, git_remote, pushed=pushed_refs)
-    return {**result, "url": url}
+    result["url"] = notify_refs_to_studio(repo, git_remote, pushed=pushed_refs)
+
+    if e:
+        raise UploadError("failed to push cache", result) from e
+    return result
 
 
 def _push(