truefoundry 0.4.10__py3-none-any.whl → 0.5.0__py3-none-any.whl

truefoundry/ml/autogen/entities/artifacts.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  artifacts.json
-#   timestamp: 2024-10-25T06:52:16+00:00
+#   timestamp: 2024-11-29T06:43:26+00:00
 
 from __future__ import annotations
 
@@ -89,7 +89,7 @@ class BlobStorageReference(BaseModel):
     )
 
 
-class ExternalModelSource(BaseModel):
+class ExternalArtifactSource(BaseModel):
     """
     +label=External
     """
@@ -98,7 +98,7 @@ class ExternalModelSource(BaseModel):
         ...,
         description="+label=Type\n+usage=Type of the source\n+value=external\n+type=External",
     )
-    uri: str = Field(..., description="+label=URI\n+usage=URI of the model source")
+    uri: str = Field(..., description="+label=URI\n+usage=URI of the Artifact source")
 
 
 class FastAIFramework(BaseModel):
@@ -226,6 +226,34 @@ class ModelConfiguration(BaseModel):
     )
 
 
+class ModelVersionEnvironment(BaseModel):
+    """
+    +label=ModelVersionEnvironment
+    """
+
+    python_version: Optional[constr(regex=r"^\d+(\.\d+){1,2}([\-\.a-z0-9]+)?$")] = (
+        Field(
+            None,
+            description="+label=Python Version\n+usage=Python version for the model version",
+        )
+    )
+    pip_packages: Optional[List[str]] = Field(
+        None,
+        description="+label=PIP Packages\n+usage=PIP packages for the model version",
+    )
+
+
+class ONNXFramework(BaseModel):
+    """
+    +docs=ONNX framework for the model version
+    +label=ONNX
+    """
+
+    type: Literal["onnx"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=onnx"
+    )
+
+
 class PaddleFramework(BaseModel):
     """
     +docs=PaddlePaddle framework for the model version
@@ -248,6 +276,17 @@ class PyTorchFramework(BaseModel):
     )
 
 
+class SerializationFormat(str, Enum):
+    """
+    +label=Serialization format
+    +usage=Serialization format used for the model
+    """
+
+    cloudpickle = "cloudpickle"
+    joblib = "joblib"
+    pickle = "pickle"
+
+
 class SklearnFramework(BaseModel):
     """
     +docs=Scikit-learn framework for the model version
@@ -257,6 +296,13 @@ class SklearnFramework(BaseModel):
     type: Literal["sklearn"] = Field(
         ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
     )
+    serialization_format: Optional[SerializationFormat] = Field(
+        None,
+        description="+label=Serialization format\n+usage=Serialization format used for the model",
+    )
+    model_filename: Optional[str] = Field(
+        None, description="+label=Model file name\n+usage=Name of the model file"
+    )
 
 
 class SpaCyFramework(BaseModel):
@@ -348,11 +394,15 @@ class TransformersFramework(BaseModel):
         None,
         description="+label=Pipeline Tag\n+usage=Pipeline tag\n+docs=Pipeline tag for the framework",
     )
+    base_model: Optional[str] = Field(
+        None,
+        description="+label=Base Model\n+usage=Base model\n+docs=Base model Id. If this is a finetuned model, this points to the base model used for finetuning",
+    )
 
 
-class TruefoundryModelSource(BaseModel):
+class TrueFoundryArtifactSource(BaseModel):
     """
-    +usage=Source for the Model
+    +usage=Source for the Artifact or Model
     +label=Upload
     """
 
@@ -361,7 +411,8 @@ class TruefoundryModelSource(BaseModel):
         description="+usage=Type of the source\n+type=Upload\n+value=truefoundry\n+label=Upload",
     )
     uri: str = Field(
-        ..., description='+label=URI\n+uiType=UploadInput\n+uiProps={"hideClear":true}'
+        ...,
+        description='+label=URI\n+uiType=ModelUploadInput\n+uiProps={"hideClear":true}',
     )
 
 
@@ -394,6 +445,13 @@ class XGBoostFramework(BaseModel):
     type: Literal["xgboost"] = Field(
         ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
     )
+    serialization_format: Optional[SerializationFormat] = Field(
+        None,
+        description="+label=Serialization format\n+usage=Serialization format used for the model",
+    )
+    model_filename: Optional[str] = Field(
+        None, description="+label=Model file name\n+usage=Name of the model file"
+    )
 
 
 class AgentOpenAPITool(BaseModel):
@@ -426,6 +484,18 @@ class AgentOpenAPIToolWithFQN(AgentOpenAPITool):
     fqn: str
 
 
+class ArtifactVersion(BaseArtifactVersion):
+    type: Literal["artifact-version"] = Field(
+        ...,
+        description='+label=Type\n+usage=Artifact Version\n+value="artifact-version"',
+    )
+    source: Union[TrueFoundryArtifactSource, ExternalArtifactSource] = Field(
+        ...,
+        description="+label=Artifact Source\n+usage=Source for the Artifact version\n+uiType=Group",
+    )
+    step: conint(ge=0) = Field(0, description="+label=Step")
+
+
 class AssistantMessage(BaseModel):
     """
     +usage=Assistant message
@@ -481,7 +551,7 @@ class ModelVersion(BaseArtifactVersion):
     type: Literal["model-version"] = Field(
         ..., description='+label=Type\n+usage=Model Version\n+value="model-version"'
     )
-    source: Union[TruefoundryModelSource, ExternalModelSource] = Field(
+    source: Union[TrueFoundryArtifactSource, ExternalArtifactSource] = Field(
         ...,
         description="+label=Model Source\n+usage=Source for the model version\n+uiType=Group",
     )
@@ -496,6 +566,7 @@ class ModelVersion(BaseArtifactVersion):
             LightGBMFramework,
             FastAIFramework,
             H2OFramework,
+            ONNXFramework,
             SpaCyFramework,
             StatsModelsFramework,
             GluonFramework,
@@ -504,7 +575,12 @@ class ModelVersion(BaseArtifactVersion):
     ] = Field(
         None, description="+label=Framework\n+usage=Framework for the model version"
     )
+    environment: Optional[ModelVersionEnvironment] = None
     step: conint(ge=0) = Field(0, description="+label=Step")
+    model_schema: Optional[Dict[str, Any]] = Field(
+        None,
+        description="+label=Model Schema\n+usage=Schema of the model\n+uiType=Hidden",
+    )
 
 
 class AgentApp(BaseModel):
@@ -514,5 +590,7 @@ class AgentApp(BaseModel):
     root_agent: constr(min_length=1)
 
 
-class ArtifactsVersion(BaseModel):
-    __root__: Union[ChatPrompt, AgentOpenAPITool, Agent, AgentApp, ModelVersion]
+class VersionedArtifactType(BaseModel):
+    __root__: Union[
+        ChatPrompt, AgentOpenAPITool, Agent, AgentApp, ModelVersion, ArtifactVersion
+    ]

truefoundry/ml/autogen/models/__init__.py

@@ -0,0 +1,4 @@
+from .schema import ColSpec, TensorSpec
+from .signature import ModelSignature, Schema, infer_signature
+
+__all__ = ["infer_signature", "ModelSignature", "Schema", "TensorSpec", "ColSpec"]

truefoundry/ml/autogen/models/exceptions.py

@@ -0,0 +1,30 @@
+INTERNAL_ERROR = 1
+INVALID_PARAMETER_VALUE = 1000
+
+
+class MlflowException(Exception):
+    """
+    Generic exception thrown to surface failure information about external-facing operations.
+    The error message associated with this exception may be exposed to clients in HTTP responses
+    for debugging purposes. If the error text is sensitive, raise a generic `Exception` object
+    instead.
+    """
+
+    def __init__(self, message, error_code=INTERNAL_ERROR, **kwargs):
+        """
+        :param message: The message describing the error that occured. This will be included in the
+                        exception's serialized JSON representation.
+        :param error_code: An appropriate error code for the error that occured; it will be included
+                           in the exception's serialized JSON representation. This should be one of
+                           the codes listed in the `mlflow.protos.databricks_pb2` proto.
+        :param kwargs: Additional key-value pairs to include in the serialized JSON representation
+                       of the MlflowException.
+        """
+        self.error_code = error_code
+        self.message = str(message)
+        self.json_kwargs = kwargs
+        super().__init__(message)
+
+    @classmethod
+    def invalid_parameter_value(cls, message, **kwargs):
+        return cls(message, error_code=INVALID_PARAMETER_VALUE, **kwargs)