Add spawner progress message

poetry.lock

@@ -1004,6 +1004,24 @@ tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""}
 [package.extras]
 testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
 
+[[package]]
+name = "pytest-asyncio"
+version = "0.23.5"
+description = "Pytest support for asyncio"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pytest-asyncio-0.23.5.tar.gz", hash = "sha256:3a048872a9c4ba14c3e90cc1aa20cbc2def7d01c7c8db3777ec281ba9c057675"},
+    {file = "pytest_asyncio-0.23.5-py3-none-any.whl", hash = "sha256:4e7093259ba018d58ede7d5315131d21923a60f8a6e9ee266ce1589685c89eac"},
+]
+
+[package.dependencies]
+pytest = ">=7.0.0,<9"
+
+[package.extras]
+docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"]
+testing = ["coverage (>=6.2)", "hypothesis (>=5.7.1)"]
+
 [[package]]
 name = "pytest-cov"
 version = "4.1.0"
@@ -1570,4 +1588,4 @@ watchdog = ["watchdog (>=2.3)"]
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "3aee48cbff66c32e4121e9a776d19c466fae94258204526f99eaca50360bbfc7"
+content-hash = "473eb2c87c8df53ff559036126f606add78179a00fc9d1853f47c44d61f844da"

pyproject.toml

@@ -25,6 +25,7 @@ pre-commit = "^3.6.0"
 pytest = "^8.0.0"
 ruff = "^0.1.14"
 pytest-cov = "^4.1.0"
+pytest-asyncio = "^0.23.5"
 flask = "^3.0.2"
 livereload = "<2.5.2"  # FIXME: python-livereload#170
 tornado = "<6.3.0"     # FIXME: python-livereload#270

src/glicid_spawner/progress.py

@@ -0,0 +1,54 @@
+"""Progress messages module."""
+
+from batchspawner import JobStatus
+
+
+def jhp(progress, message):
+    """Jupyterhub progress message dictionary."""
+    return {'progress': progress, 'message': message}
+
+
+class Process:
+    """Progress value."""
+
+    ERROR = jhp(0, '💀 Oops something when wrong')
+    SUBMIT = jhp(10, '📝 Job submitted')
+    PENDING = jhp(20, '⏱️ Your job is pending in queue')
+    INIT = jhp(40, '📦 The resources were allocated')
+    SETUP = jhp(60, '🏗️ Setting up your environment (it should take a minute or two)')
+    CONNECT = jhp(80, '📡 Be ready you should be connected at any moment')
+    TOO_LONG = jhp(95, '🧐 Your instance takes longer than usual but it should be ready soon')
+
+
+class ElapseTime:
+    """Elapse time steps when the job is running."""
+
+    SUBMIT = 0
+    PENDING = 0
+    INIT = 10
+    SETUP = 30
+    CONNECT = 60
+
+
+def get_progress(job_status: JobStatus, elapse_time: int) -> dict:  # noqa: PLR0911 (too-many-return-statements)
+    """Progress and message based on job status and elapse time."""
+    match job_status:
+        case JobStatus.NOTFOUND:
+            return Process.SUBMIT
+
+        case JobStatus.PENDING:
+            return Process.PENDING
+
+        case JobStatus.RUNNING:
+            if elapse_time < ElapseTime.INIT:
+                return Process.INIT
+
+            if elapse_time < ElapseTime.SETUP:
+                return Process.SETUP
+
+            if elapse_time < ElapseTime.CONNECT:
+                return Process.CONNECT
+
+            return Process.TOO_LONG
+
+    return Process.ERROR

src/glicid_spawner/spawner.py

@@ -1,11 +1,13 @@
 """GLiCID spawner module."""
 
+import asyncio
 import re
 
-from batchspawner import SlurmSpawner
-from traitlets import Unicode, default
+from batchspawner import JobStatus, SlurmSpawner
+from traitlets import Integer, Unicode, default
 
 from .form import options_form, options_from_form
+from .progress import ElapseTime, get_progress
 
 
 class GlicidSpawner(SlurmSpawner):
@@ -44,3 +46,27 @@ class GlicidSpawner(SlurmSpawner):
     def options_from_form(self, formdata) -> dict:
         """Export options from form."""
         return options_from_form(formdata)
+
+    progress_rate = Integer(
+        5, help='Interval in seconds at which progress is polled for messages'
+    ).tag(config=True)
+
+    async def progress(self):
+        """Progress bar feedback."""
+        elapse_time = 0
+
+        while True:
+            if self.state_isrunning():
+                job_status = JobStatus.RUNNING
+                elapse_time += self.progress_rate
+            elif self.state_ispending():
+                job_status = JobStatus.PENDING
+            else:
+                job_status = JobStatus.NOTFOUND
+
+            yield get_progress(job_status, elapse_time)
+
+            if elapse_time >= ElapseTime.CONNECT:
+                break
+
+            await asyncio.sleep(self.progress_rate)

tests/test_progress.py

@@ -0,0 +1,34 @@
+"""Test progress messages module."""
+
+import pytest
+from batchspawner import JobStatus
+from glicid_spawner.progress import get_progress, jhp
+
+
+def test_progress_dict():
+    """Test jupyterhub progress message dictionary."""
+    progress = jhp(1, 'Foo')
+
+    assert isinstance(progress, dict)
+    assert progress['progress'] == 1
+    assert progress['message'] == 'Foo'
+
+
+@pytest.mark.parametrize(
+    'job_status, elapse_time, expected',
+    [
+        (JobStatus.UNKNOWN, 0, 0),  # ERROR
+        (JobStatus.NOTFOUND, 0, 10),  # SUBMIT
+        (JobStatus.PENDING, 0, 20),  # PENDING
+        (JobStatus.RUNNING, 0, 40),  # INIT
+        (JobStatus.RUNNING, 10, 60),  # SETUP
+        (JobStatus.RUNNING, 30, 80),  # CONNECT
+        (JobStatus.RUNNING, 60, 95),  # TOO_LONG
+    ],
+)
+def test_progress_msg(job_status, elapse_time, expected):
+    """Test progress getter."""
+    progress = get_progress(job_status, elapse_time)
+
+    assert progress['progress'] == expected
+    assert isinstance(progress['message'], str)

tests/test_spawner.py

@@ -1,17 +1,20 @@
 """Test GLiCID spawner module."""
 
 from collections import namedtuple
+from itertools import repeat
 
 import glicid_spawner.spawner
+import pytest
 from batchspawner import SlurmSpawner
 from glicid_spawner import GlicidSpawner
+from glicid_spawner.spawner import asyncio
 
 User = namedtuple('User', 'name')
 
 
 def test_spawner_config():
     """Test spawner configuration."""
-    spawner = GlicidSpawner()
+    spawner = GlicidSpawner(progress_rate=10)
 
     assert isinstance(spawner, GlicidSpawner)
     assert isinstance(spawner, SlurmSpawner)
@@ -19,6 +22,7 @@ def test_spawner_config():
     assert spawner.batchspawner_singleuser_cmd == 'glicid-spawner-singleuser'
 
     assert spawner.req_qos == 'short'
+    assert spawner.progress_rate == 10
 
 
 def test_spawner_parse_job_id():
@@ -47,3 +51,48 @@ def test_spawner_options_form(monkeypatch):
     assert spawner.user.name == 'john-doe'
     assert spawner.options_form == "options_form('john-doe')"
     assert spawner.options_from_form({'foo': 123}) == "options_from_form({'foo': 123})"
+
+
+@pytest.mark.asyncio
+async def test_spawner_progress(monkeypatch):
+    """Test spawner progress messages."""
+
+    def isrunning():
+        """Running generator values."""
+        yield False  # NOTFOUND
+        yield False  # PENDING
+        yield from repeat(True)  # RUNNING
+
+    def ispending():
+        """Pending generator values."""
+        yield False  # NOTFOUND
+        yield True  # PENDING
+        yield from repeat(False)  # RUNNING (never used)
+
+    # Mock Job status generator
+    iter_ispending = iter(ispending())
+    iter_isrunning = iter(isrunning())
+
+    monkeypatch.setattr(
+        glicid_spawner.GlicidSpawner, 'state_isrunning', lambda _: next(iter_isrunning)
+    )
+    monkeypatch.setattr(
+        glicid_spawner.GlicidSpawner, 'state_ispending', lambda _: next(iter_ispending)
+    )
+
+    async def mock_sleep(_):
+        """Mock asyncio sleep."""
+
+    monkeypatch.setattr(asyncio, 'sleep', mock_sleep)
+
+    spawner = GlicidSpawner(progress_rate=20)
+
+    progress = [msg['progress'] async for msg in spawner.progress()]
+
+    assert progress == [
+        10,  # submit
+        20,  # pending
+        60,  # running | elapse time = 20 -> setup
+        80,  # running | elapse time = 40 -> connect
+        95,  # running | elapse time = 60 -> too long
+    ]