use ssh multiplexing to create SSH workers with tunneling (JuliaLang/julia#34295)

* use ssh multiplexing to create SSH workers with tunneling

* rename funcname

* add test to check existence of ssh multiplexing master socket

* add doc for ssh multiplexing in SSHWorker

* Fix typo

* add :multiplex option to make SSH multiplexing configurable

Author: s417-lama

src/cluster.jl

@@ -28,6 +28,8 @@ The `userdata` field is used to store information for each worker by external ma
 
 Some fields are used by `SSHManager` and similar managers:
   * `tunnel` -- `true` (use tunneling), `false` (do not use tunneling), or [`nothing`](@ref) (use default for the manager)
+  * `multiplex` -- `true` (use SSH multiplexing for tunneling) or `false`
+  * `forward` -- the forwarding option used for `-L` option of ssh
   * `bind_addr` -- the address on the remote host to bind to
   * `sshflags` -- flags to use in establishing the SSH connection
   * `max_parallel` -- the maximum number of workers to connect to in parallel on the host
@@ -58,6 +60,8 @@ mutable struct WorkerConfig
 
     # SSHManager / SSH tunnel connections to workers
     tunnel::Union{Bool, Nothing}
+    multiplex::Union{Bool, Nothing}
+    forward::Union{AbstractString, Nothing}
     bind_addr::Union{AbstractString, Nothing}
     sshflags::Union{Cmd, Nothing}
     max_parallel::Union{Int, Nothing}
@@ -548,7 +552,7 @@ function launch_n_additional_processes(manager, frompid, fromconfig, cnt, launch
             (bind_addr, port) = address
 
             wconfig = WorkerConfig()
-            for x in [:host, :tunnel, :sshflags, :exeflags, :exename, :enable_threaded_blas]
+            for x in [:host, :tunnel, :multiplex, :sshflags, :exeflags, :exename, :enable_threaded_blas]
                 Base.setproperty!(wconfig, x, Base.getproperty(fromconfig, x))
             end
             wconfig.bind_addr = bind_addr

src/managers.jl

@@ -71,6 +71,8 @@ Keyword arguments:
 * `tunnel`: if `true` then SSH tunneling will be used to connect to the worker from the
   master process. Default is `false`.
 
+* `multiplex`: if `true` then SSH multiplexing is used for SSH tunneling. Default is `false`.
+
 * `sshflags`: specifies additional ssh options, e.g. ```sshflags=\`-i /home/foo/bar.pem\````
 
 * `max_parallel`: specifies the maximum number of workers connected to in parallel at a
@@ -113,9 +115,9 @@ This timeout can be controlled via environment variable `JULIA_WORKER_TIMEOUT`.
 The value of `JULIA_WORKER_TIMEOUT` on the master process specifies the number of seconds a
 newly launched worker waits for connection establishment.
 """
-function addprocs(machines::AbstractVector; tunnel=false, sshflags=``, max_parallel=10, kwargs...)
+function addprocs(machines::AbstractVector; tunnel=false, multiplex=false, sshflags=``, max_parallel=10, kwargs...)
     check_addprocs_args(kwargs)
-    addprocs(SSHManager(machines); tunnel=tunnel, sshflags=sshflags, max_parallel=max_parallel, kwargs...)
+    addprocs(SSHManager(machines); tunnel=tunnel, multiplex=multiplex, sshflags=sshflags, max_parallel=max_parallel, kwargs...)
 end
 
 
@@ -149,6 +151,8 @@ function launch_on_machine(manager::SSHManager, machine, cnt, params, launched,
     dir = params[:dir]
     exename = params[:exename]
     exeflags = params[:exeflags]
+    tunnel = params[:tunnel]
+    multiplex = params[:multiplex]
 
     # machine could be of the format [user@]host[:port] bind_addr[:bind_port]
     # machine format string is split on whitespace
@@ -178,6 +182,20 @@ function launch_on_machine(manager::SSHManager, machine, cnt, params, launched,
     end
     sshflags = `$(params[:sshflags]) $portopt`
 
+    if tunnel
+        # First it checks if ssh multiplexing has been already enabled and the master process is running.
+        # If it's already running, later ssh sessions also use the same ssh multiplexing session even if
+        # `multiplex` is not explicitly specified; otherwise the tunneling session launched later won't
+        # go to background and hang. This is because of OpenSSH implementation.
+        if success(`ssh $sshflags -O check $host`)
+            multiplex = true
+        elseif multiplex
+            # automatically create an SSH multiplexing session at the next SSH connection
+            controlpath = "~/.ssh/julia-%r@%h:%p"
+            sshflags = `$sshflags -o ControlMaster=auto -o ControlPath=$controlpath -o ControlPersist=no`
+        end
+    end
+
     # Build up the ssh command
 
     # the default worker timeout
@@ -211,7 +229,8 @@ function launch_on_machine(manager::SSHManager, machine, cnt, params, launched,
     wconfig = WorkerConfig()
     wconfig.io = io.out
     wconfig.host = host
-    wconfig.tunnel = params[:tunnel]
+    wconfig.tunnel = tunnel
+    wconfig.multiplex = multiplex
     wconfig.sshflags = sshflags
     wconfig.exeflags = exeflags
     wconfig.exename = exename
@@ -256,25 +275,32 @@ end
 
 
 """
-    ssh_tunnel(user, host, bind_addr, port, sshflags) -> localport
+    ssh_tunnel(user, host, bind_addr, port, sshflags, multiplex) -> localport
 
 Establish an SSH tunnel to a remote worker.
 Return a port number `localport` such that `localhost:localport` connects to `host:port`.
 """
-function ssh_tunnel(user, host, bind_addr, port, sshflags)
+function ssh_tunnel(user, host, bind_addr, port, sshflags, multiplex)
     port = Int(port)
     cnt = ntries = 100
-    # if we cannot do port forwarding, bail immediately
+
     # the connection is forwarded to `port` on the remote server over the local port `localport`
-    # the -f option backgrounds the ssh session
-    # `sleep 60` command specifies that an alloted time of 60 seconds is allowed to start the
-    # remote julia process and establish the network connections specified by the process topology.
-    # If no connections are made within 60 seconds, ssh will exit and an error will be printed on the
-    # process that launched the remote process.
-    ssh = `ssh -T -a -x -o ExitOnForwardFailure=yes`
     while cnt > 0
         localport = next_tunnel_port()
-        if success(detach(`$ssh -f $sshflags $user@$host -L $localport:$bind_addr:$port sleep 60`))
+        if multiplex
+            # It assumes that an ssh multiplexing session has been already started by the remote worker.
+            cmd = `ssh $sshflags -O forward -L $localport:$bind_addr:$port $user@$host`
+        else
+            # if we cannot do port forwarding, fail immediately
+            # the -f option backgrounds the ssh session
+            # `sleep 60` command specifies that an alloted time of 60 seconds is allowed to start the
+            # remote julia process and establish the network connections specified by the process topology.
+            # If no connections are made within 60 seconds, ssh will exit and an error will be printed on the
+            # process that launched the remote process.
+            ssh = `ssh -T -a -x -o ExitOnForwardFailure=yes`
+            cmd = detach(`$ssh -f $sshflags $user@$host -L $localport:$bind_addr:$port sleep 60`)
+        end
+        if success(cmd)
             return localport
         end
         cnt -= 1
@@ -427,9 +453,11 @@ function connect(manager::ClusterManager, pid::Int, config::WorkerConfig)
         sem = tunnel_hosts_map[pubhost]
 
         sshflags = notnothing(config.sshflags)
+        multiplex = something(config.multiplex, false)
         acquire(sem)
         try
-            (s, bind_addr) = connect_to_worker(pubhost, bind_addr, port, user, sshflags)
+            (s, bind_addr, forward) = connect_to_worker_with_tunnel(pubhost, bind_addr, port, user, sshflags, multiplex)
+            config.forward = forward
         finally
             release(sem)
         end
@@ -515,9 +543,23 @@ function connect_to_worker(host::AbstractString, port::Integer)
 end
 
 
-function connect_to_worker(host::AbstractString, bind_addr::AbstractString, port::Integer, tunnel_user::AbstractString, sshflags)
-    s = connect("localhost", ssh_tunnel(tunnel_user, host, bind_addr, UInt16(port), sshflags))
-    (s, bind_addr)
+function connect_to_worker_with_tunnel(host::AbstractString, bind_addr::AbstractString, port::Integer, tunnel_user::AbstractString, sshflags, multiplex)
+    localport = ssh_tunnel(tunnel_user, host, bind_addr, UInt16(port), sshflags, multiplex)
+    s = connect("localhost", localport)
+    forward = "$localport:$bind_addr:$port"
+    (s, bind_addr, forward)
+end
+
+
+function cancel_ssh_tunnel(config::WorkerConfig)
+    host = notnothing(config.host)
+    sshflags = notnothing(config.sshflags)
+    tunnel = something(config.tunnel, false)
+    multiplex = something(config.multiplex, false)
+    if tunnel && multiplex
+        forward = notnothing(config.forward)
+        run(`ssh $sshflags -O cancel -L $forward $host`)
+    end
 end
 
 
@@ -531,7 +573,12 @@ It should cause the remote worker specified by `pid` to exit.
 on `pid`.
 """
 function kill(manager::ClusterManager, pid::Int, config::WorkerConfig)
-    remote_do(exit, pid) # For TCP based transports this will result in a close of the socket
-                       # at our end, which will result in a cleanup of the worker.
+    remote_do(exit, pid)
+    nothing
+end
+
+function kill(manager::SSHManager, pid::Int, config::WorkerConfig)
+    remote_do(exit, pid)
+    cancel_ssh_tunnel(config)
     nothing
 end

test/distributed_exec.jl

@@ -720,6 +720,14 @@ if Sys.isunix() # aka have ssh
     @test length(new_pids) == num_workers
     test_n_remove_pids(new_pids)
 
+    print("\nssh addprocs with tunnel (SSH multiplexing)\n")
+    new_pids = addprocs_with_testenv([("localhost", num_workers)]; tunnel=true, multiplex=true, sshflags=sshflags)
+    @test length(new_pids) == num_workers
+    controlpath = joinpath(homedir(), ".ssh", "julia-$(ENV["USER"])@localhost:22")
+    @test issocket(controlpath)
+    test_n_remove_pids(new_pids)
+    @test :ok == timedwait(()->!issocket(controlpath), 10.0; pollint=0.5)
+
     print("\nAll supported formats for hostname\n")
     h1 = "localhost"
     user = ENV["USER"]