diff --git a/src/noteburst/handlers/external.py b/src/noteburst/handlers/external.py index a5115da..a8ed096 100644 --- a/src/noteburst/handlers/external.py +++ b/src/noteburst/handlers/external.py @@ -20,7 +20,7 @@ class Index(BaseModel): """Metadata about the application.""" - metadata: SafirMetadata = Field(..., title="Package metadata") + metadata: Annotated[SafirMetadata, Field(title="Package metadata")] @external_router.get( diff --git a/src/noteburst/handlers/v1/models.py b/src/noteburst/handlers/v1/models.py index 83f664b..fd3b211 100644 --- a/src/noteburst/handlers/v1/models.py +++ b/src/noteburst/handlers/v1/models.py @@ -4,7 +4,7 @@ import json from datetime import datetime -from typing import Any +from typing import Annotated, Any from arq.jobs import JobStatus from fastapi import Request @@ -31,8 +31,8 @@ class NotebookError(BaseModel): """Information about an exception that occurred during notebook exec.""" - name: str = Field(description="The name of the exception.") - message: str = Field(description="The exception's message.") + name: Annotated[str, Field(description="The name of the exception.")] + message: Annotated[str, Field(description="The exception's message.")] @classmethod def from_nbexec_error( @@ -52,56 +52,72 @@ class NotebookResponse(BaseModel): result and source notebooks. """ - job_id: str = Field(title="The job ID") - - kernel_name: str = kernel_name_field - - enqueue_time: datetime = Field( - title="Time when the job was added to the queue (UTC)" - ) - - status: JobStatus = Field( - title="The current status of the notebook execution job" - ) - - self_url: AnyHttpUrl = Field(title="The URL of this resource") - - source: str | None = Field( - None, - title="The content of the source ipynb file (JSON-encoded string)", - description="This field is null unless the source is requested.", - ) - - start_time: datetime | None = Field( - None, - title="Time when the notebook execution started (UTC)", - description="This field is present if the result is available.", - ) - - finish_time: datetime | None = Field( - None, - title="Time when the notebook execution completed (UTC)", - description="This field is present only if the result is available.", - ) - - success: bool | None = Field( - None, - title="Whether the execution was successful or not", - description="This field is present if the result is available.", - ) - - ipynb: str | None = Field( - None, - title="The contents of the executed Jupyter notebook", - description="The ipynb is a JSON-encoded string. This field is " - "present if the result is available.", - ) - - ipynb_error: NotebookError | None = Field( - None, - title="The error that occurred during notebook execution", - description="This field is null if an exeception did not occur.", - ) + job_id: Annotated[str, Field(title="The job ID")] + + kernel_name: Annotated[str, kernel_name_field] + + enqueue_time: Annotated[ + datetime, Field(title="Time when the job was added to the queue (UTC)") + ] + + status: Annotated[ + JobStatus, + Field(title="The current status of the notebook execution job"), + ] + + self_url: Annotated[AnyHttpUrl, Field(title="The URL of this resource")] + + source: Annotated[ + str | None, + Field( + title="The content of the source ipynb file (JSON-encoded string)", + description="This field is null unless the source is requested.", + ), + ] = None + + start_time: Annotated[ + datetime | None, + Field( + title="Time when the notebook execution started (UTC)", + description="This field is present if the result is available.", + ), + ] = None + + finish_time: Annotated[ + datetime | None, + Field( + title="Time when the notebook execution completed (UTC)", + description=( + "This field is present only if the result is available." + ), + ), + ] = None + + success: Annotated[ + bool | None, + Field( + title="Whether the execution was successful or not", + description="This field is present if the result is available.", + ), + ] = None + + ipynb: Annotated[ + str | None, + Field( + title="The contents of the executed Jupyter notebook", + description="The ipynb is a JSON-encoded string. This field is " + "present if the result is available.", + ), + ] = None + + ipynb_error: Annotated[ + NotebookError | None, + Field( + None, + title="The error that occurred during notebook execution", + description="This field is null if an exeception did not occur.", + ), + ] = None @classmethod async def from_job_metadata( @@ -112,6 +128,7 @@ async def from_job_metadata( include_source: bool = False, job_result: JobResult | None = None, ) -> NotebookResponse: + """Create a NotebookResponse from a job.""" if job_result is not None and job_result.success: nbexec_result = NotebookExecutionResult.model_validate_json( job_result.result @@ -143,30 +160,35 @@ async def from_job_metadata( class PostNotebookRequest(BaseModel): """The ``POST /notebooks/`` request body.""" - ipynb: str | dict[str, Any] = Field( - ..., - title="The contents of a Jupyter notebook", - description="If a string, the content is parsed as JSON. " - "Alternatively, the content can be submitted pre-parsed as " - "an object.", - ) - - kernel_name: str = kernel_name_field - - enable_retry: bool = Field( - True, - title="Enable retries on failures", - description=( - "If true (default), noteburst will retry notebook " - "execution if the notebook fails, with an increasing back-off " - "time between tries. This is useful for dealing with transient " - "issues. However, if you are using Noteburst for continuous " - "integration of notebooks, disabling retries provides faster " - "feedback." + ipynb: Annotated[ + str | dict[str, Any], + Field( + title="The contents of a Jupyter notebook", + description="If a string, the content is parsed as JSON. " + "Alternatively, the content can be submitted pre-parsed as " + "an object.", + ), + ] + + kernel_name: Annotated[str, kernel_name_field] + + enable_retry: Annotated[ + bool, + Field( + title="Enable retries on failures", + description=( + "If true (default), noteburst will retry notebook " + "execution if the notebook fails, with an increasing back-off " + "time between tries. This is useful for dealing with " + "transient issues. However, if you are using Noteburst for " + "continuous integration of notebooks, disabling retries " + "provides faster feedback." + ), ), - ) + ] = True def get_ipynb_as_str(self) -> str: + """Get the ipynb as a JSON-encoded string.""" if isinstance(self.ipynb, str): return self.ipynb else: diff --git a/src/noteburst/jupyterclient/jupyterlab.py b/src/noteburst/jupyterclient/jupyterlab.py index 37a782a..fa9ea83 100644 --- a/src/noteburst/jupyterclient/jupyterlab.py +++ b/src/noteburst/jupyterclient/jupyterlab.py @@ -10,7 +10,7 @@ from collections.abc import AsyncGenerator, AsyncIterator from dataclasses import dataclass from random import SystemRandom -from typing import Any, Self +from typing import Annotated, Any, Self from urllib.parse import urljoin, urlparse from uuid import uuid4 @@ -356,33 +356,38 @@ def __str__(self) -> str: class NotebookExecutionErrorModel(BaseModel): - """The error from the /user/:username/rubin/execute endpoint.""" + """The error from the ``/user/:username/rubin/execute`` endpoint.""" - traceback: str = Field(description="The exeception traceback.") + traceback: Annotated[str, Field(description="The exeception traceback.")] - ename: str = Field(description="The exception name.") + ename: Annotated[str, Field(description="The exception name.")] - evalue: str = Field(description="The exception value.") + evalue: Annotated[str, Field(description="The exception value.")] - err_msg: str = Field(description="The exception message.") + err_msg: Annotated[str, Field(description="The exception message.")] class NotebookExecutionResult(BaseModel): """The result of the /user/:username/rubin/execute endpoint.""" - notebook: str = Field( - description="The notebook that was executed, as a JSON string." - ) + notebook: Annotated[ + str, + Field(description="The notebook that was executed, as a JSON string."), + ] - resources: dict[str, Any] = Field( - description=( - "The resources used to execute the notebook, as a JSON string." - ) - ) + resources: Annotated[ + dict[str, Any], + Field( + description=( + "The resources used to execute the notebook, as a JSON string." + ) + ), + ] - error: NotebookExecutionErrorModel | None = Field( - None, description="The error that occurred during execution." - ) + error: Annotated[ + NotebookExecutionErrorModel | None, + Field(description="The error that occurred during execution."), + ] = None class JupyterClient: diff --git a/src/noteburst/jupyterclient/labcontroller.py b/src/noteburst/jupyterclient/labcontroller.py index c074c5c..faa5984 100644 --- a/src/noteburst/jupyterclient/labcontroller.py +++ b/src/noteburst/jupyterclient/labcontroller.py @@ -2,10 +2,11 @@ from __future__ import annotations +from typing import Annotated from urllib.parse import urljoin import httpx -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field from noteburst.config import config @@ -13,42 +14,51 @@ class JupyterImage(BaseModel): """A model for a JupyterLab image in a `LabControllerImages` resource.""" - reference: str = Field( - ..., - examples=["lighthouse.ceres/library/sketchbook:latest_daily"], - title="Full Docker registry path for lab image", - description="cf. https://docs.docker.com/registry/introduction/", - ) - - name: str = Field( - ..., - examples=["Latest Daily (Daily 2077_10_23)"], - title="Human-readable version of image tag", - ) - - digest: str | None = Field( - None, - examples=[ - "sha256:e693782192ecef4f7846ad2b21" - "b1574682e700747f94c5a256b5731331a2eec2" - ], - title="unique digest of image contents", - ) - - tag: str = Field( - title="Image tag", - ) - - size: int | None = Field( - None, - examples=[8675309], - title="Size in bytes of image. None if image size is unknown", - ) - prepulled: bool = Field( - False, - examples=[False], - title="Whether image is prepulled to all eligible nodes", - ) + reference: Annotated[ + str, + Field( + examples=["lighthouse.ceres/library/sketchbook:latest_daily"], + title="Full Docker registry path for lab image", + description="cf. https://docs.docker.com/registry/introduction/", + ), + ] + + name: Annotated[ + str, + Field( + examples=["Latest Daily (Daily 2077_10_23)"], + title="Human-readable version of image tag", + ), + ] + + digest: Annotated[ + str | None, + Field( + examples=[ + "sha256:e693782192ecef4f7846ad2b21" + "b1574682e700747f94c5a256b5731331a2eec2" + ], + title="unique digest of image contents", + ), + ] = None + + tag: Annotated[str, Field(title="Image tag")] + + size: Annotated[ + int | None, + Field( + examples=[8675309], + title="Size in bytes of image. None if image size is unknown", + ), + ] = None + + prepulled: Annotated[ + bool, + Field( + examples=[False], + title="Whether image is prepulled to all eligible nodes", + ), + ] = False def underscore_to_dash(x: str) -> str: @@ -59,23 +69,31 @@ def underscore_to_dash(x: str) -> str: class LabControllerImages(BaseModel): """A model for the ``GET /nublado/spawner/v1/images`` response.""" - recommended: JupyterImage | None = Field( - None, title="The recommended image" - ) + recommended: Annotated[ + JupyterImage | None, Field(title="The recommended image") + ] = None - latest_weekly: JupyterImage | None = Field( - None, title="The latest weekly release image" - ) + latest_weekly: Annotated[ + JupyterImage | None, Field(title="The latest weekly release image") + ] = None - latest_daily: JupyterImage | None = Field( - None, title="The latest daily release image" - ) + latest_daily: Annotated[ + JupyterImage | None, Field(title="The latest daily release image") + ] = None - latest_release: JupyterImage | None = Field( - None, title="The latest release image" - ) + latest_release: Annotated[ + JupyterImage | None, Field(title="The latest release image") + ] = None + + all: Annotated[ + list[JupyterImage], Field(default_factory=list, title="All images") + ] - all: list[JupyterImage] = Field(default_factory=list, title="All images") + model_config = ConfigDict( + populate_by_name=True, + alias_generator=underscore_to_dash, + ) + """Pydantic model configuration.""" def get_by_reference(self, reference: str) -> JupyterImage | None: """Get the JupyterImage with a corresponding reference. @@ -96,10 +114,6 @@ def get_by_reference(self, reference: str) -> JupyterImage | None: return None - class Config: - allow_population_by_field_name = True - alias_generator = underscore_to_dash - class LabControllerError(Exception): """Unable to get image information from the JupyterLab Controller.""" diff --git a/src/noteburst/worker/identity.py b/src/noteburst/worker/identity.py index 770591b..ea4227a 100644 --- a/src/noteburst/worker/identity.py +++ b/src/noteburst/worker/identity.py @@ -8,11 +8,12 @@ from dataclasses import dataclass from pathlib import Path +from typing import Annotated import structlog import yaml from aioredlock import Aioredlock, Lock, LockError -from pydantic import BaseModel, RootModel +from pydantic import BaseModel, Field, RootModel from noteburst.config import WorkerConfig @@ -22,14 +23,19 @@ class IdentityModel(BaseModel): configuration file. """ - username: str - """The username of the user account.""" - - uid: str | None = None - """The UID of the user account. + username: Annotated[ + str, Field(description="The username of the user account.") + ] - This can be `None` if the authentication system assigns the UID. - """ + uid: Annotated[ + str | None, + Field( + description=( + "The UID of the user account. This can be `None` if the " + "authentication system assigns the UID." + ) + ), + ] = None class IdentityConfigModel(RootModel):