truefoundry 0.5.1rc6__py3-none-any.whl → 0.5.1rc8__py3-none-any.whl

truefoundry/ml/autogen/client/models/validate_external_storage_root_request_dto.py

@@ -0,0 +1,71 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+
+from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr
+
+
+class ValidateExternalStorageRootRequestDto(BaseModel):
+    """
+    ValidateExternalStorageRootRequestDto
+    """
+
+    storage_root_uri: StrictStr = Field(...)
+    experiment_id: StrictStr = Field(...)
+    __properties = ["storage_root_uri", "experiment_id"]
+
+    class Config:
+        """Pydantic configuration"""
+
+        allow_population_by_field_name = True
+        validate_assignment = True
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.dict(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> ValidateExternalStorageRootRequestDto:
+        """Create an instance of ValidateExternalStorageRootRequestDto from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self):
+        """Returns the dictionary representation of the model using alias"""
+        _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> ValidateExternalStorageRootRequestDto:
+        """Create an instance of ValidateExternalStorageRootRequestDto from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return ValidateExternalStorageRootRequestDto.parse_obj(obj)
+
+        _obj = ValidateExternalStorageRootRequestDto.parse_obj(
+            {
+                "storage_root_uri": obj.get("storage_root_uri"),
+                "experiment_id": obj.get("experiment_id"),
+            }
+        )
+        return _obj

truefoundry/ml/autogen/client/models/validate_external_storage_root_response_dto.py

@@ -0,0 +1,69 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+from typing import Optional
+
+from truefoundry.pydantic_v1 import BaseModel, Field, StrictBool, StrictStr
+
+
+class ValidateExternalStorageRootResponseDto(BaseModel):
+    """
+    ValidateExternalStorageRootResponseDto
+    """
+
+    is_valid: StrictBool = Field(...)
+    message: Optional[StrictStr] = None
+    __properties = ["is_valid", "message"]
+
+    class Config:
+        """Pydantic configuration"""
+
+        allow_population_by_field_name = True
+        validate_assignment = True
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.dict(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> ValidateExternalStorageRootResponseDto:
+        """Create an instance of ValidateExternalStorageRootResponseDto from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self):
+        """Returns the dictionary representation of the model using alias"""
+        _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> ValidateExternalStorageRootResponseDto:
+        """Create an instance of ValidateExternalStorageRootResponseDto from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return ValidateExternalStorageRootResponseDto.parse_obj(obj)
+
+        _obj = ValidateExternalStorageRootResponseDto.parse_obj(
+            {"is_valid": obj.get("is_valid"), "message": obj.get("message")}
+        )
+        return _obj

truefoundry/ml/autogen/client/models/xg_boost_framework.py

@@ -18,12 +18,12 @@ import pprint
 import re  # noqa: F401
 from typing import Optional
 
-from truefoundry.ml.autogen.client.models.serialization_format import (
-    SerializationFormat,
-)
 from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
     XGBoostModelSchema,
 )
+from truefoundry.ml.autogen.client.models.xg_boost_serialization_format import (
+    XGBoostSerializationFormat,
+)
 from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, validator
 
 
@@ -36,10 +36,7 @@ class XGBoostFramework(BaseModel):
         default=...,
         description="+label=Type +usage=Type of the framework +value=xgboost",
     )
-    serialization_format: Optional[SerializationFormat] = Field(
-        default=None,
-        description="+label=Serialization format +usage=Serialization format used for the model",
-    )
+    serialization_format: Optional[XGBoostSerializationFormat] = None
     model_filepath: Optional[StrictStr] = Field(
         default=None,
         description="+label=Model file path +usage=Relative path to the model file",

truefoundry/ml/autogen/client/models/xg_boost_model_schema.py

@@ -18,16 +18,15 @@ import pprint
 import re  # noqa: F401
 from typing import Any, Dict
 
-from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
-from truefoundry.pydantic_v1 import BaseModel, Field, conlist
+from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, conlist, validator
 
 
 class XGBoostModelSchema(BaseModel):
     """
-    XGBoostModelSchema
+    +label=XGBoost Model Schema  # noqa: E501
     """
 
-    infer_method_name: InferMethodName = Field(
+    infer_method_name: StrictStr = Field(
         default=...,
         description="+label=Inference Method Name +usage=Name of the method used for inference",
     )
@@ -39,6 +38,13 @@ class XGBoostModelSchema(BaseModel):
     )
     __properties = ["infer_method_name", "inputs", "outputs"]
 
+    @validator("infer_method_name")
+    def infer_method_name_validate_enum(cls, value):
+        """Validates the enum"""
+        if value not in ("predict",):
+            raise ValueError("must be one of enum values ('predict')")
+        return value
+
     class Config:
         """Pydantic configuration"""
 

truefoundry/ml/autogen/client/models/xg_boost_serialization_format.py

@@ -0,0 +1,36 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+import json
+import re  # noqa: F401
+
+from aenum import Enum
+
+
+class XGBoostSerializationFormat(str, Enum):
+    """
+    +label=Serialization format +usage=Serialization format used for XGBoost models
+    """
+
+    """
+    allowed enum values
+    """
+    CLOUDPICKLE = "cloudpickle"
+    JOBLIB = "joblib"
+    PICKLE = "pickle"
+    JSON = "json"
+
+    @classmethod
+    def from_json(cls, json_str: str) -> XGBoostSerializationFormat:
+        """Create an instance of XGBoostSerializationFormat from a JSON string"""
+        return XGBoostSerializationFormat(json.loads(json_str))

truefoundry/ml/autogen/client_README.md

@@ -133,6 +133,7 @@ Class | Method | HTTP request | Description
 *MlfoundryArtifactsApi* | [**update_artifact_version_post**](truefoundry/ml/autogen/client/docs/MlfoundryArtifactsApi.md#update_artifact_version_post) | **POST** /api/2.0/mlflow/mlfoundry-artifacts/artifact-versions/update | Update Artifact Version
 *MlfoundryArtifactsApi* | [**update_dataset_post**](truefoundry/ml/autogen/client/docs/MlfoundryArtifactsApi.md#update_dataset_post) | **POST** /api/2.0/mlflow/mlfoundry-artifacts/datasets/update | Update Dataset
 *MlfoundryArtifactsApi* | [**update_model_version_post**](truefoundry/ml/autogen/client/docs/MlfoundryArtifactsApi.md#update_model_version_post) | **POST** /api/2.0/mlflow/mlfoundry-artifacts/model-versions/update | Update Model Version
+*MlfoundryArtifactsApi* | [**validate_external_storage_root_path_post**](truefoundry/ml/autogen/client/docs/MlfoundryArtifactsApi.md#validate_external_storage_root_path_post) | **POST** /api/2.0/mlflow/mlfoundry-artifacts/artifact-versions/validate-storage-root | Validate External Storage Root Path
 *RunArtifactsApi* | [**list_run_artifacts_get**](truefoundry/ml/autogen/client/docs/RunArtifactsApi.md#list_run_artifacts_get) | **GET** /api/2.0/mlflow/artifacts/list | List Run Artifacts
 *RunsApi* | [**create_run_post**](truefoundry/ml/autogen/client/docs/RunsApi.md#create_run_post) | **POST** /api/2.0/mlflow/runs/create | Create Run
 *RunsApi* | [**delete_run_post**](truefoundry/ml/autogen/client/docs/RunsApi.md#delete_run_post) | **POST** /api/2.0/mlflow/runs/delete | Delete Run
@@ -299,12 +300,12 @@ Class | Method | HTTP request | Description
  - [RunTagDto](truefoundry/ml/autogen/client/docs/RunTagDto.md)
  - [SearchRunsRequestDto](truefoundry/ml/autogen/client/docs/SearchRunsRequestDto.md)
  - [SearchRunsResponseDto](truefoundry/ml/autogen/client/docs/SearchRunsResponseDto.md)
- - [SerializationFormat](truefoundry/ml/autogen/client/docs/SerializationFormat.md)
  - [SetExperimentTagRequestDto](truefoundry/ml/autogen/client/docs/SetExperimentTagRequestDto.md)
  - [SetTagRequestDto](truefoundry/ml/autogen/client/docs/SetTagRequestDto.md)
  - [SignedURLDto](truefoundry/ml/autogen/client/docs/SignedURLDto.md)
  - [SklearnFramework](truefoundry/ml/autogen/client/docs/SklearnFramework.md)
  - [SklearnModelSchema](truefoundry/ml/autogen/client/docs/SklearnModelSchema.md)
+ - [SklearnSerializationFormat](truefoundry/ml/autogen/client/docs/SklearnSerializationFormat.md)
  - [Source](truefoundry/ml/autogen/client/docs/Source.md)
  - [Source1](truefoundry/ml/autogen/client/docs/Source1.md)
  - [SpaCyFramework](truefoundry/ml/autogen/client/docs/SpaCyFramework.md)
@@ -329,10 +330,13 @@ Class | Method | HTTP request | Description
  - [UpdateRunResponseDto](truefoundry/ml/autogen/client/docs/UpdateRunResponseDto.md)
  - [Url](truefoundry/ml/autogen/client/docs/Url.md)
  - [UserMessage](truefoundry/ml/autogen/client/docs/UserMessage.md)
+ - [ValidateExternalStorageRootRequestDto](truefoundry/ml/autogen/client/docs/ValidateExternalStorageRootRequestDto.md)
+ - [ValidateExternalStorageRootResponseDto](truefoundry/ml/autogen/client/docs/ValidateExternalStorageRootResponseDto.md)
  - [ValidationError](truefoundry/ml/autogen/client/docs/ValidationError.md)
  - [ValidationErrorLocInner](truefoundry/ml/autogen/client/docs/ValidationErrorLocInner.md)
  - [XGBoostFramework](truefoundry/ml/autogen/client/docs/XGBoostFramework.md)
  - [XGBoostModelSchema](truefoundry/ml/autogen/client/docs/XGBoostModelSchema.md)
+ - [XGBoostSerializationFormat](truefoundry/ml/autogen/client/docs/XGBoostSerializationFormat.md)
 
 
 <a id="documentation-for-authorization"></a>

truefoundry/ml/autogen/entities/artifacts.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  artifacts.json
-#   timestamp: 2024-12-05T14:45:34+00:00
+#   timestamp: 2024-12-09T09:04:12+00:00
 
 from __future__ import annotations
 
@@ -60,19 +60,6 @@ class BaseArtifactVersion(BaseModel):
     )
 
 
-class BaseModelSchema(BaseModel):
-    infer_method_name: str = Field(
-        ...,
-        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
-    )
-    inputs: List[Dict[str, Any]] = Field(
-        ..., description="+label= Input Schema\n+usage=Schema of the input"
-    )
-    outputs: List[Dict[str, Any]] = Field(
-        ..., description="+label= Output Schema\n+usage=Schema of the output"
-    )
-
-
 class MimeType(str, Enum):
     """
     +label=MIME Type
@@ -289,17 +276,6 @@ class PyTorchFramework(BaseModel):
     )
 
 
-class SerializationFormat(str, Enum):
-    """
-    +label=Serialization format
-    +usage=Serialization format used for the model
-    """
-
-    cloudpickle = "cloudpickle"
-    joblib = "joblib"
-    pickle = "pickle"
-
-
 class InferMethodName(str, Enum):
     """
     +label=Inference Method Name
@@ -310,11 +286,32 @@ class InferMethodName(str, Enum):
     predict_proba = "predict_proba"
 
 
-class SklearnModelSchema(BaseModelSchema):
+class SklearnModelSchema(BaseModel):
+    """
+    +label=Sklearn Model Schema
+    """
+
     infer_method_name: InferMethodName = Field(
         ...,
         description="+label=Inference Method Name\n+usage=Name of the method used for inference",
     )
+    inputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Input Schema\n+usage=Schema of the input"
+    )
+    outputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Output Schema\n+usage=Schema of the output"
+    )
+
+
+class SklearnSerializationFormat(str, Enum):
+    """
+    +label=Serialization format
+    +usage=Serialization format used for sklearn models
+    """
+
+    cloudpickle = "cloudpickle"
+    joblib = "joblib"
+    pickle = "pickle"
 
 
 class SpaCyFramework(BaseModel):
@@ -448,11 +445,33 @@ class UserMessage(BaseModel):
     )
 
 
-class XGBoostModelSchema(BaseModelSchema):
-    infer_method_name: InferMethodName = Field(
+class XGBoostModelSchema(BaseModel):
+    """
+    +label=XGBoost Model Schema
+    """
+
+    infer_method_name: Literal["predict"] = Field(
         ...,
         description="+label=Inference Method Name\n+usage=Name of the method used for inference",
     )
+    inputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Input Schema\n+usage=Schema of the input"
+    )
+    outputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Output Schema\n+usage=Schema of the output"
+    )
+
+
+class XGBoostSerializationFormat(str, Enum):
+    """
+    +label=Serialization format
+    +usage=Serialization format used for XGBoost models
+    """
+
+    cloudpickle = "cloudpickle"
+    joblib = "joblib"
+    pickle = "pickle"
+    json = "json"
 
 
 class AgentOpenAPITool(BaseModel):
@@ -561,10 +580,7 @@ class SklearnFramework(BaseModel):
         None,
         description="+label=Model file path\n+usage=Relative path to the model file",
     )
-    serialization_format: Optional[SerializationFormat] = Field(
-        None,
-        description="+label=Serialization format\n+usage=Serialization format used for the model",
-    )
+    serialization_format: Optional[SklearnSerializationFormat] = None
     model_schema: Optional[SklearnModelSchema] = None
 
 
@@ -577,10 +593,7 @@ 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",
-    )
+    serialization_format: Optional[XGBoostSerializationFormat] = None
     model_filepath: Optional[str] = Field(
         None,
         description="+label=Model file path\n+usage=Relative path to the model file",

truefoundry/ml/model_framework.py

@@ -20,7 +20,10 @@ from truefoundry.common.utils import (
     get_python_version_major_minor,
     list_pip_packages_installed,
 )
-from truefoundry.ml.autogen.client import SerializationFormat
+from truefoundry.ml.autogen.client import (
+    SklearnSerializationFormat,
+    XGBoostSerializationFormat,
+)
 from truefoundry.ml.autogen.entities import artifacts as autogen_artifacts
 from truefoundry.ml.autogen.models import infer_signature
 from truefoundry.ml.enums import ModelFramework
@@ -35,46 +38,13 @@ if TYPE_CHECKING:
 
 # Map serialization format to corresponding pip packages
 SERIALIZATION_FORMAT_TO_PACKAGES_NAME_MAP = {
-    SerializationFormat.JOBLIB: ["joblib"],
-    SerializationFormat.CLOUDPICKLE: ["cloudpickle"],
+    SklearnSerializationFormat.JOBLIB: ["joblib"],
+    SklearnSerializationFormat.CLOUDPICKLE: ["cloudpickle"],
+    XGBoostSerializationFormat.JOBLIB: ["joblib"],
+    XGBoostSerializationFormat.CLOUDPICKLE: ["cloudpickle"],
 }
 
 
-class _SerializationFormatLoaderRegistry:
-    def __init__(self):
-        # An OrderedDict is used to maintain the order of loaders based on priority
-        # The loaders are added in the following order:
-        #     1. joblib (if available)
-        #     2. cloudpickle (if available)
-        #     3. pickle (default fallback)
-        # This ensures that when looking up a loader, it follows the correct loading priority.
-        self._loader_map: Dict[SerializationFormat, Callable[[bytes], object]] = (
-            OrderedDict()
-        )
-        try:
-            from joblib import load as joblib_load
-
-            self._loader_map[SerializationFormat.JOBLIB] = joblib_load
-        except ImportError:
-            pass
-
-        try:
-            from cloudpickle import load as cloudpickle_load
-
-            self._loader_map[SerializationFormat.CLOUDPICKLE] = cloudpickle_load
-        except ImportError:
-            pass
-
-        # Add pickle loader as a fallback
-        self._loader_map[SerializationFormat.PICKLE] = pickle_load
-
-    def get_loader_map(self) -> Dict[SerializationFormat, Callable[[bytes], object]]:
-        return self._loader_map
-
-
-_serialization_format_loader_registry = _SerializationFormatLoaderRegistry()
-
-
 class FastAIFramework(autogen_artifacts.FastAIFramework):
     """FastAI model Framework"""
 
@@ -180,6 +150,87 @@ ModelFrameworkType = Union[
 ]
 
 
+class _SerializationFormatLoaderRegistry:
+    def __init__(self, framework: Type[Union[SklearnFramework, XGBoostFramework]]):
+        # An OrderedDict is used to maintain the order of loaders based on priority
+        # The loaders are added in the following order:
+        #     1. joblib (if available)
+        #     2. cloudpickle (if available)
+        #     3. pickle (default fallback)
+        # This ensures that when looking up a loader, it follows the correct loading priority.
+        self._loader_map: Dict[
+            Union[SklearnSerializationFormat, XGBoostSerializationFormat],
+            Callable[[bytes], object],
+        ] = OrderedDict()
+        format_class: Union[SklearnSerializationFormat, XGBoostSerializationFormat] = (
+            SklearnSerializationFormat
+            if framework == SklearnFramework
+            else XGBoostSerializationFormat
+        )
+        is_xgboost = framework == XGBoostFramework
+
+        try:
+            from joblib import load as joblib_load
+
+            self._loader_map[format_class.JOBLIB] = joblib_load
+        except ImportError:
+            pass
+
+        try:
+            from cloudpickle import load as cloudpickle_load
+
+            self._loader_map[format_class.CLOUDPICKLE] = cloudpickle_load
+
+        except ImportError:
+            pass
+
+        if is_xgboost:
+            try:
+                from xgboost import Booster
+
+                booster = Booster()
+                self._loader_map[format_class.JSON] = booster.load_model
+            except ImportError:
+                pass
+
+        # Add pickle loader as a fallback
+        self._loader_map[format_class.PICKLE] = pickle_load
+
+    def get_loader_map(
+        self,
+    ) -> Dict[
+        Union[SklearnSerializationFormat, XGBoostSerializationFormat],
+        Callable[[bytes], object],
+    ]:
+        return self._loader_map
+
+    def _detect_model_serialization_format(
+        self,
+        model_file_path: str,
+    ) -> Optional[Union[SklearnSerializationFormat, XGBoostSerializationFormat]]:
+        """
+        The function will attempt to load the model using each different serialization format's loader and return the first successful one.
+
+        Args:
+            model_file_path (str): The path to the file to be loaded.
+
+        Returns:
+            Optional[Union[SklearnSerializationFormat, XGBoostSerializationFormat]]: The serialization format if successfully loaded, None otherwise.
+        """
+        # Attempt to load the model using each framework
+        for (
+            serialization_format,
+            loader,
+        ) in self._loader_map.items():
+            try:
+                with open(model_file_path, "rb") as f:
+                    loader(f)
+                return serialization_format
+            except Exception:
+                continue
+        return None
+
+
 class _ModelFramework(BaseModel):
     __root__: ModelFrameworkType = Field(discriminator="type")
 
@@ -259,32 +310,6 @@ def _get_required_framework_pip_packages(framework: "ModelFrameworkType") -> Lis
     return _MODEL_FRAMEWORK_TO_PIP_PACKAGES.get(framework.__class__, [])
 
 
-def _detect_model_serialization_format(
-    model_file_path: str,
-) -> Optional[SerializationFormat]:
-    """
-    The function will attempt to load the model using each framework's loader and return the first successful one.
-
-    Args:
-        model_file_path (str): The path to the file to be loaded.
-
-    Returns:
-        Optional[SerializationFormat]: The serialization format if successfully loaded, None otherwise.
-    """
-    # Attempt to load the model using each framework
-    for (
-        serialization_format,
-        loader,
-    ) in _serialization_format_loader_registry.get_loader_map().items():
-        try:
-            with open(model_file_path, "rb") as f:
-                loader(f)
-            return serialization_format
-        except Exception:
-            continue
-    return None
-
-
 def _fetch_framework_specific_pip_packages(
     framework: "ModelFrameworkType",
 ) -> List[str]:
@@ -431,12 +456,15 @@ def auto_update_model_framework_details(
                 if os.path.isfile(model_file_or_folder)
                 else os.path.join(model_file_or_folder, framework.model_filepath)
             )
-            framework.serialization_format = (
-                framework.serialization_format
-                or _detect_model_serialization_format(
-                    model_file_path=absolute_model_filepath
+            if not framework.serialization_format:
+                loader_registry = _SerializationFormatLoaderRegistry(
+                    framework=framework
+                )
+                framework.serialization_format = (
+                    loader_registry._detect_model_serialization_format(
+                        model_file_path=absolute_model_filepath
+                    )
                 )
-            )
 
 
 def sklearn_infer_schema(

{truefoundry-0.5.1rc6.dist-info → truefoundry-0.5.1rc8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.5.1rc6
+Version: 0.5.1rc8
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com