From 4c9410c5fed528e2618ba44c3a176c6bb2d04810 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sun, 11 Jul 2021 09:43:57 +0300
Subject: [PATCH] Fix auto mount SSH_AUTH_SOCK into docker (issue #45)

---
 clearml_agent/commands/worker.py | 11 ++++++++++-
 clearml_agent/definitions.py     |  1 +
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index 9c81071..40fbe93 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -51,6 +51,7 @@ from clearml_agent.definitions import (
     ENV_AWS_SECRET_KEY,
     ENV_AZURE_ACCOUNT_KEY,
     ENV_AGENT_DISABLE_SSH_MOUNT,
+    ENV_SSH_AUTH_SOCK,
 )
 from clearml_agent.definitions import WORKING_REPOSITORY_DIR, PIP_EXTRA_INDICES
 from clearml_agent.errors import APIError, CommandFailedError, Sigterm
@@ -2730,7 +2731,15 @@ class Worker(ServiceCommandSection):
         if temp_config.get("agent.venvs_cache.path", None):
             temp_config.put("agent.venvs_cache.path", '/root/.clearml/venvs-cache')
 
-        if ENV_AGENT_DISABLE_SSH_MOUNT.get():
+        if (ENV_SSH_AUTH_SOCK.get() or '').strip():
+            self._host_ssh_cache = None
+            ssh_auth_sock_env = 'SSH_AUTH_SOCK={}'.format(ENV_SSH_AUTH_SOCK.get())
+            if not self._extra_docker_arguments or ssh_auth_sock_env not in self._extra_docker_arguments:
+                self._extra_docker_arguments = (self._extra_docker_arguments or []) + [
+                    '-v', '{}:{}'.format(ENV_SSH_AUTH_SOCK.get(), ENV_SSH_AUTH_SOCK.get()),
+                    '-e', ssh_auth_sock_env,
+                ]
+        elif ENV_AGENT_DISABLE_SSH_MOUNT.get():
             self._host_ssh_cache = None
         else:
             self._host_ssh_cache = mkdtemp(prefix='clearml_agent.ssh.')
diff --git a/clearml_agent/definitions.py b/clearml_agent/definitions.py
index 6b0d8e2..39e300c 100644
--- a/clearml_agent/definitions.py
+++ b/clearml_agent/definitions.py
@@ -135,6 +135,7 @@ ENV_AGENT_GIT_USER = EnvironmentConfig('CLEARML_AGENT_GIT_USER', 'TRAINS_AGENT_G
 ENV_AGENT_GIT_PASS = EnvironmentConfig('CLEARML_AGENT_GIT_PASS', 'TRAINS_AGENT_GIT_PASS')
 ENV_AGENT_GIT_HOST = EnvironmentConfig('CLEARML_AGENT_GIT_HOST', 'TRAINS_AGENT_GIT_HOST')
 ENV_AGENT_DISABLE_SSH_MOUNT = EnvironmentConfig('CLEARML_AGENT_DISABLE_SSH_MOUNT', type=bool)
+ENV_SSH_AUTH_SOCK = EnvironmentConfig('SSH_AUTH_SOCK')
 ENV_TASK_EXECUTE_AS_USER = EnvironmentConfig('CLEARML_AGENT_EXEC_USER', 'TRAINS_AGENT_EXEC_USER')
 ENV_TASK_EXTRA_PYTHON_PATH = EnvironmentConfig('CLEARML_AGENT_EXTRA_PYTHON_PATH', 'TRAINS_AGENT_EXTRA_PYTHON_PATH')
 ENV_DOCKER_HOST_MOUNT = EnvironmentConfig('CLEARML_AGENT_K8S_HOST_MOUNT', 'CLEARML_AGENT_DOCKER_HOST_MOUNT',