From 1f53c4fd1b852bab5fc35ef64477a890eaa80939 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Mon, 7 Feb 2022 20:04:08 +0200
Subject: [PATCH] Fix agent fails to check out code from main branch when
 branch/commit is not explicitly specified

---
 clearml_agent/commands/worker.py | 52 ++++++++++++++++----------------
 clearml_agent/helper/repo.py     | 22 ++++++++++----
 2 files changed, 42 insertions(+), 32 deletions(-)

diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index ea536e4..2ac77bb 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -2135,32 +2135,32 @@ class Worker(ServiceCommandSection):
 
         cwd = vcs.location if vcs and vcs.location else directory
 
-        if is_cached and not standalone_mode:
-            # reinstalling git / local packages
-            package_api = copy(self.package_api)
-            OnlyExternalRequirements.cwd = package_api.cwd = cwd
-            package_api.requirements_manager = self._get_requirements_manager(
-                base_interpreter=package_api.requirements_manager.get_interpreter(),
-                requirement_substitutions=[OnlyExternalRequirements]
-            )
-            # make sure we run the handlers
-            cached_requirements = \
-                {k: package_api.requirements_manager.replace(requirements[k] or '')
-                 for k in requirements}
-            if str(cached_requirements.get('pip', '')).strip() \
-                    or str(cached_requirements.get('conda', '')).strip():
-                package_api.load_requirements(cached_requirements)
-            # make sure we call the correct freeze
-            requirements_manager = package_api.requirements_manager
-
-        elif not is_cached and not standalone_mode:
-            self.install_requirements(
-                execution,
-                repo_info,
-                requirements_manager=requirements_manager,
-                cached_requirements=requirements,
-                cwd=cwd,
-            )
+        if not standalone_mode:
+            if is_cached:
+                # reinstalling git / local packages
+                package_api = copy(self.package_api)
+                OnlyExternalRequirements.cwd = package_api.cwd = cwd
+                package_api.requirements_manager = self._get_requirements_manager(
+                    base_interpreter=package_api.requirements_manager.get_interpreter(),
+                    requirement_substitutions=[OnlyExternalRequirements]
+                )
+                # make sure we run the handlers
+                cached_requirements = \
+                    {k: package_api.requirements_manager.replace(requirements[k] or '')
+                     for k in requirements}
+                if str(cached_requirements.get('pip', '')).strip() \
+                        or str(cached_requirements.get('conda', '')).strip():
+                    package_api.load_requirements(cached_requirements)
+                # make sure we call the correct freeze
+                requirements_manager = package_api.requirements_manager
+            else:
+                self.install_requirements(
+                    execution,
+                    repo_info,
+                    requirements_manager=requirements_manager,
+                    cached_requirements=requirements,
+                    cwd=cwd,
+                )
 
         # do not update the task packages if we are using conda,
         # it will most likely make the task environment unreproducible
diff --git a/clearml_agent/helper/repo.py b/clearml_agent/helper/repo.py
index fa0c8d6..d3369ec 100644
--- a/clearml_agent/helper/repo.py
+++ b/clearml_agent/helper/repo.py
@@ -108,7 +108,7 @@ class VCS(object):
         )
         self.url = url
         self.location = Text(location)
-        self.revision = revision
+        self._revision = revision
         self.log = self.session.get_logger(__name__)
 
     @property
@@ -390,7 +390,7 @@ class VCS(object):
         """
         Checkout repository at specified revision
         """
-        self.call("checkout", self.revision, *self.checkout_flags, cwd=self.location)
+        self.call("checkout", self._revision, *self.checkout_flags, cwd=self.location)
 
     @abc.abstractmethod
     def pull(self):
@@ -519,7 +519,7 @@ class VCS(object):
 
 class Git(VCS):
     executable_name = "git"
-    main_branch = "master"
+    main_branch = ("master", "main")
     clone_flags = ("--quiet", "--recursive")
     checkout_flags = ("--force",)
     COMMAND_ENV = {
@@ -531,7 +531,9 @@ class Git(VCS):
 
     @staticmethod
     def remote_branch_name(branch):
-        return "origin/{}".format(branch)
+        return [
+            "origin/{}".format(b) for b in ([branch] if isinstance(branch, str) else branch)
+        ]
 
     def executable_not_found_error_help(self):
         return 'Cannot find "{}" executable. {}'.format(
@@ -553,7 +555,15 @@ class Git(VCS):
         """
         Checkout repository at specified revision
         """
-        self.call("checkout", self.revision, *self.checkout_flags, cwd=self.location)
+        revisions = [self._revision] if isinstance(self._revision, str) else self._revision
+        for i, revision in enumerate(revisions):
+            try:
+                self.call("checkout", revision, *self.checkout_flags, cwd=self.location)
+                break
+            except subprocess.CalledProcessError:
+                if i == len(revisions) - 1:
+                    raise
+
         try:
             self.call("submodule", "update", "--recursive", cwd=self.location)
         except:  # noqa
@@ -593,7 +603,7 @@ class Hg(VCS):
             "pull",
             self.url_with_auth,
             cwd=self.location,
-            *(("-r", self.revision) if self.revision else ())
+            *(("-r", self._revision) if self._revision else ())
         )
 
     info_commands = dict(