From 77b4c64f88ce961aa4b51b7ec045a35f53ac71df Mon Sep 17 00:00:00 2001 From: Benoit Seignovert Date: Thu, 22 Feb 2024 17:41:58 +0100 Subject: [PATCH] Add spawner progress message --- poetry.lock | 20 ++++++++++++- pyproject.toml | 1 + src/glicid_spawner/progress.py | 54 ++++++++++++++++++++++++++++++++++ src/glicid_spawner/spawner.py | 30 +++++++++++++++++-- tests/test_progress.py | 34 +++++++++++++++++++++ tests/test_spawner.py | 51 +++++++++++++++++++++++++++++++- 6 files changed, 186 insertions(+), 4 deletions(-) create mode 100644 src/glicid_spawner/progress.py create mode 100644 tests/test_progress.py diff --git a/poetry.lock b/poetry.lock index b3415ab..05f1f6b 100644 --- a/poetry.lock +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 0e537ce..9539539 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/glicid_spawner/progress.py b/src/glicid_spawner/progress.py new file mode 100644 index 0000000..b52d5d8 --- /dev/null +++ b/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 diff --git a/src/glicid_spawner/spawner.py b/src/glicid_spawner/spawner.py index 7073489..03504ad 100644 --- a/src/glicid_spawner/spawner.py +++ b/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) diff --git a/tests/test_progress.py b/tests/test_progress.py new file mode 100644 index 0000000..ca67f53 --- /dev/null +++ b/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) diff --git a/tests/test_spawner.py b/tests/test_spawner.py index 59aa2b7..e45b128 100644 --- a/tests/test_spawner.py +++ b/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 + ]