• The University of Tokyo / RIKEN
  • tsuzukums.k.u-tokyo.ac.jp

job管理

Jobの管理について考えていたら、Jobの実行と実行後の管理は分けない方が綺麗だなという気持ちになってきた。 たぶん実行済みjobの管理を分けようと思った元凶の一つは、実行結果の統計をパッととれる環境が整っていないことで、 ただそれを解決するためには普段の実験コードの方を書き直すべき。

なぜ今の状況が生まれてしまったかというと、諸事情でDBを使えない環境で計算しがちというのが大きそう。 とりあえずrecordioに実行結果のprotobufを記録とかしておいて、外部(自宅のサーバとか)に投げるところまでを支援する形にしておくのが無難そう。

観光しながらぼんやりと考えたのは下みたいなインターフェース。

#!/usr/bin/python
# -*- coding: utf-8 -*-


import os
import mango as mg


def register_jobs():
    with mg.scope('job0'):
        mg.register(
            ['echo', '0'],
            name='echo-0',
            cpus=1,
            gpus=0,
            retry=3,
            env=os.environ)
    
    with mg.scope('job1'):
        for i in range(1, 4):
            mg.register(
                ['echo', '{}'.format(i)],
                name='echo-{}'.format(i),
                cpus=1,
                gpus=2,
                priority=5 - i,
                deps=['job0'])
    
    with mg.scope('job2'):
        mg.register(
            ['echo', '4'],
            name='echo-4',
            cpus=1,
            gpus=1,
            retry=1,
            deps=['job1/echo-3']
            env=os.environ)

if __name__ == '__main__':
    register_jobs()
    mg.run(name='example')

今使ってる実験管理scriptだとyamlファイルに色々記述してそれをコンパイル?して実行するんだけど、 ちょっとしたDSLになってしまって煩雑になったので、Pythonに任せることにした。

上のコードは大きくjob group(ネストできる)がjob0/job1/job2の3つあるprogram。 job1は全てjob0に依存しているので、まず0がechoされる。 続いてjob1の中のworker(job0に依存)がどれもexecutableになる。 ただし、priorityがあるので1, 2, 3の順に出力される。 job2はjob1/echo-3に依存しているので、これの後に実行され、最後に4がechoされてprogram終了。

pythonに乗っかりつつこのくらい自由度が高いとハイパラのグリッドサーチから環境の構築までほとんど学習コストなしでサクッとかけそう。

あとはterminalから

$ alias mgo="python -m mango"
$ mgo stat
> ID     NAME    STATUS DATE
  000001 example C      2019/05/02
$ mgo stat 000001
> NAME    STATUS DATE
  job0    C      2019/05/02
  job1    C      2019/05/02
  job2    C      2019/05/02
$ mgo stat 000001/job1
> NAME    STATUS DATE
  echo-1  C      2019/05/02
  echo-2  C      2019/05/02
  echo-3  C      2019/05/02
$ mgo stdout 000001/job1/echo-3
> 3

みたいに階層構造を使ってスッキリとjobの実行結果を見通せるといい(bash-completionがないと辛そう)。 もっと色々とほしい機能はあるけど、とりあえずこれだけ。

ところでmanagementの語感からmangoという名前にしてみたけれど、python-mangoというパッケージがすでに存在しているので名前を考え直す必要がある。 面倒。