add ssh-key-path to ydbd_slice (#14463)

Author: iddqdex

ydb/tools/ydbd_slice/__init__.py

@@ -310,7 +310,7 @@ def deduce_components_from_args(args, cluster_details):
     return result
 
 
-def deduce_nodes_from_args(args, walle_provider, ssh_user):
+def deduce_nodes_from_args(args, walle_provider, ssh_user, ssh_key_path):
     cluster_hosts = safe_load_cluster_details(args.cluster, walle_provider).hosts_names
     result = cluster_hosts
 
@@ -326,7 +326,7 @@ def deduce_nodes_from_args(args, walle_provider, ssh_user):
         sys.exit("unable to deduce hosts")
 
     logger.info("use nodes '%s'", result)
-    return nodes.Nodes(result, args.dry_run, ssh_user=ssh_user, queue_size=args.cmd_queue_size)
+    return nodes.Nodes(result, args.dry_run, ssh_user=ssh_user, ssh_key_path=ssh_key_path, queue_size=args.cmd_queue_size)
 
 
 def ya_build(arcadia_root, artifact, opts, dry_run):
@@ -514,6 +514,12 @@ def ssh_args():
         help="user for ssh interaction with slice. Default value is $USER "
         "(which equals {user} now)".format(user=current_user),
     )
+    args.add_argument(
+        "--ssh-key-path",
+        metavar="SSH_KEY_PATH",
+        help="Path to ssh private key"
+        "(which equals {user} now)".format(user=current_user),
+    )
     return args
 
 
@@ -625,7 +631,7 @@ def dispatch_run(func, args, walle_provider, need_confirmation=False):
     cluster_details = safe_load_cluster_details(args.cluster, walle_provider)
     components = deduce_components_from_args(args, cluster_details)
 
-    nodes = deduce_nodes_from_args(args, walle_provider, args.ssh_user)
+    nodes = deduce_nodes_from_args(args, walle_provider, args.ssh_user, args.ssh_key_path)
 
     temp_dir = deduce_temp_dir_from_args(args)
     clear_tmp = not args.dry_run and args.temp_dir is None

ydb/tools/ydbd_slice/nodes.py

@@ -9,13 +9,14 @@
 
 
 class Nodes(object):
-    def __init__(self, nodes, dry_run=False, ssh_user=None, queue_size=0):
+    def __init__(self, nodes, dry_run=False, ssh_user=None, queue_size=0, ssh_key_path=None):
         assert isinstance(nodes, list)
         assert len(nodes) > 0
         assert isinstance(nodes[0], str)
         self._nodes = nodes
         self._dry_run = bool(dry_run)
         self._ssh_user = ssh_user
+        self._ssh_key_path = ssh_key_path
         self._logger = logger.getChild(self.__class__.__name__)
         self._queue = queue.Queue(queue_size)
         self._qsize = queue_size
@@ -24,12 +25,15 @@ def __init__(self, nodes, dry_run=False, ssh_user=None, queue_size=0):
     def nodes_list(self):
         return self._nodes
 
-    def _get_ssh_command_prefix(self):
+    def _get_ssh_command_prefix(self, remote=False):
         command = []
         command.extend(['ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null', '-A'])
         if (self._ssh_user):
             command.extend(['-l', self._ssh_user])
 
+        if not remote and self._ssh_key_path:
+            command.extend(['-i', self._ssh_key_path])
+
         return command
 
     def _check_async_execution(self, running_jobs, check_retcode=True, results=None):
@@ -144,9 +148,9 @@ def _copy_between_nodes(self, hub, hub_path, hosts, remote_path):
             if self._dry_run:
                 continue
             cmd = self._get_ssh_command_prefix() + [dst]
-            rsh = " ".join(self._get_ssh_command_prefix())
+            rsh = " ".join(self._get_ssh_command_prefix(remote=True))
             cmd.extend([
-                "sudo", "SSH_AUTH_SOCK=$SSH_AUTH_SOCK", "rsync", "-avqW", "--del", "--no-o", "--no-g",
+                "sudo", "--preserve-env=SSH_AUTH_SOCK", "rsync", "-avqW", "--del", "--no-o", "--no-g",
                 "--rsh='{}'".format(rsh),
                 src, remote_path
             ])