Merge pull request #39 from lsst-sqre/tickets/DM-35203

DM-35203: Make the worker's image selection configurable

CHANGELOG.rst

@@ -2,11 +2,15 @@
 Change log
 ##########
 
-Unreleased
-==========
+0.4.0 (2022-06-15)
+==================
 
 - The worker identity configuration can now omit the ``uid`` field for environments where Gafaelfawr is able to assign a UID (e.g. through an LDAP backend).
-- A new ``NOTEBURST_WORKER_TOKEN_LIFETIME`` environment variable enables you to configure the lifetime of the workers' authentication tokens. The default matches the existing behavior, 28 days.
+- New configurations for workers:
+    - The new ``NOTEBURST_WORKER_TOKEN_LIFETIME`` environment variable enables you to configure the lifetime of the workers' authentication tokens. The default matches the existing behavior, 28 days.
+    - ``NOTEBURST_WORKER_TOKEN_SCOPES`` environment variable enables you to set what token scopes the nublado2 bot users should have, as a comma-separated list.
+    - ``NOTEBURST_WORKER_IMAGE_SELECTOR`` allows you to specify what stream of Nublado image to select. Can be ``recommended``, ``weekly`` or ``reference``. If the latter, you can specify the specific Docker Image with ``NOTEBURST_WORKER_IMAGE_REFERENCE``.
+    - The ``NOTEBURST_WORKER_KEEPALIVE`` configuration controls whether the worker keep alive function is run (to default the Nublado pod culler), and at what frequencey. Set to ``disabled`` to disable; ``fast`` to run every 30 seconds; or ``normal`` to run every 5 minutes.
 - Noteburst now uses the arq client and dependency from Safir 3.2, which was originally developed from Noteburst.
 
 0.3.0 (2022-05-24)

src/noteburst/config.py

@@ -4,11 +4,18 @@
 
 from enum import Enum
 from pathlib import Path
-from typing import List
+from typing import Any, List, Mapping, Optional
 from urllib.parse import urlparse
 
 from arq.connections import RedisSettings
-from pydantic import BaseSettings, Field, HttpUrl, RedisDsn, SecretStr
+from pydantic import (
+    BaseSettings,
+    Field,
+    HttpUrl,
+    RedisDsn,
+    SecretStr,
+    validator,
+)
 from safir.arq import ArqMode
 
 __all__ = ["Config", "Profile", "LogLevel"]
@@ -34,6 +41,32 @@ class LogLevel(str, Enum):
     CRITICAL = "CRITICAL"
 
 
+class JupyterImageSelector(str, Enum):
+    """Possible ways of selecting a JupyterLab image."""
+
+    recommended = "recommended"
+    """Currently recommended image."""
+
+    weekly = "weekly"
+    """Current weekly image."""
+
+    reference = "reference"
+    """Select a specific image by reference."""
+
+
+class WorkerKeepAliveSetting(str, Enum):
+    """Modes for the worker keep-alive function."""
+
+    disabled = "disabled"
+    """Do not run a keep-alive function."""
+
+    fast = "fast"
+    """Run the keep-alive function at a high frequency (every 30 seconds)."""
+
+    normal = "normal"
+    """Run the keep-alive function at a slower frequencey (i.e. 5 minutes)."""
+
+
 class Config(BaseSettings):
 
     name: str = Field("Noteburst", env="SAFIR_NAME")
@@ -108,11 +141,63 @@ class WorkerConfig(Config):
         description="Worker auth token lifetime in seconds.",
     )
 
+    worker_token_scopes: str = Field(
+        "exec:notebook",
+        env="NOTEBURST_WORKER_TOKEN_SCOPES",
+        description=(
+            "Worker (nublado2 pod) token scopes as a comma-separated string."
+        ),
+    )
+
+    image_selector: JupyterImageSelector = Field(
+        JupyterImageSelector.recommended,
+        env="NOTEBURST_WORKER_IMAGE_SELECTOR",
+        description="Method for selecting a Jupyter image to run.",
+    )
+
+    image_reference: Optional[str] = Field(
+        None,
+        env="NOTEBURST_WORKER_IMAGE_REFERENCE",
+        description=(
+            "Docker image reference, if NOTEBURST_WORKER_IMAGE_SELECTOR is "
+            "``reference``."
+        ),
+    )
+
+    worker_keepalive: WorkerKeepAliveSetting = Field(
+        WorkerKeepAliveSetting.normal, env="NOTEBURST_WORKER_KEEPALIVE"
+    )
+
     @property
     def aioredlock_redis_config(self) -> List[str]:
         """Redis configurations for aioredlock."""
         return [str(self.identity_lock_redis_url)]
 
+    @validator("image_reference")
+    def is_image_ref_set(
+        cls, v: Optional[str], values: Mapping[str, Any]
+    ) -> Optional[str]:
+        """Validate that image_reference is set if image_selector is
+        set to reference.
+        """
+        if (
+            v is None
+            and values["image_selector"] == JupyterImageSelector.reference
+        ):
+            raise ValueError(
+                "Set NOTEBURST_WORKER_IMAGE_REFERENCE since "
+                "NOTEBURST_WORKER_IMAGE_SELECTOR is ``reference``."
+            )
+
+        return v
+
+    @property
+    def parsed_worker_token_scopes(self) -> List[str]:
+        """Sequence of worker token scopes, parsed from the comma-separated
+        list in `worker_token_scopes`.
+        """
+        return [t.strip() for t in self.worker_token_scopes.split(",") if t]
+
 
 config = Config()
 """Configuration for noteburst."""

src/noteburst/jupyterclient/jupyterlab.py

@@ -9,7 +9,6 @@
 import random
 import string
 from dataclasses import dataclass
-from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -26,6 +25,7 @@
 import websockets.typing
 from websockets.exceptions import WebSocketException
 
+from noteburst.config import JupyterImageSelector
 from noteburst.config import config as noteburst_config
 
 from .cachemachine import CachemachineClient, JupyterImage
@@ -37,7 +37,6 @@
     from .user import AuthenticatedUser
 
 __all__ = [
-    "JupyterImageSelector",
     "SpawnProgressMessage",
     "JupyterSpawnProgress",
     "JupyterLabSession",
@@ -48,14 +47,6 @@
 ]
 
 
-class JupyterImageSelector(Enum):
-    """Possible ways of selecting a JupyterLab image."""
-
-    RECOMMENDED = "recommended"
-    LATEST_WEEKLY = "latest-weekly"
-    BY_REFERENCE = "by-reference"
-
-
 @dataclass(frozen=True)
 class SpawnProgressMessage:
     """A progress message from JupyterLab spawning."""
@@ -254,7 +245,8 @@ class JupyterConfig:
     image_reference: Optional[str] = None
     """Docker reference to the JupyterLab image to spawn.
 
-    May be null if ``image_selector`` is `JupyterImageSelector.BY_REFERENCE`.
+    May be null if ``image_selector`` is is not
+    `JupyterImageSelector.reference`.
     """
 
     image_size: str = "Large"
@@ -545,11 +537,11 @@ async def spawn_progress(self) -> AsyncIterator[SpawnProgressMessage]:
 
     async def _get_spawn_image(self) -> JupyterImage:
         """Determine what image to spawn."""
-        if self.config.image_selector == JupyterImageSelector.RECOMMENDED:
+        if self.config.image_selector == JupyterImageSelector.recommended:
             image = await self.cachemachine.get_recommended()
-        elif self.config.image_selector == JupyterImageSelector.LATEST_WEEKLY:
+        elif self.config.image_selector == JupyterImageSelector.weekly:
             image = await self.cachemachine.get_latest_weekly()
-        elif self.config.image_selector == JupyterImageSelector.BY_REFERENCE:
+        elif self.config.image_selector == JupyterImageSelector.reference:
             assert self.config.image_reference
             image = JupyterImage.from_reference(self.config.image_reference)
         else:

src/noteburst/worker/main.py

@@ -2,19 +2,18 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict
+from typing import Any, Dict, List
 
 import httpx
 import structlog
 from arq import cron
 from safir.logging import configure_logging
 
-from noteburst.config import WorkerConfig
+from noteburst.config import WorkerConfig, WorkerKeepAliveSetting
 from noteburst.jupyterclient.jupyterlab import (
     JupyterClient,
     JupyterConfig,
     JupyterError,
-    JupyterImageSelector,
 )
 from noteburst.jupyterclient.user import User
 
@@ -50,7 +49,8 @@ async def startup(ctx: Dict[Any, Any]) -> None:
     ctx["http_client"] = http_client
 
     jupyter_config = JupyterConfig(
-        image_selector=JupyterImageSelector.RECOMMENDED
+        image_selector=config.image_selector,
+        image_reference=config.image_reference,
     )
 
     identity = await identity_manager.get_identity()
@@ -60,7 +60,7 @@ async def startup(ctx: Dict[Any, Any]) -> None:
 
         user = User(username=identity.username, uid=identity.uid)
         authed_user = await user.login(
-            scopes=["exec:notebook"],
+            scopes=config.parsed_worker_token_scopes,
             http_client=http_client,
             token_lifetime=config.worker_token_lifetime,
         )
@@ -117,7 +117,17 @@ async def shutdown(ctx: Dict[Any, Any]) -> None:
 
 # For info on ignoring the type checking here, see
 # https://github.com/samuelcolvin/arq/issues/249
-cron_jobs = [cron(keep_alive, second={0, 30}, unique=False)]  # type: ignore
+cron_jobs: List[cron] = []  # type: ignore
+if config.worker_keepalive == WorkerKeepAliveSetting.fast:
+    f = cron(keep_alive, second={0, 30}, unique=False)  # type: ignore
+    cron_jobs.append(f)
+elif config.worker_keepalive == WorkerKeepAliveSetting.normal:
+    f = cron(
+        keep_alive,  # type: ignore
+        minute={0, 5, 10, 15, 20, 25, 30, 35, 40, 45, 50, 55},
+        unique=False,
+    )
+    cron_jobs.append(f)
 
 
 class WorkerSettings:

tests/jupyterclient/jupyterclient_test.py

@@ -9,11 +9,8 @@
 import respx
 import structlog
 
-from noteburst.jupyterclient.jupyterlab import (
-    JupyterClient,
-    JupyterConfig,
-    JupyterImageSelector,
-)
+from noteburst.config import JupyterImageSelector
+from noteburst.jupyterclient.jupyterlab import JupyterClient, JupyterConfig
 from noteburst.jupyterclient.user import User
 from tests.support.gafaelfawr import mock_gafaelfawr
 
@@ -36,7 +33,7 @@ async def test_jupyterclient(
     logger = structlog.get_logger(__name__)
 
     jupyter_config = JupyterConfig(
-        image_selector=JupyterImageSelector.RECOMMENDED
+        image_selector=JupyterImageSelector.recommended
     )
 
     async with httpx.AsyncClient() as http_client: