diff --git a/systemdspawner/systemdspawner.py b/systemdspawner/systemdspawner.py index 2f4ada1..cebdc9f 100644 --- a/systemdspawner/systemdspawner.py +++ b/systemdspawner/systemdspawner.py @@ -3,7 +3,7 @@ import time import subprocess import shlex -from traitlets import Bool, Int, Unicode, List +from traitlets import Bool, Int, Unicode, List, Dict, Set from tornado import gen from jupyterhub.spawner import Spawner @@ -70,6 +70,21 @@ class SystemdSpawner(Spawner): """ ).tag(config=True) + users_without_quotas = Set( + None, + allow_none=True, + help=''' + Set of users that will have unlimited mem_limit and cpu_limit, it will overwrite special_user_quotas even if + it's set for the user. + ''' + ).tag(config=True) + + special_user_quotas = Dict( + None, + allow_none=True, + help='Dict with users as key and cpu_limit/mem_limit as value. It will overwrite mem_limit/cpu_limit.' + ).tag(config=True) + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # All traitlets configurables are configured by now @@ -83,7 +98,6 @@ def __init__(self, *args, **kwargs): self.log.debug('user:%s Initialized spawner with unit %s', self.user.name, self.unit_name) - def _expand_user_vars(self, string): """ Expand user related variables in a given string @@ -180,6 +194,17 @@ def start(self): cmd.append('--setenv=SHELL={shell}'.format(shell=self.default_shell)) + if self.users_without_quotas and self.user.name in self.users_without_quotas: + # If users_without_quotas is set for the user, it will overwrite mem_limit and cpu_limit + self.mem_limit = None + self.cpu_limit = None + elif self.special_user_quotas and self.user.name in self.special_user_quotas: + # If the user has a specific value for mem_limit or cpu_limit, it will overwrite the default value + if 'mem_limit' in self.special_user_quotas[self.user.name]: + self.mem_limit = self.special_user_quotas[self.user.name]['mem_limit'] + if 'cpu_limit' in self.special_user_quotas[self.user.name]: + self.cpu_limit = self.special_user_quotas[self.user.name]['cpu_limit'] + if self.mem_limit is not None: # FIXME: Detect & use proper properties for v1 vs v2 cgroups cmd.extend([