β

Python资源限制上下文

Yuan 509 阅读

在很多监控应用中,python脚本越来越多的被应用在生产环境,那么基于python的资源限制管理就变得很重要,主要问题有:

  1. 实现基础的资源限制功能(CPU优先级、内存使用、文件句柄数等等)。
  2. 可以灵活使用,避免重复劳动,代码复用性好。
  3. 尽量无视外部环境,在语言内部解决。

对于问题1和3,使用python stl自带的resource limit即可,但是对于问题2,需要我们好好思考。

对于脚本语言,对于shell中的ulimit内嵌命令来说,是以进程为单位,限制本进程的资源使用情况,对于单个脚本文件内部分代码的资源限制实现并不优雅。同时,在代码复用性上来说,也很难做出能多次复用的资源限制模板以供调用。

基于此,使用python的上下文系统是一个很好的选择,python的上下文系统可以很轻松的实现了以代码快为单位的环境配置与恢复,同时易于嵌套复用,对此情境非常适合。

代码:

#limit memory, num of file, core and cpu nice level
#Author: Yuan

import resource
import sys
import signal
import contextlib
import os

@contextlib.contextmanager
def limit_resource(**kargs):
	"""
	When limit resource, use this function in a with statement
	And this function only change soft limit
	"""

	dict = {"CORE": resource.RLIMIT_CORE,
#		"NICE": resource.RLIMIT_NICE,
		"NOFILE": resource.RLIMIT_NOFILE,
		"RSS": resource.RLIMIT_RSS,
		"CPU": resource.RLIMIT_CPU}

	backup_dict = {}

	print("Running PID: ", os.getpid())
	print("Backuping and limiting...")
	for (k,v) in kargs.items():
		if v >= 0 or v == -1:
			soft, hard = resource.getrlimit(dict[k])
			backup_dict[k] = soft,hard
			print("Backuping and setting soft and hard value: ",k,soft,hard," to " ,v,backup_dict[k][1])
			resource.setrlimit(dict[k], (v, backup_dict[k][1]))
			soft, hard = resource.getrlimit(dict[k])
			print("After setting: ", soft, hard)

	yield
	print("Releasing...")
	for (k,v) in kargs.items():
		if v >= 0 or v == -1:
			resource.setrlimit(dict[k], (backup_dict[k][0], backup_dict[k][1]))

#Using limit resource
with limit_resource(CORE=1, NICE=-2, NOFILE=4, RSS=10, CPU=-1):
	print("working...")
	while True:
		pass

我们可以看到,我们可以把资源限制上下文作为lib文件,引用到其他脚本文件中,通过简单的with语句,可以方便的实现灵活的、可嵌套的资源限制功能。

作者:Yuan
原代码
原文地址:Python资源限制上下文, 感谢原作者分享。