numbaで一定時間ごとに進捗を書き出す

経緯

numbaを使用するとき、多くの場合で長時間計算しっぱなしだと思います。
そんな時には、進捗度合いを知りたいでしょう。
現在もnumba_progressというものを他の人が作ってくれいているので、それ自体は可能です。

ですが、そんな状況の時、多くはターミナルからの実行や、ノートブック上からの実行ではなく、linuxならnohupとかを使用して、ターミナルから出ても勝手にまわり続けるようにするんじゃないかと思います。
その場合、numba-progressはtqdmとほぼ同じで、例えばファイルにlogを出力する場合大量にprogressbarが出力されます。

こんな感じ、一応numbaを使わないなら解決できそうです。
numbaに関してはなさそうだった(自分が見つけられていないだけかも)ので、自分で作ってみました。
こんな感じ

コード

ちょっと長いので折りたたみ

import numpy as np
from numba import int64
import time,datetime,logging,threading
from numba.experimental import jitclass
from numba_progress.numba_atomic import atomic_add
        
class loop_log:
    def __init__(self,logger=None,interval=5,tot_loop=None,header="",footer=""):    
        if logger is None:
            self.logger=logging.getLogger("loop_log")
            self.logger.setLevel(logging.DEBUG)
            for h in self.logger.handlers[:]:
                logger.removeHandler(h)
                h.close()
            st_handler = logging.StreamHandler()
            st_handler.setLevel(logging.INFO)
            _format = "%(asctime)s %(name)s [%(levelname)s] : %(message)s"
            st_handler.setFormatter(logging.Formatter(_format))
            self.logger.addHandler(st_handler)
        else:
            self.logger=logger

        self.is_in_progress=False
        self.interval=interval
        self.counter=jit_counter()
        if tot_loop is None:
            self.is_tot_loop=False
        else:
            if not isinstance(tot_loop,int):
                raise TypeError(f"Expected variable 'tot_loop' to be of type int, but got {type(tot_loop)}.")
            self.is_tot_loop=True
            self.tot_loop=tot_loop
            self.str_tot_loop=str(tot_loop)
        self.header=header
        self.footer=footer
    
    def out_log(self):
        time.sleep(self.interval)
        while self.is_in_progress:
            dtime=datetime.datetime.now()-self.start_time
            if self.is_tot_loop:
                count=self.counter.count[0]
                if count!=0:
                    left_time=(self.tot_loop-count)*dtime/count
                    txt=f"{str(count).rjust(len(self.str_tot_loop))}/{self.str_tot_loop}  {dtime} : {left_time}"
                else:
                    txt=f"{str(count).rjust(len(self.str_tot_loop))}/{self.str_tot_loop}  {dtime} : inf"
            else:
                txt=f"{self.counter.count[0]}  {dtime}"
            self.logger.info(self.header+txt+self.footer)
            time.sleep(self.interval)
        
    def __enter__(self):
        self.is_in_progress=True
        self.start_time=datetime.datetime.now()
        thread = threading.Thread(target=self.out_log)
        thread.start()
        return self.counter
        
    def __exit__(self,*args):
        self.is_in_progress=False

spec = [('count', int64[:])]
@jitclass(spec)
class jit_counter:
    def __init__(self):
        self.count=np.zeros(1,dtype="int64")
    def update(self,n):
        atomic_add(self.count,0,n)

使い方

numba progressに揃えてみました。
上のコードをlog_progress.pyとか名前つけて保存して使ってください。
numba_progressのatomic_addを使用するので、numba_progressのinstallは必須です。

from log_progress import loop_log

num_iterations = 100

@njit(nogil=True, parallel=True)
def numba_function(num_iterations, progress_proxy):
    for i in prange(num_iterations):
        #<DO CUSTOM WORK HERE>
        progress_proxy.update(1)

with loop_log() as progress:
    numba_function(num_iterations, progress)

ちなみにnumba progressの使いかた

from numba import njit, prange
from numba_progress import ProgressBar

num_iterations = 100

@njit(nogil=True, parallel=True)
def numba_function(num_iterations, progress_proxy):
    for i in prange(num_iterations):
        #<DO CUSTOM WORK HERE>
        progress_proxy.update(1)

with ProgressBar(total=num_iterations) as progress:
    numba_function(num_iterations, progress)

logger,interval,tot_loop,header,footerのオプションがあります。

• logger
  • 標準モジュールloggingのloggerを入れてくれればそれにlogを渡します。
  • ファイル出力したい場合はFileHandlerのあるloggerを入れてね
• interval
  • 何秒ごとに更新するかです。初期値は5秒
• tot_loop
  • 分母です。
• header
  • 表示の最初にstrをくっつける
• footer
• 表示の最後にstrをくっつける

自身がnumba_progressを使用していてupdate以外使用したいと思ったことがないため、
それ以外は未実装です(何があるかも知らない)。