Julia のワーカをクラスタのワーカノードで起動する

Juliaには分散計算の機構がはじめから組み込まれていて、リモートノードでワーカプロセスを起動することも簡単にできる。SSHが設定されていれば

```Julia
addprocs([("HOSTNAME 127.0.0.1", 1)], tunnel=true, exename=JULIA_PATH, topology=:master_worker)
```

のようにすれば良い。ただバッチキューイングシステムで管理されているような大規模クラスタだと、そもそもワーカノードに直接SSHすることができないのでそう簡単には行かない。`exename`で実行ファイルが指定できることを利用して、バッチキューイングシステムで起動する機構をラップしてやればよい。

juliaのプロセスは、ワーカモードで起動すると次のように動作する。

- 標準入力からクッキーと呼ばれる認証に用いるランダムな文字列を読み込む
- 標準出力に`julia_worker:PORT#HOSTIP`というような文字列を出力する
- マスタプロセスからの接続をPORTで待受けて通信する

```
   +--------+     cookie  ->   +--------+
   |        |  ------------->  |        |
   | master |    <- portstr    | worker |
   |        |  ------------->  |        |
   +--------+    <- comm ->    +--------+
```



この挙動をラッパプロセスで再現し、ワーカノードで起動したプロセスに中継してやればい。つまり、ラッパプロセスは次のことをすれば良い。

- マスタプロセスに対してワーカプロセスであるかのように動作し、
- ワーカプロセスに対してはマスタプロセスであるかのように動作し、
- マスタプロセスとワーカプロセスとの通信を透過的に中継する

```
   +--------+     cookie  ->   +--------+     cookie  ->   +--------+
   |        |  ------------->  |        |  ------------->  |        |
   | master |    <- portstr    | wrapper|    <- portstr    | worker |
   |        |  ------------->  |        |  ------------->  |        |
   +--------+    <- comm ->    +--------+    <- comm ->    +--------+
```

具体的な挙動は次の様になる。

1. マスタからの通信を受け付けるソケットを開く
2. ワーカノード上で、Juliaのプロセスをワーカモードでqrshで起動する。
3. 標準入力からクッキーを読み込む
4. ワーカプロセスにクッキーを送信
5. ワーカプロセスから、ホストとポートの文字列を受信
6. ホストとポートの文字列を1．で開いたポートに書き換えて、標準出力に書き出す
7. ワーカプロセスに対するソケットを開く
8. 1.のソケットと7のソケットを双方向にフォワード

複数のソケットのフォワードはCで書くと`select`とか使わなければならずかなり面倒なのだけど、Juliaではコルーチンを使えばいいので楽勝。

コードはこんな感じ。

``` Julia
#!/usr/bin/env julia
import Sockets

GROUP_ID      = "GID"
RESOURCE_TYPE = "RTYPE=1"

function invoke_worker()
    QRSH = ["qrsh", "-g", GROUP_ID, "-l", RESOURCE_TYPE]
    CMD = "source ~/.bashrc; module load julia; julia --worker "
    
    return open(`$QRSH bash -c \"$CMD\"`, write = true, read = true)
end

worker_re = r"julia_worker:([0-9]+)#([0-9]+\.[0-9]+\.[0-9]+\.[0-9]+)"
function parse_connect_string(worker_str)
    m = match(worker_re, worker_str)
    if isnothing(m)
        @error "$worker_str does not match expectation\n"
        exit()
    end
    return m[2], parse(Int64, m[1])
end

function create_connect_string(host, port)
   return "julia_worker:$port#$host"
end

function forward(inStream, outStream)
    while true
        buf = readavailable(inStream)
        if length(buf) == 0
           break
        end
        write(outStream, buf)
        flush(outStream)
    end
end


function find_and_listen()::Tuple{Sockets.TCPServer, Int}
    candidate = 9000
    for i in 1:100
        try
            candidate += 1
            s = Sockets.listen(candidate)
            return s, candidate
        catch e
            if ! isa(e, Base.IOError)
               rethrow(e)
            end
        end
    end
    @error "tried 100 ports and found no available port, exit."
    exit()
end

function accept_forward(master_ss, worker)
   s = Sockets.accept(master_ss)
   @info "accept connection from master, start forwarding"

   @async forward(s, worker)
   forward(worker, s)
end


function main()
    ss, my_port = find_and_listen()
    @info "listening on $my_port"

    proc = invoke_worker()
    @info "worker launched"

    cookie = readline(stdin, keep=true)
    write(proc.in, cookie)
    close(proc.in)
    @info "forward cookie $cookie"

    worker_str = readline(proc.out)
    @info "worker_str is $worker_str"
    host, worker_port = parse_connect_string(worker_str)
    @info "worker is running at $host:$worker_port"

    worker = Sockets.connect(host, worker_port)
    @info "connect to the worker"

    worker_str = create_connect_string("127.0.0.1", my_port)
    @info "return worker string: $worker_str"
    write(stdout, worker_str)
    write(stdout, "\n")
    flush(stdout)

    accept_forward(ss, worker)
end

main()
```
これを適当な場所に置いておいて、`exename`で指定して起動すればいい。

## 所感
Julia、数値演算をメインターゲットにしているはずなのだけど、意外にこの手のプログラムも書きやすい。入出力関数の仕様がエキセントリックなので、調べるのが大変だけど。`read`関数でバッファにデータがあるにも関わらず、読み込みがブロックすると誰が予想できるだろうか。こういうのはやめてほしいなあ。