2016年7月23日

[python] 限制 subprocess 的 CPU & memory 使用量

暑假到了,又一批要進入研究所的新生要來報到,由於實驗室的傳統是所有新生都要接受 programming training project,而我們實驗室的 training 是老闆會幾篇經典的論文要新生實做出來,雖然標準年年都會調整,不過今年老闆明確的說他要檢查新生寫的程式,然後通過的標準除了要檢驗結果的正確性,也要檢查執行時間 & 記憶體使用量。而這次的 training project 老闆指定其中一題要把我現在做的東西切割一小塊出來讓他們實做看看,所以這次我必須要寫一個 verifier 去做自動化的驗證


首先呢,由於我那部份的題目是利用 DP (dynamic programming) 來達到 performance 的要求,所以對於記憶體的使用量要特別注意,不然極有可能程式一有 bug 就把記憶體吃光光了;另一方面,實作方式不好的話有可能會讓 run time 長到受不了;綜合這兩者,如果是讓程式完整執行完才回報最大的記憶體使用量與執行時間不是個好方法,因為如果要測試的程式有問題是會死在那邊的。因此我們可以了解如果要達到自動化驗證,我們需要的是在測試程式執行的過程中就針對他可執行的最長時間 與最大的記憶體使用量加以限制,類似 ACM online judge 的方式。

首先是測試程式,首先是我用 subprocess 另外開了一個 process 專門讓測試程式能夠執行給訂的測試資料,同時將程式的輸出資料  pipe 到我指定的地方去以便遍後續分析。另一方面, subprocess 能夠直接指定他最多只能執行多久,像我這次的使用方式是這樣:
  1. 利用 subprocess.Popen 指定執行方式與環境設置
  2. 藉由 subprocess.wait(TL) 讓 main process 等待測試程式執行完畢,但是最多只等 TL 秒,超過的話會丟出 subprocess.TimeoutExpired 的 exception
所以一但我接到這個 exception 就可以直接把執行測試程式的 process 給砍了。而記憶體使用量的問題就比較麻煩了,python 是有提供原生的 package 可以用,我找到且覺得比較容易的是 resource,不過實際上也還是遇上了些麻煩。雖然 resource 可以設定與取得各類資源的使用限制 (其實很多種資源都能設定,可以直接參考官方 manual 的說明,像是執行時間也能設定),不過其實最後要看 OS 的支援程度,像我這次想設計最大的記憶體使用量時,因為覺得最接近的是 resource.RLIMIT_RSS 這個選項所以就直接用了,後來發現有設定跟沒設定差不多後去查了 stackoverflow 才知道原來新的 Linux 版本不知道這個選項了 Orz 所以後來是選 resource.RLIMIT_AS ,雖然明顯不同,但相對接近了。

設定的方法其實滿簡單的,可以參考這篇討論,不過簡單來說就是在 subprocess.Popen 的參數中加上 "preexec_fn=setlimits" 這個參數,其中 setlimits 是設定可用資源上限的 function,大致上會長得像這樣
def setrlimit(mem_limit):
    resource.setrlimit(resource.RLIMIT_AS, (mem_limit, resource.RLIM_INFINITY)))
基本上的概念就是利用 resource.setrlimit 去控制可用的資源上限 (所以上面這行的可以一直增加想要的功能),相關我想網路上就有很多討論,這邊就直接略過囉。比較值得注意的地方應該就如果超過記憶體使用上面的話會丟 MemoryError 這個 exception。

這次整個 script 雖然我是採用 python 來寫,不過這點純粹是個人覺得 python 寫起來比較輕鬆 XD 其他語言也許也有類似功能,不過這就不在我這次討論的範圍內了,但依舊歡迎大家討論其他語言是否有更方便的方式

沒有留言:

張貼留言