From 98b069488ef9161d34dee5e7fa7ea50c5d06dc5b Mon Sep 17 00:00:00 2001
From: Ean Garvey <87458719+monorimet@users.noreply.github.com>
Date: Mon, 3 Apr 2023 20:36:23 -0500
Subject: [PATCH] Add tank_version.json (#1272)

---
 conftest.py               |  4 ++--
 shark/shark_downloader.py | 17 +++++++++--------
 tank_version.json         |  3 +++
 3 files changed, 14 insertions(+), 10 deletions(-)
 create mode 100644 tank_version.json

diff --git a/conftest.py b/conftest.py
index 95442d76..f11bdf0f 100644
--- a/conftest.py
+++ b/conftest.py
@@ -71,8 +71,8 @@ def pytest_addoption(parser):
     parser.addoption(
         "--tank_prefix",
         type=str,
-        default="nightly",
-        help="Prefix to gs://shark_tank/ model directories from which to download SHARK tank artifacts. Default is 'latest'.",
+        default=None,
+        help="Prefix to gs://shark_tank/ model directories from which to download SHARK tank artifacts. Default is nightly.",
     )
     parser.addoption(
         "--benchmark_dispatches",
diff --git a/shark/shark_downloader.py b/shark/shark_downloader.py
index 5c99f021..b87bd3c5 100644
--- a/shark/shark_downloader.py
+++ b/shark/shark_downloader.py
@@ -150,11 +150,14 @@ def get_git_revision_short_hash() -> str:
     if shark_args.shark_prefix is not None:
         prefix_kw = shark_args.shark_prefix
     else:
-        prefix_kw = (
-            subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])
-            .decode("ascii")
-            .strip()
-        )
+        import json
+
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        src = os.path.join(dir_path, "..", "tank_version.json")
+        with open(src, "r") as f:
+            data = json.loads(f.read())
+            prefix_kw = data["version"]
+    print(f"Checking for updates from gs://shark_tank/{prefix_kw}")
     return prefix_kw
 
 
@@ -186,9 +189,6 @@ def get_sharktank_prefix():
     return tank_prefix
 
 
-shark_args.shark_prefix = get_sharktank_prefix()
-
-
 # Downloads the torch model from gs://shark_tank dir.
 def download_model(
     model_name,
@@ -201,6 +201,7 @@ def download_model(
     model_name = model_name.replace("/", "_")
     dyn_str = "_dynamic" if dynamic else ""
     os.makedirs(WORKDIR, exist_ok=True)
+    shark_args.shark_prefix = get_sharktank_prefix()
     if import_args["batch_size"] != 1:
         model_dir_name = (
             model_name
diff --git a/tank_version.json b/tank_version.json
new file mode 100644
index 00000000..6826ad37
--- /dev/null
+++ b/tank_version.json
@@ -0,0 +1,3 @@
+{
+	"version": "2023-03-31_02d52bb"
+}