Skip to content

Commit

Permalink
Move strings into variables to be reused across the Spawners
Browse files Browse the repository at this point in the history
  • Loading branch information
@raethlein
raethlein committed Dec 2, 2019
1 parent b957c2a commit b624bdf
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 35 deletions.
28 changes: 11 additions & 17 deletions resources/mlhubspawner/mlhubspawner/mlhubkubernetesspawner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,10 +60,10 @@ def get_env(self):
#if self.user_options.get('gpus'):
# env['NVIDIA_VISIBLE_DEVICES'] = self.user_options.get('gpus')

if self.user_options.get('cpu_limit'):
env["MAX_NUM_THREADS"] = self.user_options.get('cpu_limit')
if self.user_options.get(utils.OPTION_CPU_LIMIT):
env[utils.OPTION_MAX_NUM_THREADS] = self.user_options.get(utils.OPTION_CPU_LIMIT)

env['SSH_JUMPHOST_TARGET'] = self.pod_name
env[utils.OPTION_SSH_JUMPHOST_TARGET] = self.pod_name

return env

Expand All @@ -73,19 +73,19 @@ def start(self):

self.saved_user_options = self.user_options

if self.user_options.get('image'):
self.image = self.user_options.get('image')
if self.user_options.get(utils.OPTION_IMAGE):
self.image = self.user_options.get(utils.OPTION_IMAGE)

# Set request explicitly to 0, otherwise Kubernetes will set it to the same amount as limit
# self.cpu_guarantee / self.mem_guarantee cannot be directly used, as they are of type ByteSpecification and, for example, 0G will be transformed to 0 which will not pass
# the 'if cpu_guarantee' check (see https://github.com/jupyterhub/kubespawner/blob/8a6d66e04768565c0fc56c790a5fc42bfee634ec/kubespawner/objects.py#L279).
# Hence, set it via extra_resource_guarantees.
self.extra_resource_guarantees = {"cpu": 0, "memory": "0G"}
if self.user_options.get('cpu_limit'):
self.cpu_limit = float(self.user_options.get('cpu_limit'))
if self.user_options.get(utils.OPTION_CPU_LIMIT):
self.cpu_limit = float(self.user_options.get(utils.OPTION_CPU_LIMIT))

if self.user_options.get('mem_limit'):
memory = str(self.user_options.get('mem_limit')) + "G"
if self.user_options.get(utils.OPTION_MEM_LIMIT):
memory = str(self.user_options.get(utils.OPTION_MEM_LIMIT)) + "G"
self.mem_limit = memory.upper().replace("GB", "G").replace("KB", "K").replace("MB", "M").replace("TB", "T")

#if self.user_options.get('is_mount_volume') == 'on':
Expand All @@ -94,8 +94,8 @@ def start(self):

# set default label 'origin' to know for sure which containers where started via the hub
#self.extra_labels['pod_name'] = self.pod_name
if self.user_options.get('days_to_live'):
days_to_live_in_seconds = int(self.user_options.get('days_to_live')) * 24 * 60 * 60 # days * hours_per_day * minutes_per_hour * seconds_per_minute
if self.user_options.get(utils.OPTION_DAYS_TO_LIVE):
days_to_live_in_seconds = int(self.user_options.get(utils.OPTION_DAYS_TO_LIVE)) * 24 * 60 * 60 # days * hours_per_day * minutes_per_hour * seconds_per_minute
expiration_timestamp = time.time() + days_to_live_in_seconds
self.extra_labels[utils.LABEL_EXPIRATION_TIMESTAMP] = str(expiration_timestamp)
else:
Expand Down Expand Up @@ -151,12 +151,6 @@ def stop(self, now=False):
except:
self.log.warn("Could not delete service with name {}".format(self.pod_name))


def get_container_metadata(self) -> str:
if self.pod_name is None or self.pod_name == '':
return ""

return utils.get_container_metadata(self)

def get_workspace_config(self) -> str:
return utils.get_workspace_config(self)
Expand Down
38 changes: 20 additions & 18 deletions resources/mlhubspawner/mlhubspawner/mlhubspawner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
INITIAL_CIDR_SECOND_OCTET = 33
INITIAL_CIDR = "{}.33.0.0/24".format(INITIAL_CIDR_FIRST_OCTET)

OPTION_SHM_SIZE = "shm_size"

def has_complete_network_information(network):
"""Convenient function to check whether the docker.Network object has all required properties.
Expand Down Expand Up @@ -134,10 +136,10 @@ def get_env(self):
if self.user_options.get('gpus'):
env['NVIDIA_VISIBLE_DEVICES'] = self.user_options.get('gpus')

if self.user_options.get('cpu_limit'):
env["MAX_NUM_THREADS"] = self.user_options.get('cpu_limit')
if self.user_options.get(utils.OPTION_CPU_LIMIT):
env[utils.OPTION_MAX_NUM_THREADS] = self.user_options.get(utils.OPTION_CPU_LIMIT)

env['SSH_JUMPHOST_TARGET'] = self.object_name
env[utils.OPTION_SSH_JUMPHOST_TARGET] = self.object_name

return env

Expand All @@ -151,22 +153,22 @@ def start(self) -> (str, int):

self.saved_user_options = self.user_options

if self.user_options.get('image'):
self.image = self.user_options.get('image')
if self.user_options.get(utils.OPTION_IMAGE):
self.image = self.user_options.get(utils.OPTION_IMAGE)

extra_host_config = {}
if self.user_options.get('cpu_limit'):
if self.user_options.get(utils.OPTION_CPU_LIMIT):
# nano_cpus cannot be bigger than the number of CPUs of the machine (this method would currently not work in a cluster, as machines could be different than the machine where the runtime-manager and this code run.
max_available_cpus = self.resource_information["cpu_count"]
limited_cpus = min(
int(self.user_options.get('cpu_limit')), max_available_cpus)
int(self.user_options.get(utils.OPTION_CPU_LIMIT)), max_available_cpus)

# the nano_cpu parameter of the Docker client expects an integer, not a float
nano_cpus = int(limited_cpus * 1e9)
extra_host_config['nano_cpus'] = nano_cpus
if self.user_options.get('mem_limit'):
extra_host_config['mem_limit'] = str(self.user_options.get(
'mem_limit')) + "gb"
if self.user_options.get(utils.OPTION_MEM_LIMIT):
extra_host_config[utils.OPTION_MEM_LIMIT] = str(self.user_options.get(
utils.OPTION_MEM_LIMIT)) + "gb"

if self.user_options.get('is_mount_volume') == 'on':
# {username} and {servername} will be automatically replaced by DockerSpawner with the right values as in template_namespace
Expand All @@ -176,20 +178,20 @@ def start(self) -> (str, int):

extra_create_kwargs = {}
# set default label 'origin' to know for sure which containers where started via the hub
extra_create_kwargs['labels'] = self.default_labels
if self.user_options.get('days_to_live'):
days_to_live_in_seconds = int(self.user_options.get('days_to_live')) * 24 * 60 * 60 # days * hours_per_day * minutes_per_hour * seconds_per_minute
extra_create_kwargs[utils.OPTION_LABELS] = self.default_labels
if self.user_options.get(utils.OPTION_DAYS_TO_LIVE):
days_to_live_in_seconds = int(self.user_options.get(utils.OPTION_DAYS_TO_LIVE)) * 24 * 60 * 60 # days * hours_per_day * minutes_per_hour * seconds_per_minute
expiration_timestamp = time.time() + days_to_live_in_seconds
extra_create_kwargs['labels'][utils.LABEL_EXPIRATION_TIMESTAMP] = str(expiration_timestamp)
extra_create_kwargs[utils.OPTION_LABELS][utils.LABEL_EXPIRATION_TIMESTAMP] = str(expiration_timestamp)
else:
extra_create_kwargs['labels'][utils.LABEL_EXPIRATION_TIMESTAMP] = str(0)
extra_create_kwargs[utils.OPTION_LABELS][utils.LABEL_EXPIRATION_TIMESTAMP] = str(0)

if self.user_options.get('shm_size'):
extra_host_config['shm_size'] = self.user_options.get('shm_size')
if self.user_options.get(OPTION_SHM_SIZE):
extra_host_config[OPTION_SHM_SIZE] = self.user_options.get('shm_size')

if self.user_options.get('gpus'):
extra_host_config['runtime'] = "nvidia"
extra_create_kwargs['labels'][utils.LABEL_NVIDIA_VISIBLE_DEVICES] = self.user_options.get('gpus')
extra_create_kwargs[utils.OPTION_LABELS][utils.LABEL_NVIDIA_VISIBLE_DEVICES] = self.user_options.get('gpus')

self.extra_host_config.update(extra_host_config)
self.extra_create_kwargs.update(extra_create_kwargs)
Expand Down
11 changes: 11 additions & 0 deletions resources/mlhubspawner/mlhubspawner/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,17 @@

ENV_HUB_NAME = os.getenv("HUB_NAME", "mlhub")

OPTION_LABELS = "labels"
OPTION_DAYS_TO_LIVE = "days_to_live"
OPTION_NANO_CPUS = "nano_cpus"
OPTION_CPU_LIMIT = "cpu_limit"
OPTION_MEM_LIMIT = "mem_limit"
OPTION_IMAGE = "image"
OPTION_SHM_SIZE = "shm_size"
OPTION_SSH_JUMPHOST_TARGET = "SSH_JUMPHOST_TARGET"
OPTION_MAX_NUM_THREADS = "MAX_NUM_THREADS"


def get_lifetime_timestamp(labels: dict) -> float:
return float(labels.get(LABEL_EXPIRATION_TIMESTAMP, '0'))

Expand Down

0 comments on commit b624bdf

Please sign in to comment.