Implement the function to get mesos slaves

tests/test_mesosslave.py

@@ -0,0 +1,22 @@
+import json
+import types
+
+import pytest
+
+from marathon.models.constraint import MarathonConstraint
+
+from ucrspawner import mesosslave
+from ucrspawner.exceptions import UCRSpawnerException
+
+def test_match(monkeypatch):
+    mesos_master_response = json.loads('{"id":"40237088-f509-46a4-a5f1-c8f88ca8e03f-S0","hostname":"mesos-slave","port":5051,"attributes":{"foo":123.0,"bar":"wow"},"pid":"slave(1)@172.24.0.5:5051","registered_time":1518339392.10703,"resources":{"disk":69716.0,"mem":999.0,"gpus":0.0,"cpus":2.0,"ports":"[31000-32000]"},"used_resources":{"disk":0.0,"mem":0.0,"gpus":0.0,"cpus":0.0},"offered_resources":{"disk":0.0,"mem":0.0,"gpus":0.0,"cpus":0.0},"reserved_resources":{},"unreserved_resources":{"disk":69716.0,"mem":999.0,"gpus":0.0,"cpus":2.0,"ports":"[31000-32000]"},"active":true,"version":"1.4.0","capabilities":["MULTI_ROLE","HIERARCHICAL_ROLE","RESERVATION_REFINEMENT"],"reserved_resources_full":{},"unreserved_resources_full":[{"name":"cpus","type":"SCALAR","scalar":{"value":2.0},"role":"*"},{"name":"mem","type":"SCALAR","scalar":{"value":999.0},"role":"*"},{"name":"disk","type":"SCALAR","scalar":{"value":69716.0},"role":"*"},{"name":"ports","type":"RANGES","ranges":{"range":[{"begin":31000,"end":32000}]},"role":"*"}],"used_resources_full":[],"offered_resources_full":[]}')
+
+    slave = mesosslave.MesosSlave(mesos_master_response)
+
+    assert slave.match(MarathonConstraint('foo', 'LIKE', '123'))
+    assert not slave.match(MarathonConstraint('foo', 'UNLIKE', '123'))
+    assert slave.match(MarathonConstraint('bar', 'LIKE', 'wo[w]'))
+    assert not slave.match(MarathonConstraint('bar', 'LIKE', 'wof'))
+    with pytest.raises(UCRSpawnerException) as excinfo:
+        slave.match(MarathonConstraint('foo', 'UNKNOWN', ''))
+    assert str(excinfo.value) == 'Unsupported constraint operator: UNKNOWN'

ucrspawner/mesosslave.py

@@ -0,0 +1,38 @@
+import re
+
+from .exceptions import UCRSpawnerException
+
+
+class MesosSlave():
+    def __init__(self, params):
+        self.hostname = params['hostname']
+        self.attributes = params['attributes']
+
+        self.disk = params['resources']['disk']
+        self.cpus = params['resources']['cpus']
+        self.mem = params['resources']['mem']
+        self.disk = params['resources']['disk']
+        self.gpus = int(params['resources']['gpus'])
+
+        self.available_cpus = params['unreserved_resources']['cpus']
+        self.available_mem = params['unreserved_resources']['mem']
+        self.available_disk = params['unreserved_resources']['disk']
+        self.available_gpus = int(params['unreserved_resources']['gpus'])
+
+    def is_available(self):
+        return self.available_cpus > 0 and \
+            self.available_mem > 0 and \
+            self.available_disk > 0 and \
+            (self.gpus == 0 or self.available_gpus > 0)
+
+    def match(self, constraint):
+        if constraint.field == 'hostname':
+            value = self.hostname
+        else:
+            value = str(self.attributes[constraint.field])
+        if constraint.operator == 'LIKE':
+            return re.compile(constraint.value).match(value) is not None
+        elif constraint.operator == 'UNLIKE':
+            return re.compile(constraint.value).match(value) is None
+        else:
+            raise UCRSpawnerException('Unsupported constraint operator: %s' % constraint.operator)

ucrspawner/ucrspawner.py

@@ -3,7 +3,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timedelta
 from textwrap import dedent
-from urllib.parse import urlparse, urlunparse
+from urllib.parse import urljoin, urlparse, urlunparse
 
 import jupyterhub
 from jupyterhub.spawner import Spawner
@@ -18,13 +18,16 @@
     MarathonDockerContainer,
 )
 
+import requests
+
 from tornado import gen
 from tornado.concurrent import run_on_executor
 
 from traitlets import Any, Float, Integer, List, Unicode, default, observe
 
-from .volumenaming import default_format_volume_name
 from .exceptions import UCRSpawnerException
+from .mesosslave import MesosSlave
+from .volumenaming import default_format_volume_name
 
 _jupyterhub_xy = '%i.%i' % (jupyterhub.version_info[:2])
 
@@ -160,7 +163,7 @@ def __init__(self, *args, **kwargs):
         super(UCRSpawner, self).__init__(*args, **kwargs)
         self.marathon = MarathonClient(self.marathon_host)
         self.get_state()
-        self.env_keep = [] # env_keep doesn't make sense in Mesos
+        self.env_keep = []  # env_keep doesn't make sense in Mesos
 
     @property
     def app_id(self):
@@ -204,7 +207,9 @@ def get_constraints(self):
         constraints = [MarathonConstraint.from_json(c) for c in self.marathon_constraints]
         unsupported = [c for c in constraints if c.operator not in ('LIKE', 'UNLIKE')]
         if len(unsupported) > 0:
-            raise UCRSpawnerException('Unsupported constraint operators: %s' % sorted(list(set([c.operator for c in unsupported]))))
+            raise UCRSpawnerException(
+                'Unsupported constraint operators: %s'
+                % sorted(list(set([c.operator for c in unsupported]))))
         return constraints
 
     def get_ip_and_port(self, app_info):
@@ -456,3 +461,16 @@ def poll(self):
                 return 0
 
         return None
+
+    def get_mesos_slaves(self):
+        info = self.marathon.get_info()
+
+        headers = {
+            'Accept': 'application/json'
+        }
+        response = requests.get(urljoin(info.marathon_config.mesos_leader_ui_url, '/slaves'), headers=headers)
+        json = response.json()
+        slaves = [MesosSlave(j) for j in json['slaves']]
+        for c in self.get_constraints():
+            slaves = [s for s in slaves if s.match(c)]
+        return slaves