Merge pull request #102 from kellyrowland/configurable-cull

Configurable cull arbitration - continuation of #90

Author: yuvipanda

jupyterhub_idle_culler/__init__.py

@@ -20,9 +20,11 @@
 from tornado.httputil import url_concat
 from tornado.ioloop import IOLoop, PeriodicCallback
 from tornado.log import LogFormatter
-from traitlets import Bool, Int, Unicode, default
+from traitlets import Bool, Callable, Int, Unicode, default
 from traitlets.config import Application
 
+from .utils import maybe_future
+
 __version__ = "1.4.1.dev"
 
 STATE_FILTER_MIN_VERSION = V("1.3.0")
@@ -78,6 +80,11 @@ def utcnow():
     return datetime.now(timezone.utc)
 
 
+def default_cull_arbiter(*, inactive, inactive_limit, server, **kwargs):
+    """Return True if time inactive exceeds limit, the classic cull check."""
+    return inactive.total_seconds() >= inactive_limit
+
+
 async def cull_idle(
     url,
     api_token,
@@ -93,6 +100,7 @@ async def cull_idle(
     api_page_size=0,
     cull_default_servers=True,
     cull_named_servers=True,
+    cull_arbiter=default_cull_arbiter,
 ):
     """Shutdown idle single-user servers
 
@@ -226,26 +234,18 @@ async def handle_server(user, server_name, server, max_age, inactive_limit):
             # for running servers
             inactive = age
 
-        # CUSTOM CULLING TEST CODE HERE
-        # Add in additional server tests here.  Return False to mean "don't
-        # cull", True means "cull immediately", or, for example, update some
-        # other variables like inactive_limit.
-        #
-        # Here, server['state'] is the result of the get_state method
-        # on the spawner.  This does *not* contain the below by
-        # default, you may have to modify your spawner to make this
-        # work.  The `user` variable is the user model from the API.
-        #
-        # if server['state']['profile_name'] == 'unlimited'
-        #     return False
-        # inactive_limit = server['state']['culltime']
-
         is_default_server = server_name == ""
         is_named_server = server_name != ""
 
+        cull_result = await maybe_future(
+            cull_arbiter(
+                inactive=inactive, inactive_limit=inactive_limit, server=server
+            )
+        )
+
         should_cull = (
             inactive is not None
-            and inactive.total_seconds() >= inactive_limit
+            and cull_result
             and (
                 (cull_default_servers and is_default_server)
                 or (cull_named_servers and is_named_server)
@@ -508,6 +508,36 @@ class IdleCuller(Application):
         config=True,
     )
 
+    cull_arbiter_hook = Callable(
+        default_cull_arbiter,
+        help=dedent("""
+            Enable custom culling logic.
+
+            By default, the idle culler compares a server's time since last
+            activity with a specified idle time limit.  This hook allows for
+            additional or more arbitrary logic when deciding whether to cull a
+            server or not.  To customize culling logic, define a callable taking
+            3 arguments:
+
+                def my_cull_arbiter(inactive, inactive_limit, server):
+                    if server['state']['profile_name'] == 'unlimited':
+                        return False
+                    return inactive.total_seconds() >= inactive_limit
+                c.IdleCuller.cull_arbiter_hook = my_cull_arbiter
+
+            - 'inactive' is the server's time since last activity, a timedelta object
+            - 'inactive_limit' is the idle timeout limit, in seconds
+            - 'server' is the server being considered for culling
+
+            This callable should return True if the server should be culled, and
+            False if it should not.  In this example, servers with a profile
+            name of "unlimited" are never culled, but all others are subject to
+            the default time limit logic.
+            """).strip(),
+    ).tag(
+        config=True,
+    )
+
     cull_every = Int(
         0,
         help=dedent("""
@@ -658,6 +688,8 @@ def start(self):
 
         api_token = os.environ["JUPYTERHUB_API_TOKEN"]
 
+        cull_arbiter = self.cull_arbiter_hook
+
         try:
             AsyncHTTPClient.configure("tornado.curl_httpclient.CurlAsyncHTTPClient")
         except ImportError as e:
@@ -683,6 +715,7 @@ def start(self):
             api_page_size=self.api_page_size,
             cull_default_servers=self.cull_default_servers,
             cull_named_servers=self.cull_named_servers,
+            cull_arbiter=cull_arbiter,
         )
         # schedule first cull immediately
         # because PeriodicCallback doesn't start until the end of the first interval

jupyterhub_idle_culler/utils.py

@@ -0,0 +1,36 @@
+"""Miscellaneous utilities"""
+
+import asyncio
+import concurrent.futures
+import inspect
+
+
+def maybe_future(obj):
+    """Return an asyncio Future
+
+    Use instead of gen.maybe_future
+
+    For our compatibility, this must accept:
+
+    - asyncio coroutine (gen.maybe_future doesn't work in tornado < 5)
+    - tornado coroutine (asyncio.ensure_future doesn't work)
+    - scalar (asyncio.ensure_future doesn't work)
+    - concurrent.futures.Future (asyncio.ensure_future doesn't work)
+    - tornado Future (works both ways)
+    - asyncio Future (works both ways)
+
+    Borrowed from https://github.com/jupyterhub/jupyterhub/blob/main/jupyterhub/utils.py
+    since duplicating the one function here is lighter than adding jupyterhub as a
+    dependency.
+    """
+    if inspect.isawaitable(obj):
+        # already awaitable, use ensure_future
+        return asyncio.ensure_future(obj)
+    elif isinstance(obj, concurrent.futures.Future):
+        return asyncio.wrap_future(obj)
+    else:
+        # could also check for tornado.concurrent.Future
+        # but with tornado >= 5.1 tornado.Future is asyncio.Future
+        f = asyncio.Future()
+        f.set_result(obj)
+        return f

tests/test_idle_culler.py

@@ -30,6 +30,14 @@ async def count_active_users(admin_request):
     return active_users
 
 
+def cull_arbiter_function(inactive, inactive_limit, server):
+    return True
+
+
+async def async_cull_arbiter_function(inactive, inactive_limit, server):
+    return False
+
+
 async def test_cull_idle(cull_idle, start_users, admin_request):
     assert await count_active_users(admin_request) == 0
     await start_users(3)
@@ -46,5 +54,28 @@ async def test_cull_idle(cull_idle, start_users, admin_request):
     assert await count_active_users(admin_request) == 0
 
 
+async def test_custom_cull_arbiter(cull_idle, start_users, admin_request):
+    assert await count_active_users(admin_request) == 0
+    await start_users(3)
+    assert await count_active_users(admin_request) == 3
+    await cull_idle(
+        inactive_limit=300, logger=app_log, cull_arbiter=cull_arbiter_function
+    )
+    # time has not passed but the cull arbiter function returns true
+    # so, everyone culled
+    assert await count_active_users(admin_request) == 0
+
+
+async def test_async_custom_cull_arbiter(cull_idle, start_users, admin_request):
+    assert await count_active_users(admin_request) == 0
+    await start_users(3)
+    assert await count_active_users(admin_request) == 3
+    # test an async arbiter that just returns false to verify that no changes are made
+    await cull_idle(
+        inactive_limit=300, logger=app_log, cull_arbiter=async_cull_arbiter_function
+    )
+    assert await count_active_users(admin_request) == 3
+
+
 def test_help():
     check_output([sys.executable, "-m", "jupyterhub_idle_culler", "--help"])