Python resource.RLIMIT_AS Examples

def report_resource_limits(logger: logging.Logger) -> None:
    resources = [
        ('CPU time (seconds)', resource.RLIMIT_CPU),
        ('Heap size (bytes)', resource.RLIMIT_DATA),
        ('Num. process', resource.RLIMIT_NPROC),
        ('Num. files', resource.RLIMIT_NOFILE),
        ('Address space', resource.RLIMIT_AS),
        ('Locked address space', resource.RLIMIT_MEMLOCK)
    ]
    resource_limits = [
        (name, resource.getrlimit(res)) for (name, res) in resources
    ]
    resource_s = '\n'.join([
        '* {}: {}'.format(res, lim) for (res, lim) in resource_limits
    ])
    logger.info("resource limits:\n%s", indent(resource_s, 2)) 

def getMemoryLimit():
        try:
            limit = getrlimit(RLIMIT_AS)[0]
            if 0 < limit:
                limit *= PAGE_SIZE
            return limit
        except ValueError:
            return None 

def setMemoryLimit(max_mem):
        if max_mem is None:
            max_mem = -1
        try:
            setrlimit(RLIMIT_AS, (max_mem, -1))
            return True
        except ValueError:
            return False 

def getMemoryLimit():
        try:
            limit = getrlimit(RLIMIT_AS)[0]
            if 0 < limit:
                limit *= PAGE_SIZE
            return limit
        except ValueError:
            return None 

def setMemoryLimit(max_mem):
        if max_mem is None:
            max_mem = -1
        try:
            setrlimit(RLIMIT_AS, (max_mem, -1))
            return True
        except ValueError:
            return False 

def test_leak_load(self):
        from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
        setrlimit(RLIMIT_STACK, (stack_size, stack_size))
        setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
        for _ in range(iterations):
            with Image.open(test_file) as im:
                im.load() 

def test_leak_save(self):
        from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
        setrlimit(RLIMIT_STACK, (stack_size, stack_size))
        setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
        for _ in range(iterations):
            with Image.open(test_file) as im:
                im.load()
                test_output = BytesIO()
                im.save(test_output, "JPEG2000")
                test_output.seek(0)
                test_output.read() 

def make_memsafe():
    if sys.platform.startswith('linux'):
        import resource
        import psutil
        for rsrc in (resource.RLIMIT_AS, resource.RLIMIT_DATA):
            freemem = psutil.virtual_memory().free
            hard = int(round(freemem *.8))
            soft = hard
            resource.setrlimit(rsrc, (soft, hard)) #limit to 80% of available memory 

def make_memsafe():
    if sys.platform.startswith('linux'):
        import resource
        import psutil
        for rsrc in (resource.RLIMIT_AS, resource.RLIMIT_DATA):
            freemem = psutil.virtual_memory().free
            hard = int(round(freemem *.8))
            soft = hard
            resource.setrlimit(rsrc, (soft, hard)) #limit to 80% of available memory 

def limitMemory(nbytes):
        setrlimit(RLIMIT_AS, (nbytes, -1)) 

def createChild(arguments, no_stdout, env=None):
    """
    Create a child process:
     - arguments: list of string where (eg. ['ls', '-la'])
     - no_stdout: if True, use null device for stdout/stderr
     - env: environment variables dictionary

    Use:
     - env={} to start with an empty environment
     - env=None (default) to copy the environment
    """

    # Fork process
    pid = fork()
    if pid:
        return pid
    else:
        # print "limit",getrlimit(RLIMIT_DATA)
        setrlimit(RLIMIT_AS, (1024 * 1024 * 1024, -1))
        # print "limit",getrlimit(RLIMIT_DATA)

        try:
            ptrace_traceme()
        except PtraceError as err:
            raise ChildError(str(err))

        _execChild(arguments, no_stdout, env)
        exit(255) 

def limit_address_space(self):
        max_memory = self.memory_limit(resource.RLIMIT_AS)
        return processutils.ProcessLimits(address_space=max_memory) 

def test_address_space(self):
        prlimit = self.limit_address_space()
        self.check_limit(prlimit, 'RLIMIT_AS', prlimit.address_space) 

def setMemoryLimit(max_mem):
        if max_mem is None:
            max_mem = -1
        try:
            setrlimit(RLIMIT_AS, (max_mem, -1))
            return True
        except ValueError:
            return False 

def setProcessLimits(x):
        # This is called after fork and before exec. Any messages
        # printed here will look like the program that we are calling
        # printed them out.

        #print("pre switch user")
        if switchUser:
            if os.geteuid() != 0:
                print("We can't switch to a different user if this script is not run as root.")
                exit(1)
            # If we are supposed to switch to a different user account to do the grading
            if os.geteuid() == 0:
                os.setreuid(autograderUid,autograderUid)
            if os.geteuid() == 0:
                print("Still root after trying to switch to autograder user?")
                exit(1)
        else:
            # If we are not switching to a different user, make sure that we don't run as root.
            if os.geteuid() == 0:
                print("Halting. Do not run submitted programs as root.")
                exit(1)
                
        #print("post switch user")
        
        #print("Preexec start")
        if os.setpgrp() == -1:  # put all processes in the same process group so we can kill it and all children it creates.
            print("Failed to set process group!")
        #print("Preexec middle")

        def limitHelper(limitType, limit):
            # limit is a string referring to a previously set environment variable.
            if limit in os.environ:
                limit = int(os.environ[limit])
                (soft, hard) = resource.getrlimit(limitType)
                #print("soft %d, hard %d, requested %d\n" % (soft, hard, limit))
                if hard > 0 and limit > hard:
                    limit = hard
                resource.setrlimit(limitType, (limit, limit))

        limitHelper(resource.RLIMIT_NPROC, "ULIMIT_NPROC")
        limitHelper(resource.RLIMIT_AS, "ULIMIT_AS")
        limitHelper(resource.RLIMIT_DATA, "ULIMIT_DATA") 

def getMemoryLimit():
        try:
            limit = getrlimit(RLIMIT_AS)[0]
            if 0 < limit:
                limit *= PAGE_SIZE
            return limit
        except ValueError:
            return None