fix: Fixed openapi dataplane.yaml (#1752)

* Fixed openapi dataplane.yaml * Geneated dataplane.py with pydantic v2 * Fixed TensorData root typing
SeldonIO · May 14, 2024 · 1eb48eb · 1eb48eb
1 parent a2c0d64
commit 1eb48eb
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 54 deletions.
diff --git a/mlserver/types/dataplane.py b/mlserver/types/dataplane.py
@@ -1,10 +1,12 @@
 # generated by datamodel-codegen:
 #   filename:  dataplane.yaml
 
+from __future__ import annotations
+
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
-from pydantic import Field, RootModel, ConfigDict
+from pydantic import ConfigDict, Field, RootModel
 
 from .base import BaseModel
 
@@ -19,22 +21,6 @@ class MetadataServerErrorResponse(BaseModel):
     error: str
 
 
-class Datatype(Enum):
-    BOOL = "BOOL"
-    UINT8 = "UINT8"
-    UINT16 = "UINT16"
-    UINT32 = "UINT32"
-    UINT64 = "UINT64"
-    INT8 = "INT8"
-    INT16 = "INT16"
-    INT32 = "INT32"
-    INT64 = "INT64"
-    FP16 = "FP16"
-    FP32 = "FP32"
-    FP64 = "FP64"
-    BYTES = "BYTES"
-
-
 class MetadataModelErrorResponse(BaseModel):
     error: str
 
@@ -43,13 +29,12 @@ class Parameters(BaseModel):
     model_config = ConfigDict(
         extra="allow",
     )
-
     content_type: Optional[str] = None
     headers: Optional[Dict[str, Any]] = None
 
 
-class TensorData(RootModel[Any]):
-    root: Any = Field(..., title="TensorData")
+class TensorData(RootModel[Union[List, str]]):
+    root: Union[List, str] = Field(..., title="TensorData")
 
     def __iter__(self):
         return iter(self.root)
@@ -63,55 +48,71 @@ def __len__(self):
 
 class RequestOutput(BaseModel):
     name: str
-    parameters: Optional["Parameters"] = None
-
-
-class ResponseOutput(BaseModel):
-    name: str
-    shape: List[int]
-    datatype: "Datatype"
-    parameters: Optional["Parameters"] = None
-    data: "TensorData"
-
-
-class InferenceResponse(BaseModel):
-    model_name: str
-    model_version: Optional[str] = None
-    id: Optional[str] = None
-    parameters: Optional["Parameters"] = None
-    outputs: List["ResponseOutput"]
+    parameters: Optional[Parameters] = None
 
 
 class InferenceErrorResponse(BaseModel):
     error: Optional[str] = None
 
 
+class Datatype(Enum):
+    BOOL = "BOOL"
+    UINT8 = "UINT8"
+    UINT16 = "UINT16"
+    UINT32 = "UINT32"
+    UINT64 = "UINT64"
+    INT8 = "INT8"
+    INT16 = "INT16"
+    INT32 = "INT32"
+    INT64 = "INT64"
+    FP16 = "FP16"
+    FP32 = "FP32"
+    FP64 = "FP64"
+    BYTES = "BYTES"
+
+
 class MetadataTensor(BaseModel):
     name: str
-    datatype: "Datatype"
+    datatype: Datatype
     shape: List[int]
-    parameters: Optional["Parameters"] = None
+    parameters: Optional[Parameters] = None
 
 
 class RequestInput(BaseModel):
     name: str
     shape: List[int]
-    datatype: "Datatype"
-    parameters: Optional["Parameters"] = None
-    data: "TensorData"
+    datatype: Datatype
+    parameters: Optional[Parameters] = None
+    data: TensorData
+
+
+class ResponseOutput(BaseModel):
+    name: str
+    shape: List[int]
+    datatype: Datatype
+    parameters: Optional[Parameters] = None
+    data: TensorData
+
+
+class InferenceResponse(BaseModel):
+    model_name: str
+    model_version: Optional[str] = None
+    id: Optional[str] = None
+    parameters: Optional[Parameters] = None
+    outputs: List[ResponseOutput]
 
 
 class MetadataModelResponse(BaseModel):
     name: str
     versions: Optional[List[str]] = None
     platform: str
-    inputs: Optional[List["MetadataTensor"]] = None
-    outputs: Optional[List["MetadataTensor"]] = None
-    parameters: Optional["Parameters"] = None
+    inputs: Optional[List[MetadataTensor]] = None
+    outputs: Optional[List[MetadataTensor]] = None
+    parameters: Optional[Parameters] = None
 
 
 class InferenceRequest(BaseModel):
     id: Optional[str] = None
-    parameters: Optional["Parameters"] = None
-    inputs: List["RequestInput"]
-    outputs: Optional[List["RequestOutput"]] = None
+    parameters: Optional[Parameters] = None
+    inputs: List[RequestInput]
+    outputs: Optional[List[RequestOutput]] = None
diff --git a/openapi/dataplane.yaml b/openapi/dataplane.yaml
@@ -247,7 +247,7 @@ paths:
         - model
 components:
   enums:
-    TensorDataTypes:
+    Datatype:
       type: string
       enum:
         - BOOL
@@ -321,7 +321,8 @@ components:
       properties:
         name:
           type: string
-        datatype: $ref: '#/components/enums/TensorDataTypes'
+        datatype:
+          $ref: '#/components/enums/Datatype'
         shape:
           type: array
           items:
@@ -415,7 +416,8 @@ components:
           type: array
           items:
             type: integer
-        datatype: $ref: '#/components/enums/TensorDataTypes'
+        datatype:
+          $ref: '#/components/enums/Datatype'
         parameters:
           $ref: '#/components/schemas/Parameters'
         data:
@@ -427,7 +429,9 @@ components:
         - data
     TensorData:
       title: TensorData
-      type: array
+      oneOf:
+        - type: array
+        - type: string
     RequestOutput:
       title: RequestOutput
       type: object
@@ -448,7 +452,8 @@ components:
           type: array
           items:
             type: integer
-        datatype: $ref: '#/components/enums/TensorDataTypes'
+        datatype:
+          $ref: '#/components/enums/Datatype'
         parameters:
           $ref: '#/components/schemas/Parameters'
         data: