From dbec0b0df4b678cf7fd6199344808403f821fb30 Mon Sep 17 00:00:00 2001
From: Egor Sklyarov <egor.sklyarov.ru@gmail.com>
Date: Mon, 5 Feb 2024 11:42:04 +0100
Subject: [PATCH] Make get_latest_runner_build optional. Allow 'latest' as a
 gateway version

---
 .github/workflows/build.yml                   |  1 +
 .github/workflows/release.yml                 |  1 +
 .../_internal/core/backends/base/compute.py   | 24 +++++++++++++------
 src/dstack/_internal/settings.py              |  1 +
 4 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 9fed9c214..afa038d51 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -207,3 +207,4 @@ jobs:
         run: |
           WHEEL=dstack_gateway-${{ env.VERSION }}-py3-none-any.whl
           aws s3 cp dist/$WHEEL "s3://dstack-gateway-downloads/stgn/$WHEEL"
+          echo "${{ env.VERSION }}" > aws s3 cp - "s3://dstack-gateway-downloads/stgn/latest-version"
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 944d16687..94c10ade0 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -248,3 +248,4 @@ jobs:
         run: |
           WHEEL=dstack_gateway-${{ env.VERSION }}-py3-none-any.whl
           aws s3 cp dist/$WHEEL "s3://dstack-gateway-downloads/release/$WHEEL"
+          echo "${{ env.VERSION }}" > aws s3 cp - "s3://dstack-gateway-downloads/release/latest-version"
diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py
index e51fd0999..bd8db3bfd 100644
--- a/src/dstack/_internal/core/backends/base/compute.py
+++ b/src/dstack/_internal/core/backends/base/compute.py
@@ -1,13 +1,13 @@
 import os
 import re
 from abc import ABC, abstractmethod
+from functools import cache
 from typing import List, Optional
 
 import git
 import requests
 import yaml
 
-from dstack import version
 from dstack._internal import settings
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
@@ -104,7 +104,10 @@ def get_shim_commands(
 def get_dstack_runner_version() -> str:
     if settings.DSTACK_VERSION is not None:
         return settings.DSTACK_VERSION
-    return os.environ.get("DSTACK_RUNNER_VERSION", None) or get_latest_runner_build() or "latest"
+    version = os.environ.get("DSTACK_RUNNER_VERSION", None)
+    if version is None and settings.DSTACK_:
+        version = get_latest_runner_build()
+    return version or "latest"
 
 
 def get_cloud_config(**config) -> str:
@@ -193,12 +196,16 @@ def get_docker_commands(authorized_keys: List[str]) -> List[str]:
     return commands
 
 
+@cache  # Restart the server to find the latest build
 def get_latest_runner_build() -> Optional[str]:
     owner_repo = "dstackai/dstack"
     workflow_id = "build.yml"
     version_offset = 150
 
-    repo = git.Repo(os.path.abspath(os.path.dirname(__file__)), search_parent_directories=True)
+    try:
+        repo = git.Repo(os.path.abspath(os.path.dirname(__file__)), search_parent_directories=True)
+    except git.InvalidGitRepositoryError:
+        return None
     for remote in repo.remotes:
         if re.search(rf"[@/]github\.com[:/]{owner_repo}\.", remote.url):
             break
@@ -222,7 +229,7 @@ def get_latest_runner_build() -> Optional[str]:
         try:
             if repo.is_ancestor(run["head_sha"], head):
                 ver = str(run["run_number"] + version_offset)
-                logger.debug(f"Found the latest runner build: %s", ver)
+                logger.debug("Found the latest runner build: %s", ver)
                 return ver
         except git.GitCommandError as e:
             if "Not a valid commit name" not in e.stderr:
@@ -232,13 +239,16 @@ def get_latest_runner_build() -> Optional[str]:
 
 def get_dstack_gateway_wheel(build: str) -> str:
     channel = "release" if settings.DSTACK_RELEASE else "stgn"
-    return f"https://dstack-gateway-downloads.s3.amazonaws.com/{channel}/dstack_gateway-{build}-py3-none-any.whl"
+    base_url = f"https://dstack-gateway-downloads.s3.amazonaws.com/{channel}"
+    if build == "latest":
+        r = requests.get(f"{base_url}/latest-version")
+        r.raise_for_status()
+        build = r.text.strip()
+    return f"{base_url}/dstack_gateway-{build}-py3-none-any.whl"
 
 
 def get_dstack_gateway_commands() -> List[str]:
     build = get_dstack_runner_version()
-    if build == "latest":
-        raise ValueError("`latest` is not appropriate version for a gateway")
     return [
         "mkdir -p /home/ubuntu/dstack",
         "python3 -m venv /home/ubuntu/dstack/blue",
diff --git a/src/dstack/_internal/settings.py b/src/dstack/_internal/settings.py
index 9247e7ba5..9d0845292 100644
--- a/src/dstack/_internal/settings.py
+++ b/src/dstack/_internal/settings.py
@@ -4,3 +4,4 @@
 
 DSTACK_VERSION = os.getenv("DSTACK_VERSION", version.__version__)
 DSTACK_RELEASE = os.getenv("DSTACK_RELEASE") is not None or version.__is_release__
+DSTACK_USE_LATEST_FROM_BRANCH = os.getenv("DSTACK_USE_LATEST_FROM_BRANCH") is not None