better forward compatibility with Pydantic v2 (#162)

* better forward compatibility with pydantic v2 * format change * fix format

Author: Tma2333-Scale

launch/__init__.py

@@ -6,9 +6,16 @@
 
 """
 
+import warnings
 from typing import Sequence
 
 import pkg_resources
+import pydantic
+
+if pydantic.VERSION.startswith("2."):
+ # HACK: Suppress warning from pydantic v2 about protected namespace, this is due to
+ # launch-python-client module is based on v1 and only does minimum to support forward compatibility
+ warnings.filterwarnings("ignore", category=UserWarning, module="pydantic")
 
 from .client import LaunchClient
 from .connection import Connection

launch/docker_image_batch_job_bundle.py

@@ -34,18 +34,18 @@ class DockerImageBatchJobBundleResponse(BaseModel):
  """The command to run inside the docker image"""
  env: Dict[str, str]
  """Environment variables to be injected into the docker image"""
- mount_location: Optional[str]
+ mount_location: Optional[str] = None
  """Location of a json-formatted file to mount inside the docker image.
  Contents get populated at runtime, and this is the method to change behavior on runtime."""
- cpus: Optional[str]
+ cpus: Optional[str] = None
  """Default number of cpus to give to the docker image"""
- memory: Optional[str]
+ memory: Optional[str] = None
  """Default amount of memory to give to the docker image"""
- storage: Optional[str]
+ storage: Optional[str] = None
  """Default amount of disk to give to the docker image"""
- gpus: Optional[int]
+ gpus: Optional[int] = None
  """Default number of gpus to give to the docker image"""
- gpu_type: Optional[str]
+ gpu_type: Optional[str] = None
  """Default type of gpu, e.g. nvidia-tesla-t4, nvidia-ampere-a10 to give to the docker image"""
 
 

launch/fine_tune.py

@@ -28,7 +28,7 @@ class CreateFineTuneResponse(BaseModel):
 class GetFineTuneResponse(BaseModel):
  id: str
  """ID of the requested job"""
- fine_tuned_model: Optional[str]
+ fine_tuned_model: Optional[str] = None
  """
  Name of the resulting fine-tuned model. This can be plugged into the
  Completion API ones the fine-tune is complete

launch/model_bundle.py

@@ -94,7 +94,7 @@ class ZipArtifactFlavor(BaseModel):
  [`CustomFramework`](./#launch.model_bundle.CustomFramework).
  """
 
- app_config: Optional[Dict[str, Any]]
+ app_config: Optional[Dict[str, Any]] = None
  """Optional configuration for the application."""
 
  location: str
@@ -112,7 +112,7 @@ class RunnableImageLike(BaseModel, ABC):
  repository: str
  tag: str
  command: List[str]
- env: Optional[Dict[str, str]]
+ env: Optional[Dict[str, str]] = None
  protocol: Literal["http"] # TODO: add support for other protocols (e.g. grpc)
  readiness_initial_delay_seconds: int = 120
 
@@ -137,15 +137,15 @@ class TritonEnhancedRunnableImageFlavor(RunnableImageLike):
 
  triton_model_repository: str
 
- triton_model_replicas: Optional[Dict[str, str]]
+ triton_model_replicas: Optional[Dict[str, str]] = None
 
  triton_num_cpu: float
 
  triton_commit_tag: str
 
- triton_storage: Optional[str]
+ triton_storage: Optional[str] = None
 
- triton_memory: Optional[str]
+ triton_memory: Optional[str] = None
 
  triton_readiness_initial_delay_seconds: int = 300
 
@@ -197,7 +197,7 @@ class ModelBundleV2Response(BaseModel):
  model_artifact_ids: List[str]
  """IDs of the Model Artifacts associated with the Model Bundle."""
 
- schema_location: Optional[str]
+ schema_location: Optional[str] = None
 
  flavor: ModelBundleFlavors = Field(..., discriminator="flavor")
  """Flavor of the Model Bundle, representing how the model bundle was packaged.

launch/pydantic_schemas.py

@@ -1,25 +1,25 @@
 from enum import Enum
-from typing import Any, Dict, Set, Type, Union
+from typing import Any, Callable, Dict, Set, Type, Union
 
+import pydantic
 from pydantic import BaseModel
 
-try:
+if hasattr(pydantic, "VERSION") and pydantic.VERSION.startswith("1."):
+ PYDANTIC_VERSION = 1
  from pydantic.schema import (
  get_flat_models_from_models,
  model_process_schema,
  )
-except ImportError:
- # We assume this is due to the user having pydantic 2.x installed.
- from pydantic.v1.schema import ( # type: ignore
- get_flat_models_from_models,
- model_process_schema,
- )
+elif hasattr(pydantic, "VERSION") and pydantic.VERSION.startswith("2."):
+ PYDANTIC_VERSION = 2
+else:
+ raise ImportError("Unsupported pydantic version.")
 
 
 REF_PREFIX = "#/components/schemas/"
 
 
-def get_model_definitions(request_schema: Type[BaseModel], response_schema: Type[BaseModel]) -> Dict[str, Any]:
+def get_model_definitions_v1(request_schema: Type[BaseModel], response_schema: Type[BaseModel]) -> Dict[str, Any]:
  """
  Gets the model schemas in jsonschema format from a sequence of Pydantic BaseModels.
  """
@@ -29,6 +29,16 @@ def get_model_definitions(request_schema: Type[BaseModel], response_schema: Type
  return get_model_definitions_from_flat_models(flat_models=flat_models, model_name_map=model_name_map)
 
 
+def get_model_definitions_v2(request_schema: Type[BaseModel], response_schema: Type[BaseModel]) -> Dict[str, Any]:
+ return {"RequestSchema": request_schema.model_json_schema(), "ResponseSchema": response_schema.model_json_schema()}
+
+
+if PYDANTIC_VERSION == 1:
+ get_model_definitions: Callable = get_model_definitions_v1
+elif PYDANTIC_VERSION == 2:
+ get_model_definitions: Callable = get_model_definitions_v2
+
+
 def get_model_definitions_from_flat_models(
  *,
  flat_models: Set[Union[Type[BaseModel], Type[Enum]]],