[WB-3544] Apply diff.patch when calling wandb.restore (wandb#1339)

vanpelt · raubitsj · web-flow · commit 15422ae5c891 · 2020-10-09T13:02:21.000-07:00
* wow, this has been busted since gorilla was released...

* Attempt to fix windows

* Just skip the test in Windows

Co-authored-by: Jeff Raubitschek &lt;jeff@wandb.com&gt;
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -667,6 +667,7 @@ def test_local_already_running(runner, docker, local_settings):
     assert "A container named wandb-local is already running" in result.output
 
 
+@pytest.mark.skipif(platform.system() == "Windows", reason="The patch in mock_server.py doesn't work in windows")
 def test_restore_no_remote(runner, mock_server, git_repo, docker, monkeypatch):
     with open("patch.txt", "w") as f:
         f.write("test")
diff --git a/tests/utils/mock_server.py b/tests/utils/mock_server.py
@@ -199,22 +199,20 @@ def set_ctx(ctx):
 
 def _bucket_config():
     return {
-        'patch': '''
-diff --git a/patch.txt b/patch.txt
-index 30d74d2..9a2c773 100644
---- a/patch.txt
-+++ b/patch.txt
-@@ -1 +1 @@
--test
-\ No newline at end of file
-+testing
-\ No newline at end of file
-        ''',
         'commit': 'HEAD',
         'github': 'https://github.com/vanpelt',
         'config': '{"foo":{"value":"bar"}}',
         'files': {
-            'edges': [{'node': {'directUrl': request.url_root + "/storage?file=metadata.json"}}]
+            'edges': [
+                {'node': {
+                    'directUrl': request.url_root + "/storage?file=wandb-metadata.json",
+                    'name': 'wandb-metadata.json'
+                }},
+                {'node': {
+                    'directUrl': request.url_root + "/storage?file=diff.patch",
+                    'name': 'diff.patch'
+                }}
+            ]
         }
     }
 
@@ -629,7 +627,7 @@ def storage():
         if request.method == "GET" and size:
             return os.urandom(size), 200
         # make sure to read the data
-        data = request.get_data()
+        request.get_data()
         if file == "wandb_manifest.json":
             return {
                 "version": 1,
@@ -639,8 +637,27 @@ def storage():
                     "digits.h5": {"digest": "TeSJ4xxXg0ohuL5xEdq2Ew==", "size": 81299},
                 },
             }
-        elif file == "metadata.json":
-            return {"docker": "test/docker", "program": "train.py", "args": ["--test", "foo"], "git": ctx.get("git", {})}
+        elif file == "wandb-metadata.json":
+            return {
+                "docker": "test/docker",
+                "program": "train.py",
+                "args": ["--test", "foo"],
+                "git": ctx.get("git", {})
+            }
+        elif file == "diff.patch":
+            # TODO: make sure the patch is valid for windows as well,
+            # and un skip the test in test_cli.py
+            return '''
+diff --git a/patch.txt b/patch.txt
+index 30d74d2..9a2c773 100644
+--- a/patch.txt
++++ b/patch.txt
+@@ -1 +1 @@
+-test
+\ No newline at end of file
++testing
+\ No newline at end of file
+'''
         return "", 200
 
     @app.route("/artifacts/<entity>/<digest>", methods=["GET", "POST"])
diff --git a/wandb/cli/cli.py b/wandb/cli/cli.py
@@ -1306,7 +1306,6 @@ def restore(ctx, run, no_git, branch, project, entity):
     commit, json_config, patch_content, metadata = api.run_config(
         project, run=run, entity=entity
     )
-    print(metadata)
     repo = metadata.get("git", {}).get("repo")
     image = metadata.get("docker")
     restore_message = (
@@ -1325,6 +1324,7 @@ def restore(ctx, run, no_git, branch, project, entity):
             )
 
     if commit and api.git.enabled:
+        wandb.termlog("Fetching origin and finding commit: {}".format(commit))
         subprocess.check_call(["git", "fetch", "--all"])
         try:
             api.git.repo.commit(commit)
@@ -1378,10 +1378,15 @@ def restore(ctx, run, no_git, branch, project, entity):
             patch_rel_path = os.path.relpath(patch_path, start=root)
             # --reject is necessary or else this fails any time a binary file
             # occurs in the diff
-            # we use .call() instead of .check_call() for the same reason
-            # TODO(adrian): this means there is no error checking here
-            subprocess.call(["git", "apply", "--reject", patch_rel_path], cwd=root)
-            wandb.termlog("Applied patch")
+            exit_code = subprocess.call(
+                ["git", "apply", "--reject", patch_rel_path], cwd=root
+            )
+            if exit_code == 0:
+                wandb.termlog("Applied patch")
+            else:
+                wandb.termerror(
+                    "Failed to apply patch, try un-staging any un-committed changes"
+                )
 
     util.mkdir_exists_ok(wandb_dir())
     config_path = os.path.join(wandb_dir(), "config.yaml")
diff --git a/wandb/internal/internal_api.py b/wandb/internal/internal_api.py
@@ -661,10 +661,10 @@ def run_config(self, project, run=None, entity=None):
                 bucket(name: $run) {
                     config
                     commit
-                    patch
-                    files(names: ["wandb-metadata.json"]) {
+                    files(names: ["wandb-metadata.json", "diff.patch"]) {
                         edges {
                             node {
+                                name
                                 directUrl
                             }
                         }
@@ -682,15 +682,19 @@ def run_config(self, project, run=None, entity=None):
             raise CommError("Run {}/{}/{} not found".format(entity, project, run))
         run = response["model"]["bucket"]
         commit = run["commit"]
-        patch = run["patch"]
         config = json.loads(run["config"] or "{}")
+        patch = None
+        metadata = {}
         if len(run["files"]["edges"]) > 0:
-            url = run["files"]["edges"][0]["node"]["directUrl"]
-            res = requests.get(url)
-            res.raise_for_status()
-            metadata = res.json()
-        else:
-            metadata = {}
+            for file_edge in run["files"]["edges"]:
+                name = file_edge["node"]["name"]
+                url = file_edge["node"]["directUrl"]
+                res = requests.get(url)
+                res.raise_for_status()
+                if name == "wandb-metadata.json":
+                    metadata = res.json()
+                elif name == "diff.patch":
+                    patch = res.text
         return (commit, config, patch, metadata)
 
     @normalize_exceptions
