Migrate from pydantic v1 to v2 API

jupyter_scheduler/backend_registry.py

@@ -1,10 +1,11 @@
 import logging
 from typing import Any, Dict, List, Optional, Type
 
+from pydantic import BaseModel
+
 from jupyter_scheduler.backends import BackendConfig, DescribeBackendResponse
 from jupyter_scheduler.environments import EnvironmentManager
 from jupyter_scheduler.orm import create_tables
-from jupyter_scheduler.pydantic_v1 import BaseModel
 
 logger = logging.getLogger(__name__)
 

jupyter_scheduler/backends.py

@@ -1,8 +1,9 @@
 from typing import Any, Dict, List, Optional
 
+from pydantic import BaseModel, ConfigDict, Field
+
 from jupyter_scheduler.base_backend import BaseBackend
 from jupyter_scheduler.models import OutputFormat
-from jupyter_scheduler.pydantic_v1 import BaseModel, Field
 
 JUPYTER_SERVER_NB_BACKEND_ID = "jupyter_server_nb"
 JUPYTER_SERVER_PY_BACKEND_ID = "jupyter_server_py"
@@ -37,8 +38,7 @@ class DescribeBackendResponse(BaseModel):
     file_extensions: List[str]
     output_formats: List[OutputFormat]
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 class JupyterServerNotebookBackend(BaseBackend):

jupyter_scheduler/executors.py

@@ -40,7 +40,7 @@ def model(self):
         if self._model is None:
             with self.db_session() as session:
                 job = session.query(Job).filter(Job.job_id == self.job_id).first()
-                self._model = DescribeJob.from_orm(job)
+                self._model = DescribeJob.model_validate(job)
         return self._model
 
     @property

jupyter_scheduler/handlers.py

@@ -6,6 +6,7 @@
 from jupyter_server.base.handlers import APIHandler
 from jupyter_server.extension.handler import ExtensionHandlerMixin
 from jupyter_server.utils import ensure_async
+from pydantic import ValidationError
 from tornado.web import HTTPError, authenticated
 
 from jupyter_scheduler.backend_registry import BackendInstance, BackendRegistry
@@ -31,7 +32,6 @@
     UpdateJob,
     UpdateJobDefinition,
 )
-from jupyter_scheduler.pydantic_v1 import ValidationError
 from jupyter_scheduler.scheduler import BaseScheduler
 
 logger = logging.getLogger(__name__)
@@ -124,7 +124,7 @@ async def get(self, job_definition_id=None):
                     500, f"Unexpected error while getting job definition details: {e}"
                 ) from e
             else:
-                self.finish(job_definition.json())
+                self.finish(job_definition.model_dump_json())
         else:
             create_time = self.get_query_argument("create_time", None)
             sort_by = compute_sort_model(self.get_query_arguments("sort_by"))
@@ -153,7 +153,7 @@ async def get(self, job_definition_id=None):
                     500, f"Unexpected error while getting job definition list: {e}"
                 ) from e
             else:
-                self.finish(list_response.json(exclude_none=True))
+                self.finish(list_response.model_dump_json(exclude_none=True))
 
     @authenticated
     async def post(self):
@@ -238,7 +238,7 @@ async def get(self, job_id=None):
                 self.log.exception(e)
                 raise HTTPError(500, f"Unexpected error while getting job details: {e}") from e
             else:
-                self.finish(job.json())
+                self.finish(job.model_dump_json())
         else:
             status = self.get_query_argument("status", None)
             start_time = self.get_query_argument("start_time", None)
@@ -295,7 +295,7 @@ async def get(self, job_id=None):
                 self.log.exception(e)
                 raise HTTPError(500, f"Unexpected error while getting jobs list: {e}") from e
             else:
-                self.finish(list_jobs_response.json(exclude_none=True))
+                self.finish(list_jobs_response.model_dump_json(exclude_none=True))
 
     @authenticated
     async def post(self):
@@ -456,7 +456,7 @@ async def get(self):
 
         response = []
         for environment in environments:
-            env = environment.dict()
+            env = environment.model_dump()
             formats = env["output_formats"]
             env["output_formats"] = [{"id": f, "label": output_formats[f]} for f in formats]
             response.append(env)
@@ -522,7 +522,7 @@ async def get(self):
                 raise HTTPError(500, "Backend registry not initialized")
 
             backends = registry.describe_backends()
-            self.finish(json.dumps([b.dict() for b in backends]))
+            self.finish(json.dumps([b.model_dump() for b in backends]))
         except SchedulerError as e:
             self.log.exception(e)
             raise HTTPError(500, f"Unexpected error while listing backends: {e}") from e

jupyter_scheduler/models.py

@@ -1,8 +1,15 @@
 import os
 from enum import Enum
-from typing import Dict, List, Optional, Union
+from typing import Annotated, Dict, List, Optional, Union
 
-from jupyter_scheduler.pydantic_v1 import BaseModel, root_validator
+from pydantic import BaseModel, BeforeValidator, ConfigDict, model_validator
+
+
+def _coerce_str(v):
+    return str(v) if v is not None else v
+
+
+CoercedStr = Annotated[str, BeforeValidator(_coerce_str)]
 
 Tags = List[str]
 EnvironmentParameterValues = Union[int, float, bool, str]
@@ -33,23 +40,23 @@ class RuntimeEnvironment(BaseModel):
     description: str
     file_extensions: List[str]  # Supported input file types
     output_formats: List[str]  # Supported output formats
-    metadata: Optional[Dict[str, str]]  # Optional metadata
-    compute_types: Optional[List[str]]
-    default_compute_type: Optional[str]  # Should be a member of the compute_types list
-    utc_only: Optional[bool]
+    metadata: Optional[Dict[str, str]] = None  # Optional metadata
+    compute_types: Optional[List[str]] = None
+    default_compute_type: Optional[str] = None  # Should be a member of the compute_types list
+    utc_only: Optional[bool] = None
 
     def __str__(self):
-        return self.json()
+        return self.model_dump_json()
 
 
 class EmailNotifications(BaseModel):
-    on_start: Optional[List[str]]
-    on_success: Optional[List[str]]
-    on_failure: Optional[List[str]]
+    on_start: Optional[List[str]] = None
+    on_success: Optional[List[str]] = None
+    on_failure: Optional[List[str]] = None
     no_alert_for_skipped_runs: bool = True
 
     def __str__(self) -> str:
-        return self.json()
+        return self.model_dump_json()
 
 
 class Status(str, Enum):
@@ -85,23 +92,24 @@ class CreateJob(BaseModel):
     """Defines the model for creating a new job"""
 
     input_uri: str
-    input_filename: str = None
+    input_filename: Optional[str] = None
     runtime_environment_name: str
-    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
+    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]] = None
     output_formats: Optional[List[str]] = None
     idempotency_token: Optional[str] = None
     job_definition_id: Optional[str] = None
-    parameters: Optional[Dict[str, str]] = None
+    parameters: Optional[Dict[str, Union[str, int, float, bool]]] = None
     tags: Optional[Tags] = None
     name: str
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
     compute_type: Optional[str] = None
     package_input_folder: Optional[bool] = None
     backend_id: Optional[str] = None
 
-    @root_validator
+    @model_validator(mode="before")
+    @classmethod
     def compute_input_filename(cls, values) -> Dict:
-        if not values["input_filename"] and values["input_uri"]:
+        if not values.get("input_filename") and values.get("input_uri"):
             values["input_filename"] = os.path.basename(values["input_uri"])
 
         return values
@@ -137,13 +145,13 @@ class JobFile(BaseModel):
 
 
 class DescribeJob(BaseModel):
-    input_filename: str = None
+    input_filename: Optional[str] = None
     runtime_environment_name: str
-    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
+    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]] = None
     output_formats: Optional[List[str]] = None
     idempotency_token: Optional[str] = None
     job_definition_id: Optional[str] = None
-    parameters: Optional[Dict[str, str]] = None
+    parameters: Optional[Dict[str, Union[str, int, float, bool]]] = None
     tags: Optional[Tags] = None
     name: str
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -162,8 +170,7 @@ class DescribeJob(BaseModel):
     packaged_files: Optional[List[str]] = []
     backend_id: Optional[str] = None
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 class SortDirection(Enum):
@@ -216,11 +223,11 @@ class DeleteJob(BaseModel):
 
 class CreateJobDefinition(BaseModel):
     input_uri: str
-    input_filename: str = None
+    input_filename: Optional[str] = None
     runtime_environment_name: str
-    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
+    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]] = None
     output_formats: Optional[List[str]] = None
-    parameters: Optional[Dict[str, str]] = None
+    parameters: Optional[Dict[str, Union[str, int, float, bool]]] = None
     tags: Optional[Tags] = None
     name: str
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -230,20 +237,21 @@ class CreateJobDefinition(BaseModel):
     package_input_folder: Optional[bool] = None
     backend_id: Optional[str] = None
 
-    @root_validator
+    @model_validator(mode="before")
+    @classmethod
     def compute_input_filename(cls, values) -> Dict:
-        if not values["input_filename"] and "input_uri" in values and values["input_uri"]:
+        if not values.get("input_filename") and "input_uri" in values and values.get("input_uri"):
             values["input_filename"] = os.path.basename(values["input_uri"])
 
         return values
 
 
 class DescribeJobDefinition(BaseModel):
-    input_filename: str = None
+    input_filename: Optional[str] = None
     runtime_environment_name: str
-    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
+    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]] = None
     output_formats: Optional[List[str]] = None
-    parameters: Optional[Dict[str, str]] = None
+    parameters: Optional[Dict[str, Union[str, int, float, bool]]] = None
     tags: Optional[Tags] = None
     name: str
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
@@ -258,15 +266,14 @@ class DescribeJobDefinition(BaseModel):
     packaged_files: Optional[List[str]] = []
     backend_id: Optional[str] = None
 
-    class Config:
-        orm_mode = True
+    model_config = ConfigDict(from_attributes=True)
 
 
 class UpdateJobDefinition(BaseModel):
-    runtime_environment_name: Optional[str]
-    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]]
+    runtime_environment_name: Optional[str] = None
+    runtime_environment_parameters: Optional[Dict[str, EnvironmentParameterValues]] = None
     output_formats: Optional[List[str]] = None
-    parameters: Optional[Dict[str, str]] = None
+    parameters: Optional[Dict[str, Union[str, int, float, bool]]] = None
     tags: Optional[Tags] = None
     name: Optional[str] = None
     url: Optional[str] = None
@@ -284,17 +291,17 @@ class ListJobDefinitionsQuery(BaseModel):
     tags: Optional[Tags] = None
     sort_by: List[SortField] = [DEFAULT_SORT]
     max_items: Optional[int] = DEFAULT_MAX_ITEMS
-    next_token: Optional[str] = None
+    next_token: Optional[CoercedStr] = None
 
 
 class ListJobDefinitionsResponse(BaseModel):
     job_definitions: List[DescribeJobDefinition] = []
     total_count: int = 0
-    next_token: Optional[str] = None
+    next_token: Optional[CoercedStr] = None
 
 
 class CreateJobFromDefinition(BaseModel):
-    parameters: Optional[Dict[str, str]] = None
+    parameters: Optional[Dict[str, Union[str, int, float, bool]]] = None
 
 
 class JobFeature(str, Enum):

jupyter_scheduler/orm.py

@@ -54,7 +54,7 @@ def process_bind_param(self, value, dialect):
             return None
 
         if isinstance(value, EmailNotifications):
-            return json.dumps(value.dict(exclude_none=True))
+            return json.dumps(value.model_dump(exclude_none=True))
         else:
             return value