Skip to content

Commit

Permalink
Add tank_version.json (#1272)
Browse files Browse the repository at this point in the history
  • Loading branch information
monorimet authored Apr 4, 2023
1 parent e0f2276 commit 98b0694
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 10 deletions.
4 changes: 2 additions & 2 deletions conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
17 changes: 9 additions & 8 deletions shark/shark_downloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
3 changes: 3 additions & 0 deletions tank_version.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
{
"version": "2023-03-31_02d52bb"
}

0 comments on commit 98b0694

Please sign in to comment.