Improved remote repo support (#3279)

Author: peterschmidt85

runner/internal/repo/manager.go

@@ -16,20 +16,26 @@ type Manager struct {
 	ctx       context.Context
 	localPath string
 	clo       git.CloneOptions
+	branch    string
 	hash      string
 }
 
 func NewManager(ctx context.Context, url, branch, hash string, singleBranch bool) *Manager {
 	ctx = log.AppendArgsCtx(ctx, "url", url, "branch", branch, "hash", hash)
 	m := &Manager{
-		ctx: ctx,
+		ctx:    ctx,
+		branch: branch,
+		hash:   hash,
 		clo: git.CloneOptions{
 			URL:               url,
 			RecurseSubmodules: git.DefaultSubmoduleRecursionDepth,
-			ReferenceName:     plumbing.NewBranchReferenceName(branch),
 			SingleBranch:      singleBranch,
 		},
-		hash: hash,
+	}
+	// Only set ReferenceName if branch is non-empty
+	// If empty, it will default to HEAD in CloneOptions.Validate()
+	if branch != "" {
+		m.clo.ReferenceName = plumbing.NewBranchReferenceName(branch)
 	}
 
 	return m
@@ -69,23 +75,41 @@ func (m *Manager) Checkout() error {
 		return fmt.Errorf("clone repo: %w", err)
 	}
 	if ref != nil {
-		branchRef, err := ref.Reference(m.clo.ReferenceName, true)
-		if err != nil {
-			return fmt.Errorf("get branch reference: %w", err)
-		}
 		var cho git.CheckoutOptions
-		if m.hash == "" || m.hash == branchRef.Hash().String() {
-			cho.Branch = m.clo.ReferenceName
+		needCheckout := false
+
+		if m.branch != "" {
+			branchRef, err := ref.Reference(m.clo.ReferenceName, true)
+			if err != nil {
+				return fmt.Errorf("get branch reference: %w", err)
+			}
+			if m.hash == "" || m.hash == branchRef.Hash().String() {
+				// Hash is empty or matches branch head: checkout branch
+				cho.Branch = m.clo.ReferenceName
+				needCheckout = true
+			} else {
+				// Hash is specified and different: checkout by hash
+				cho.Hash = plumbing.NewHash(m.hash)
+				needCheckout = true
+			}
 		} else {
-			cho.Hash = plumbing.NewHash(m.hash)
+			// Branch is empty: checkout by hash if specified, otherwise HEAD is already checked out
+			if m.hash != "" {
+				cho.Hash = plumbing.NewHash(m.hash)
+				needCheckout = true
+			}
+			// If hash is also empty, HEAD is already checked out by clone, no need to checkout again
 		}
-		workTree, err := ref.Worktree()
-		if err != nil {
-			return fmt.Errorf("get worktree: %w", err)
-		}
-		err = workTree.Checkout(&cho)
-		if err != nil {
-			return fmt.Errorf("checkout: %w", err)
+
+		if needCheckout {
+			workTree, err := ref.Worktree()
+			if err != nil {
+				return fmt.Errorf("get worktree: %w", err)
+			}
+			err = workTree.Checkout(&cho)
+			if err != nil {
+				return fmt.Errorf("checkout: %w", err)
+			}
 		}
 	} else {
 		log.Warning(m.ctx, "git clone ref==nil")

src/dstack/_internal/cli/services/configurators/run.py

@@ -54,7 +54,7 @@
 from dstack._internal.core.services.diff import diff_models
 from dstack._internal.core.services.repos import (
     InvalidRepoCredentialsError,
-    get_repo_creds_and_default_branch,
+    get_repo_creds,
     load_repo,
 )
 from dstack._internal.utils.common import local_time
@@ -562,8 +562,7 @@ def get_repo(
             local_path = Path.cwd()
             legacy_local_path = True
         if url:
-            # "master" is a dummy value, we'll fetch the actual default branch later
-            repo = RemoteRepo.from_url(repo_url=url, repo_branch="master")
+            repo = RemoteRepo.from_url(repo_url=url)
             repo_head = self.api.repos.get(repo_id=repo.repo_id, with_creds=True)
         elif local_path:
             if legacy_local_path:
@@ -618,7 +617,7 @@ def get_repo(
                 init = True
 
             try:
-                repo_creds, default_repo_branch = get_repo_creds_and_default_branch(
+                repo_creds = get_repo_creds(
                     repo_url=repo.repo_url,
                     identity_file=git_identity_file,
                     private_key=git_private_key,
@@ -627,16 +626,11 @@ def get_repo(
             except InvalidRepoCredentialsError as e:
                 raise CLIError(*e.args) from e
 
-            if repo_branch is None and repo_hash is None:
-                repo_branch = default_repo_branch
-                if repo_branch is None:
-                    raise CLIError(
-                        "Failed to automatically detect remote repo branch."
-                        " Specify branch or hash."
-                    )
-            repo = RemoteRepo.from_url(
-                repo_url=repo.repo_url, repo_branch=repo_branch, repo_hash=repo_hash
-            )
+            # repo_branch and repo_hash are taken from the repo_spec
+            if repo_branch is not None:
+                repo.run_repo_data.repo_branch = repo_branch
+            if repo_hash is not None:
+                repo.run_repo_data.repo_hash = repo_hash
 
         if init:
             self.api.repos.init(

src/dstack/_internal/core/models/repos/remote.py

@@ -135,14 +135,12 @@ def from_url(
 
         Args:
             repo_url: The URL of a remote Git repo.
-            repo_branch: The name of the remote branch. Must be specified if `hash` is not specified.
-            repo_hash: The hash of the revision. Must be specified if `branch` is not specified.
+            repo_branch: The name of the remote branch.
+            repo_hash: The hash of the revision.
 
         Returns:
             A remote repo instance.
         """
-        if repo_branch is None and repo_hash is None:
-            raise ValueError("Either `repo_branch` or `repo_hash` must be specified.")
         return RemoteRepo(
             repo_url=repo_url,
             repo_branch=repo_branch,

src/dstack/_internal/core/services/repos.py

@@ -26,86 +26,74 @@ class InvalidRepoCredentialsError(DstackError):
     pass
 
 
-def get_repo_creds_and_default_branch(
+def get_repo_creds(
     repo_url: str,
     identity_file: Optional[PathLike] = None,
     private_key: Optional[str] = None,
     oauth_token: Optional[str] = None,
-) -> tuple[RemoteRepoCreds, Optional[str]]:
+) -> RemoteRepoCreds:
     url = GitRepoURL.parse(repo_url, get_ssh_config=get_host_config)
 
     # no auth
     with suppress(InvalidRepoCredentialsError):
-        creds, default_branch = _get_repo_creds_and_default_branch_https(url)
-        logger.debug(
-            "Git repo %s is public. Using no auth. Default branch: %s", repo_url, default_branch
-        )
-        return creds, default_branch
+        creds = _get_repo_creds_https(url)
+        logger.debug("Git repo %s is public. Using no auth.", repo_url)
+        return creds
 
     # ssh key provided by the user or pulled from the server
     if identity_file is not None or private_key is not None:
         if identity_file is not None:
             private_key = _read_private_key(identity_file)
-            creds, default_branch = _get_repo_creds_and_default_branch_ssh(
-                url, identity_file, private_key
-            )
+            creds = _get_repo_creds_ssh(url, identity_file, private_key)
             logger.debug(
-                "Git repo %s is private. Using identity file: %s. Default branch: %s",
+                "Git repo %s is private. Using identity file: %s.",
                 repo_url,
                 identity_file,
-                default_branch,
             )
-            return creds, default_branch
+            return creds
         elif private_key is not None:
             with NamedTemporaryFile("w+", 0o600) as f:
                 f.write(private_key)
                 f.flush()
-                creds, default_branch = _get_repo_creds_and_default_branch_ssh(
-                    url, f.name, private_key
-                )
+                creds = _get_repo_creds_ssh(url, f.name, private_key)
                 masked_key = "***" + private_key[-10:] if len(private_key) > 10 else "***MASKED***"
                 logger.debug(
                     "Git repo %s is private. Using private key: %s. Default branch: %s",
                     repo_url,
                     masked_key,
-                    default_branch,
                 )
-                return creds, default_branch
+                return creds
         else:
             assert False, "should not reach here"
 
     # oauth token provided by the user or pulled from the server
     if oauth_token is not None:
-        creds, default_branch = _get_repo_creds_and_default_branch_https(url, oauth_token)
+        creds = _get_repo_creds_https(url, oauth_token)
         masked_token = (
             len(oauth_token[:-4]) * "*" + oauth_token[-4:]
             if len(oauth_token) > 4
             else "***MASKED***"
         )
         logger.debug(
-            "Git repo %s is private. Using provided OAuth token: %s. Default branch: %s",
+            "Git repo %s is private. Using provided OAuth token: %s.",
             repo_url,
             masked_token,
-            default_branch,
         )
-        return creds, default_branch
+        return creds
 
     # key from ssh config
     identities = get_host_config(url.original_host).get("identityfile")
     if identities:
         _identity_file = identities[0]
         with suppress(InvalidRepoCredentialsError):
             _private_key = _read_private_key(_identity_file)
-            creds, default_branch = _get_repo_creds_and_default_branch_ssh(
-                url, _identity_file, _private_key
-            )
+            creds = _get_repo_creds_ssh(url, _identity_file, _private_key)
             logger.debug(
-                "Git repo %s is private. Using SSH config identity file: %s. Default branch: %s",
+                "Git repo %s is private. Using SSH config identity file: %s.",
                 repo_url,
                 _identity_file,
-                default_branch,
             )
-            return creds, default_branch
+            return creds
 
     # token from gh config
     if os.path.exists(gh_config_path):
@@ -114,48 +102,44 @@ def get_repo_creds_and_default_branch(
         _oauth_token = gh_hosts.get(url.host, {}).get("oauth_token")
         if _oauth_token is not None:
             with suppress(InvalidRepoCredentialsError):
-                creds, default_branch = _get_repo_creds_and_default_branch_https(url, _oauth_token)
+                creds = _get_repo_creds_https(url, _oauth_token)
                 masked_token = (
                     len(_oauth_token[:-4]) * "*" + _oauth_token[-4:]
                     if len(_oauth_token) > 4
                     else "***MASKED***"
                 )
                 logger.debug(
-                    "Git repo %s is private. Using GitHub config token: %s from %s. Default branch: %s",
+                    "Git repo %s is private. Using GitHub config token: %s from %s.",
                     repo_url,
                     masked_token,
                     gh_config_path,
-                    default_branch,
                 )
-                return creds, default_branch
+                return creds
 
     # default user key
     if os.path.exists(default_ssh_key):
         with suppress(InvalidRepoCredentialsError):
             _private_key = _read_private_key(default_ssh_key)
-            creds, default_branch = _get_repo_creds_and_default_branch_ssh(
-                url, default_ssh_key, _private_key
-            )
+            creds = _get_repo_creds_ssh(url, default_ssh_key, _private_key)
             logger.debug(
-                "Git repo %s is private. Using default identity file: %s. Default branch: %s",
+                "Git repo %s is private. Using default identity file: %s.",
                 repo_url,
                 default_ssh_key,
-                default_branch,
             )
-            return creds, default_branch
+            return creds
 
     raise InvalidRepoCredentialsError(
         "No valid default Git credentials found. Pass valid `--token` or `--git-identity`."
     )
 
 
-def _get_repo_creds_and_default_branch_ssh(
+def _get_repo_creds_ssh(
     url: GitRepoURL, identity_file: PathLike, private_key: str
-) -> tuple[RemoteRepoCreds, Optional[str]]:
+) -> RemoteRepoCreds:
     _url = url.as_ssh()
     env = _make_git_env_for_creds_check(identity_file=identity_file)
     try:
-        default_branch = _get_repo_default_branch(_url, env)
+        _get_repo_default_branch(_url, env)
     except GitCommandError as e:
         message = f"Cannot access `{_url}` using the `{identity_file}` private SSH key"
         raise InvalidRepoCredentialsError(message) from e
@@ -164,16 +148,14 @@ def _get_repo_creds_and_default_branch_ssh(
         private_key=private_key,
         oauth_token=None,
     )
-    return creds, default_branch
+    return creds
 
 
-def _get_repo_creds_and_default_branch_https(
-    url: GitRepoURL, oauth_token: Optional[str] = None
-) -> tuple[RemoteRepoCreds, Optional[str]]:
+def _get_repo_creds_https(url: GitRepoURL, oauth_token: Optional[str] = None) -> RemoteRepoCreds:
     _url = url.as_https()
     env = _make_git_env_for_creds_check()
     try:
-        default_branch = _get_repo_default_branch(url.as_https(oauth_token), env)
+        _get_repo_default_branch(url.as_https(oauth_token), env)
     except GitCommandError as e:
         message = f"Cannot access `{_url}`"
         if oauth_token is not None:
@@ -185,7 +167,7 @@ def _get_repo_creds_and_default_branch_https(
         private_key=None,
         oauth_token=oauth_token,
     )
-    return creds, default_branch
+    return creds
 
 
 def _make_git_env_for_creds_check(identity_file: Optional[PathLike] = None) -> dict[str, str]:

src/dstack/api/_public/repos.py

@@ -14,7 +14,7 @@
 from dstack._internal.core.services.configs import ConfigManager
 from dstack._internal.core.services.repos import (
     InvalidRepoCredentialsError,
-    get_repo_creds_and_default_branch,
+    get_repo_creds,
     load_repo,
 )
 from dstack._internal.utils.logging import get_logger
@@ -76,7 +76,7 @@ def init(
         if creds is None and isinstance(repo, RemoteRepo):
             assert repo.repo_url is not None
             try:
-                creds, _ = get_repo_creds_and_default_branch(
+                creds = get_repo_creds(
                     repo_url=repo.repo_url,
                     identity_file=git_identity_file,
                     oauth_token=oauth_token,