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 @@
 
 """
 
+importwarnings
 fromtypingimportSequence
 
 importpkg_resources
+importpydantic
+
+ifpydantic.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 .clientimportLaunchClient
 from .connectionimportConnection

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):
 classGetFineTuneResponse(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 @@
 fromenumimportEnum
-fromtypingimportAny, Dict, Set, Type, Union
+fromtypingimportAny, Callable, Dict, Set, Type, Union
 
+importpydantic
 frompydanticimportBaseModel
 
-try:
+ifhasattr(pydantic, "VERSION") andpydantic.VERSION.startswith("1."):
+PYDANTIC_VERSION=1
 frompydantic.schemaimport (
 get_flat_models_from_models,
 model_process_schema,
  )
-exceptImportError:
-# We assume this is due to the user having pydantic 2.x installed.
-frompydantic.v1.schemaimport ( # type: ignore
-get_flat_models_from_models,
-model_process_schema,
- )
+elifhasattr(pydantic, "VERSION") andpydantic.VERSION.startswith("2."):
+PYDANTIC_VERSION=2
+else:
+raiseImportError("Unsupported pydantic version.")
 
 
 REF_PREFIX="#/components/schemas/"
 
 
-defget_model_definitions(request_schema: Type[BaseModel], response_schema: Type[BaseModel]) ->Dict[str, Any]:
+defget_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
 returnget_model_definitions_from_flat_models(flat_models=flat_models, model_name_map=model_name_map)
 
 
+defget_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()}
+
+
+ifPYDANTIC_VERSION==1:
+get_model_definitions: Callable=get_model_definitions_v1
+elifPYDANTIC_VERSION==2:
+get_model_definitions: Callable=get_model_definitions_v2
+
+
 defget_model_definitions_from_flat_models(
 *,
 flat_models: Set[Union[Type[BaseModel], Type[Enum]]],