Add support for in-process workers to LocalManager

Author: JamesWrigley

src/DistributedNext.jl

@@ -222,6 +222,11 @@ function Base.close(ctx::ClusterContext)
     if !isnothing(ctx.stdlib_watcher_timer)
         close(ctx.stdlib_watcher_timer::Timer)
     end
+
+    # Close all tracked sockets
+    @lock ctx.map_sock_wrkr for sock in keys(ctx.map_sock_wrkr[])
+        close(sock)
+    end
 end
 
 const CTX = ScopedValue(ClusterContext())

src/cluster.jl

@@ -231,6 +231,7 @@ mutable struct LocalProcess
     bind_port_hint::Int
     bind_port::Int
     cookie::String
+    in_process::Bool  # true when running as a task, not a separate OS process
     LocalProcess() = new(1)
 end
 
@@ -260,7 +261,7 @@ should not be relied upon to always pick the fastest interface.
 It does not return.
 """
 start_worker(cookie::AbstractString=readline(stdin); kwargs...) = start_worker(stdout, cookie; kwargs...)
-function start_worker(out::IO, cookie::AbstractString=readline(stdin); close_stdin::Bool=true, stderr_to_stdout::Bool=true)
+function start_worker(out::IO, cookie::AbstractString=readline(stdin); close_stdin::Bool=true, stderr_to_stdout::Bool=true, exit_on_close::Bool=true)
     init_multi()
 
     if close_stdin # workers will not use it
@@ -278,8 +279,15 @@ function start_worker(out::IO, cookie::AbstractString=readline(stdin); close_std
         sock = listen(interface, CTX[].lproc.bind_port)
     end
     errormonitor(@async while isopen(sock)
-        client = accept(sock)
-        process_messages(client, client, true)
+        try
+            client = accept(sock)
+            process_messages(client, client, true)
+        catch ex
+            # An IOError is thrown when the socket is closed
+            if !(ex isa Base.IOError)
+                rethrow()
+            end
+        end
     end)
     print(out, "julia_worker:")  # print header
     print(out, "$(CTX[].lproc.bind_port)#") # print port
@@ -298,13 +306,17 @@ function start_worker(out::IO, cookie::AbstractString=readline(stdin); close_std
         # To prevent hanging processes on remote machines, newly launched workers exit if the
         # master process does not connect in time.
         check_master_connect()
-        while true; wait(); end
+        while isopen(out)
+            sleep(0.1)
+        end
     catch err
         print(stderr, "unhandled exception on $(myid()): $(err)\nexiting.\n")
     end
 
     close(sock)
-    exit(0)
+    if exit_on_close
+        exit(0)
+    end
 end
 
 

src/managers.jl

@@ -458,8 +458,10 @@ end
 # LocalManager
 struct LocalManager <: ClusterManager
     np::Int
-    restrict::Bool  # Restrict binding to 127.0.0.1 only
+    restrict::Bool    # Restrict binding to 127.0.0.1 only
+    in_process::Bool  # Run workers as tasks in the same process (internal)
 end
+LocalManager(np, restrict) = LocalManager(np, restrict, false)
 
 """
     addprocs(np::Integer=Sys.CPU_THREADS; restrict=true, kwargs...) -> List of process identifiers
@@ -527,13 +529,30 @@ function launch(manager::LocalManager, params::Dict, launched::Array, c::Conditi
     end
 
     for i in 1:manager.np
-        cmd = `$(julia_cmd(exename)) $exeflags --bind-to $bind_to $(get_worker_arg())`
-        io = open(detach(setenv(addenv(cmd, env), dir=dir)), "r+")
-        write_cookie(io)
-
         wconfig = WorkerConfig()
-        wconfig.process = io
-        wconfig.io = io.out
+
+        if manager.in_process
+            cookie = cluster_cookie()
+            worker_ctx = ClusterContext()
+            worker_ctx.lproc.in_process = true
+            pipe = Pipe()
+            Base.link_pipe!(pipe; reader_supports_async=true, writer_supports_async=true)
+
+            task = Threads.@spawn @with CTX => worker_ctx begin
+                start_worker(pipe.in, cookie; close_stdin=false, stderr_to_stdout=false, exit_on_close=false)
+            end
+
+            wconfig.io = pipe.out
+            wconfig.userdata = (; ctx=worker_ctx, task, pipe)
+        else
+            cmd = `$(julia_cmd(exename)) $exeflags --bind-to $bind_to $(get_worker_arg())`
+            proc = open(detach(setenv(addenv(cmd, env), dir=dir)), "r+")
+
+            write_cookie(proc)
+            wconfig.io = proc.out
+            wconfig.process = proc
+        end
+
         wconfig.enable_threaded_blas = params[:enable_threaded_blas]
         push!(launched, wconfig)
     end
@@ -542,7 +561,7 @@ function launch(manager::LocalManager, params::Dict, launched::Array, c::Conditi
 end
 
 function manage(manager::LocalManager, id::Integer, config::WorkerConfig, op::Symbol)
-    if op === :interrupt
+    if op === :interrupt && !manager.in_process
         kill(config.process::Process, 2)
     end
 end
@@ -756,6 +775,15 @@ function kill(manager::SSHManager, pid::Int, config::WorkerConfig)
 end
 
 function kill(manager::LocalManager, pid::Int, config::WorkerConfig; profile_wait = 6, exit_timeout = 15, term_timeout = 15)
+    if manager.in_process
+        (; ctx, task, pipe) = config.userdata
+        # Close the pipe, which causes the start_worker() task to exit
+        close(pipe)
+        wait(task)
+        close(ctx)
+        return nothing
+    end
+
     # profile_wait = 6 is 1s for profile, 5s for the report to show
     # First, try sending `exit()` to the remote over the usual control channels
     remote_do(exit, pid)

src/process_messages.jl

@@ -228,7 +228,7 @@ function message_handler_loop(r_stream::IO, w_stream::IO, incoming::Bool)
             set_worker_state(werr, oldstate != WorkerState_terminating ? WorkerState_exterminated : WorkerState_terminated)
 
             # If unhandleable error occurred talking to pid 1, exit
-            if wpid == 1
+            if wpid == 1 && !ctx.lproc.in_process
                 if isopen(w_stream)
                     @error "Fatal error on process $(myid())" exception=e,catch_backtrace()
                 end