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

truefoundry/ml/autogen/client/__init__.py

@@ -363,9 +363,6 @@ from truefoundry.ml.autogen.client.models.search_runs_request_dto import (
 from truefoundry.ml.autogen.client.models.search_runs_response_dto import (
     SearchRunsResponseDto,
 )
-from truefoundry.ml.autogen.client.models.serialization_format import (
-    SerializationFormat,
-)
 from truefoundry.ml.autogen.client.models.set_experiment_tag_request_dto import (
     SetExperimentTagRequestDto,
 )
@@ -373,6 +370,9 @@ from truefoundry.ml.autogen.client.models.set_tag_request_dto import SetTagReque
 from truefoundry.ml.autogen.client.models.signed_url_dto import SignedURLDto
 from truefoundry.ml.autogen.client.models.sklearn_framework import SklearnFramework
 from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
+from truefoundry.ml.autogen.client.models.sklearn_serialization_format import (
+    SklearnSerializationFormat,
+)
 from truefoundry.ml.autogen.client.models.source import Source
 from truefoundry.ml.autogen.client.models.source1 import Source1
 from truefoundry.ml.autogen.client.models.spa_cy_framework import SpaCyFramework
@@ -423,6 +423,12 @@ from truefoundry.ml.autogen.client.models.update_run_response_dto import (
 )
 from truefoundry.ml.autogen.client.models.url import Url
 from truefoundry.ml.autogen.client.models.user_message import UserMessage
+from truefoundry.ml.autogen.client.models.validate_external_storage_root_request_dto import (
+    ValidateExternalStorageRootRequestDto,
+)
+from truefoundry.ml.autogen.client.models.validate_external_storage_root_response_dto import (
+    ValidateExternalStorageRootResponseDto,
+)
 from truefoundry.ml.autogen.client.models.validation_error import ValidationError
 from truefoundry.ml.autogen.client.models.validation_error_loc_inner import (
     ValidationErrorLocInner,
@@ -431,3 +437,6 @@ from truefoundry.ml.autogen.client.models.xg_boost_framework import XGBoostFrame
 from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
     XGBoostModelSchema,
 )
+from truefoundry.ml.autogen.client.models.xg_boost_serialization_format import (
+    XGBoostSerializationFormat,
+)

truefoundry/ml/autogen/client/api/mlfoundry_artifacts_api.py

@@ -167,6 +167,12 @@ from truefoundry.ml.autogen.client.models.update_dataset_request_dto import (
 from truefoundry.ml.autogen.client.models.update_model_version_request_dto import (
     UpdateModelVersionRequestDto,
 )
+from truefoundry.ml.autogen.client.models.validate_external_storage_root_request_dto import (
+    ValidateExternalStorageRootRequestDto,
+)
+from truefoundry.ml.autogen.client.models.validate_external_storage_root_response_dto import (
+    ValidateExternalStorageRootResponseDto,
+)
 from truefoundry.pydantic_v1 import (
     StrictInt,
     StrictStr,
@@ -7215,3 +7221,161 @@ class MlfoundryArtifactsApi:
             collection_formats=_collection_formats,
             _request_auth=_params.get("_request_auth"),
         )
+
+    @validate_arguments
+    def validate_external_storage_root_path_post(
+        self,
+        validate_external_storage_root_request_dto: ValidateExternalStorageRootRequestDto,
+        **kwargs,
+    ) -> ValidateExternalStorageRootResponseDto:  # noqa: E501
+        """Validate External Storage Root Path  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+
+        >>> thread = api.validate_external_storage_root_path_post(validate_external_storage_root_request_dto, async_req=True)
+        >>> result = thread.get()
+
+        :param validate_external_storage_root_request_dto: (required)
+        :type validate_external_storage_root_request_dto: ValidateExternalStorageRootRequestDto
+        :param async_req: Whether to execute the request asynchronously.
+        :type async_req: bool, optional
+        :param _request_timeout: timeout setting for this request.
+               If one number provided, it will be total request
+               timeout. It can also be a pair (tuple) of
+               (connection, read) timeouts.
+        :return: Returns the result object.
+                 If the method is called asynchronously,
+                 returns the request thread.
+        :rtype: ValidateExternalStorageRootResponseDto
+        """
+        kwargs["_return_http_data_only"] = True
+        if "_preload_content" in kwargs:
+            message = "Error! Please call the validate_external_storage_root_path_post_with_http_info method with `_preload_content` instead and obtain raw data from ApiResponse.raw_data"  # noqa: E501
+            raise ValueError(message)
+        return self.validate_external_storage_root_path_post_with_http_info(
+            validate_external_storage_root_request_dto, **kwargs
+        )  # noqa: E501
+
+    @validate_arguments
+    def validate_external_storage_root_path_post_with_http_info(
+        self,
+        validate_external_storage_root_request_dto: ValidateExternalStorageRootRequestDto,
+        **kwargs,
+    ) -> ApiResponse:  # noqa: E501
+        """Validate External Storage Root Path  # noqa: E501
+
+        This method makes a synchronous HTTP request by default. To make an
+        asynchronous HTTP request, please pass async_req=True
+
+        >>> thread = api.validate_external_storage_root_path_post_with_http_info(validate_external_storage_root_request_dto, async_req=True)
+        >>> result = thread.get()
+
+        :param validate_external_storage_root_request_dto: (required)
+        :type validate_external_storage_root_request_dto: ValidateExternalStorageRootRequestDto
+        :param async_req: Whether to execute the request asynchronously.
+        :type async_req: bool, optional
+        :param _preload_content: if False, the ApiResponse.data will
+                                 be set to none and raw_data will store the
+                                 HTTP response body without reading/decoding.
+                                 Default is True.
+        :type _preload_content: bool, optional
+        :param _return_http_data_only: response data instead of ApiResponse
+                                       object with status code, headers, etc
+        :type _return_http_data_only: bool, optional
+        :param _request_timeout: timeout setting for this request. If one
+                                 number provided, it will be total request
+                                 timeout. It can also be a pair (tuple) of
+                                 (connection, read) timeouts.
+        :param _request_auth: set to override the auth_settings for an a single
+                              request; this effectively ignores the authentication
+                              in the spec for a single request.
+        :type _request_auth: dict, optional
+        :type _content_type: string, optional: force content-type for the request
+        :return: Returns the result object.
+                 If the method is called asynchronously,
+                 returns the request thread.
+        :rtype: tuple(ValidateExternalStorageRootResponseDto, status_code(int), headers(HTTPHeaderDict))
+        """
+
+        _params = locals()
+
+        _all_params = ["validate_external_storage_root_request_dto"]
+        _all_params.extend(
+            [
+                "async_req",
+                "_return_http_data_only",
+                "_preload_content",
+                "_request_timeout",
+                "_request_auth",
+                "_content_type",
+                "_headers",
+            ]
+        )
+
+        # validate the arguments
+        for _key, _val in _params["kwargs"].items():
+            if _key not in _all_params:
+                raise ApiTypeError(
+                    "Got an unexpected keyword argument '%s'"
+                    " to method validate_external_storage_root_path_post" % _key
+                )
+            _params[_key] = _val
+        del _params["kwargs"]
+
+        _collection_formats = {}
+
+        # process the path parameters
+        _path_params = {}
+
+        # process the query parameters
+        _query_params = []
+        # process the header parameters
+        _header_params = dict(_params.get("_headers", {}))
+        # process the form parameters
+        _form_params = []
+        _files = {}
+        # process the body parameter
+        _body_params = None
+        if _params["validate_external_storage_root_request_dto"] is not None:
+            _body_params = _params["validate_external_storage_root_request_dto"]
+
+        # set the HTTP header `Accept`
+        _header_params["Accept"] = self.api_client.select_header_accept(
+            ["application/json"]
+        )  # noqa: E501
+
+        # set the HTTP header `Content-Type`
+        _content_types_list = _params.get(
+            "_content_type",
+            self.api_client.select_header_content_type(["application/json"]),
+        )
+        if _content_types_list:
+            _header_params["Content-Type"] = _content_types_list
+
+        # authentication setting
+        _auth_settings = ["HTTPBearer", "APIKeyCookie"]  # noqa: E501
+
+        _response_types_map = {
+            "200": "ValidateExternalStorageRootResponseDto",
+            "422": "HTTPValidationError",
+        }
+
+        return self.api_client.call_api(
+            "/api/2.0/mlflow/mlfoundry-artifacts/artifact-versions/validate-storage-root",
+            "POST",
+            _path_params,
+            _query_params,
+            _header_params,
+            body=_body_params,
+            post_params=_form_params,
+            files=_files,
+            response_types_map=_response_types_map,
+            auth_settings=_auth_settings,
+            async_req=_params.get("async_req"),
+            _return_http_data_only=_params.get("_return_http_data_only"),  # noqa: E501
+            _preload_content=_params.get("_preload_content", True),
+            _request_timeout=_params.get("_request_timeout"),
+            collection_formats=_collection_formats,
+            _request_auth=_params.get("_request_auth"),
+        )

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

@@ -334,9 +334,6 @@ from truefoundry.ml.autogen.client.models.search_runs_request_dto import (
 from truefoundry.ml.autogen.client.models.search_runs_response_dto import (
     SearchRunsResponseDto,
 )
-from truefoundry.ml.autogen.client.models.serialization_format import (
-    SerializationFormat,
-)
 from truefoundry.ml.autogen.client.models.set_experiment_tag_request_dto import (
     SetExperimentTagRequestDto,
 )
@@ -344,6 +341,9 @@ from truefoundry.ml.autogen.client.models.set_tag_request_dto import SetTagReque
 from truefoundry.ml.autogen.client.models.signed_url_dto import SignedURLDto
 from truefoundry.ml.autogen.client.models.sklearn_framework import SklearnFramework
 from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
+from truefoundry.ml.autogen.client.models.sklearn_serialization_format import (
+    SklearnSerializationFormat,
+)
 from truefoundry.ml.autogen.client.models.source import Source
 from truefoundry.ml.autogen.client.models.source1 import Source1
 from truefoundry.ml.autogen.client.models.spa_cy_framework import SpaCyFramework
@@ -394,6 +394,12 @@ from truefoundry.ml.autogen.client.models.update_run_response_dto import (
 )
 from truefoundry.ml.autogen.client.models.url import Url
 from truefoundry.ml.autogen.client.models.user_message import UserMessage
+from truefoundry.ml.autogen.client.models.validate_external_storage_root_request_dto import (
+    ValidateExternalStorageRootRequestDto,
+)
+from truefoundry.ml.autogen.client.models.validate_external_storage_root_response_dto import (
+    ValidateExternalStorageRootResponseDto,
+)
 from truefoundry.ml.autogen.client.models.validation_error import ValidationError
 from truefoundry.ml.autogen.client.models.validation_error_loc_inner import (
     ValidationErrorLocInner,
@@ -402,3 +408,6 @@ from truefoundry.ml.autogen.client.models.xg_boost_framework import XGBoostFrame
 from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
     XGBoostModelSchema,
 )
+from truefoundry.ml.autogen.client.models.xg_boost_serialization_format import (
+    XGBoostSerializationFormat,
+)

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

@@ -18,10 +18,10 @@ 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.sklearn_model_schema import SklearnModelSchema
+from truefoundry.ml.autogen.client.models.sklearn_serialization_format import (
+    SklearnSerializationFormat,
+)
 from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, validator
 
 
@@ -38,10 +38,7 @@ class SklearnFramework(BaseModel):
         default=None,
         description="+label=Model file path +usage=Relative path to the model file",
     )
-    serialization_format: Optional[SerializationFormat] = Field(
-        default=None,
-        description="+label=Serialization format +usage=Serialization format used for the model",
-    )
+    serialization_format: Optional[SklearnSerializationFormat] = None
     model_schema: Optional[SklearnModelSchema] = None
     __properties = ["type", "model_filepath", "serialization_format", "model_schema"]
 

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

@@ -24,7 +24,7 @@ from truefoundry.pydantic_v1 import BaseModel, Field, conlist
 
 class SklearnModelSchema(BaseModel):
     """
-    SklearnModelSchema
+    +label=Sklearn Model Schema  # noqa: E501
     """
 
     infer_method_name: InferMethodName = Field(

truefoundry/ml/autogen/client/models/{serialization_format.py → sklearn_serialization_format.py}

@@ -17,9 +17,9 @@ import re  # noqa: F401
 from aenum import Enum
 
 
-class SerializationFormat(str, Enum):
+class SklearnSerializationFormat(str, Enum):
     """
-    +label=Serialization format +usage=Serialization format used for the model
+    +label=Serialization format +usage=Serialization format used for sklearn models
     """
 
     """
@@ -30,6 +30,6 @@ class SerializationFormat(str, Enum):
     PICKLE = "pickle"
 
     @classmethod
-    def from_json(cls, json_str: str) -> SerializationFormat:
-        """Create an instance of SerializationFormat from a JSON string"""
-        return SerializationFormat(json.loads(json_str))
+    def from_json(cls, json_str: str) -> SklearnSerializationFormat:
+        """Create an instance of SklearnSerializationFormat from a JSON string"""
+        return SklearnSerializationFormat(json.loads(json_str))

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.1rc7.dist-info}/METADATA

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

{truefoundry-0.5.1rc6.dist-info → truefoundry-0.5.1rc7.dist-info}/RECORD

@@ -111,7 +111,7 @@ truefoundry/ml/__init__.py,sha256=VtP0B3DNFFV_1b4WHkAo_XRMvWHmyeEwDu984t4cy8s,20
 truefoundry/ml/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=FksxhUpRHb9pgWZmAB16DhXqkAL6UIAPA1S3RJUApQU,46201
 truefoundry/ml/autogen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/autogen/client/__init__.py,sha256=uf2QeBvpQi38iGSPkAmSRki2b8nJIBx-3AnJOo9_mHM,19571
+truefoundry/ml/autogen/client/__init__.py,sha256=Da3NzHW9aCi9zjQZteTm1aPTuyHDVqrFDtzwsdqIPbc,19981
 truefoundry/ml/autogen/client/api/__init__.py,sha256=NyMBxBmIzB1o5LzZZwz9LiydHB3-hPqc_sevsnY6Jrw,746
 truefoundry/ml/autogen/client/api/auth_api.py,sha256=zpWzJhUmW6HHMY_atlUf0B25k77E1kue2hmix5I5Ih0,7017
 truefoundry/ml/autogen/client/api/deprecated_api.py,sha256=mu5x_skNcwz8v1Tr6VP72-tVP7pmBV9muyKy_2NH3F0,38824
@@ -119,14 +119,14 @@ truefoundry/ml/autogen/client/api/experiments_api.py,sha256=mRKS8qGzcFJUpTWjfQoF
 truefoundry/ml/autogen/client/api/generate_code_snippet_api.py,sha256=C8iABP-pYIyIoTlL8PFi3ZBG_6B_dkQf09eErNeoj7A,22316
 truefoundry/ml/autogen/client/api/health_api.py,sha256=IAPhRAo9CLUT5ipVR1fCf-qDx57UR0wg5ekhtUl8lug,11554
 truefoundry/ml/autogen/client/api/metrics_api.py,sha256=q3L38eD-2hu4_9YvcSdnWDYXD2V8il-X9reinOAABek,14908
-truefoundry/ml/autogen/client/api/mlfoundry_artifacts_api.py,sha256=6AXyTrxQVnBIxELij-4WkdqS-2AJUXZqrxBG_bBuCII,306208
+truefoundry/ml/autogen/client/api/mlfoundry_artifacts_api.py,sha256=5FSOinHzGTbNh-F6KdpxFlVwPrtKVKn4jzr4Dwp9bpw,313554
 truefoundry/ml/autogen/client/api/run_artifacts_api.py,sha256=x-vVnY2LEFChZxiiFauswRWwFz6Qqh30PKXjzuTvxmc,8799
 truefoundry/ml/autogen/client/api/runs_api.py,sha256=-aghrZ2VYuZOw_vBtOzWDsnK7Ji29oZQxK2CLRgyo2w,119232
 truefoundry/ml/autogen/client/api_client.py,sha256=M31IycWorn10FYS8WbO_Whaipr0tVu2pVWFCfbS7XCo,29167
 truefoundry/ml/autogen/client/api_response.py,sha256=KRyvecPMXF05PaxILHZ8JHoP4rgKBjKONMgG83aU-rM,844
 truefoundry/ml/autogen/client/configuration.py,sha256=V1oaEnxt-NfpaNmp-EZpf2glovzVhM2coWYt8HBNB4M,15723
 truefoundry/ml/autogen/client/exceptions.py,sha256=XbCbDHhYT3BVejdoGNPgEa4oS56ypkwFdxk1iOc_tFY,5355
-truefoundry/ml/autogen/client/models/__init__.py,sha256=iX7kGG97Ly93SQfSwFFbggcXAf2OnO9ubNPl0OhN7W4,18193
+truefoundry/ml/autogen/client/models/__init__.py,sha256=4mkEcTiiMIGu16f2idqmFFaeT3T7E1qDIorrhhXReEk,18603
 truefoundry/ml/autogen/client/models/add_custom_metrics_to_model_version_request_dto.py,sha256=_ISDspicTGjBCYYXubKfRYYSSQVyW3AvG-jFh47-Zfc,2163
 truefoundry/ml/autogen/client/models/add_features_to_model_version_request_dto.py,sha256=OT1-98DyWNfAHz_EgD2gMX2XkrGQ4Re945fhoAl8qSE,2099
 truefoundry/ml/autogen/client/models/agent.py,sha256=lUAbs092yo9hZmqzX6atKebDNf1uEwTf5jLvjOdvoeM,3872
@@ -270,12 +270,12 @@ truefoundry/ml/autogen/client/models/run_response_dto.py,sha256=iayBnhwmByVC3iiw
 truefoundry/ml/autogen/client/models/run_tag_dto.py,sha256=CdWwi9P-CXk1vvXWRHN2tHqRckgH-Igur1fowsS-I80,1795
 truefoundry/ml/autogen/client/models/search_runs_request_dto.py,sha256=TxUqdYG5aKj6qdqsqNDe_ih45o0cMw3jwGu2Hrv1X48,2852
 truefoundry/ml/autogen/client/models/search_runs_response_dto.py,sha256=u0l11v7p7zdxtaqV_RBoMtB2Zy8pJbvtMwZiL4B4qMg,2636
-truefoundry/ml/autogen/client/models/serialization_format.py,sha256=CYq9DR1yipEGBA00XjBWyeJf_PIT8_l_F1tm6QafgRM,821
 truefoundry/ml/autogen/client/models/set_experiment_tag_request_dto.py,sha256=nrmi_NxLD1fI2gwlpdqFSMnBS11gRkjS4_GQFHgBcXs,2118
 truefoundry/ml/autogen/client/models/set_tag_request_dto.py,sha256=IRgAdMcWBxmjNV6nZJej4pcNfLmZwrelEZ3otwt7eeE,2144
 truefoundry/ml/autogen/client/models/signed_url_dto.py,sha256=9oHoXBj07xTdc04rqOqJO3eOjQWXCyWPhfHg-6qX60w,1897
-truefoundry/ml/autogen/client/models/sklearn_framework.py,sha256=eWpv-CNyxsD9pIBKga568DpQvCn5Nn_GvNMzY1ElRW4,3424
-truefoundry/ml/autogen/client/models/sklearn_model_schema.py,sha256=RJIbkzvIDxkFTyygQLvUfeg7sCrkhwm9Euo7FaNnrLw,2526
+truefoundry/ml/autogen/client/models/sklearn_framework.py,sha256=1-NDqQL31S6O2qYm_EUTxQT2KFF1XCborUuTMGkGZjE,3318
+truefoundry/ml/autogen/client/models/sklearn_model_schema.py,sha256=ZSmwGRw0R-4iAFUiYjhGNqq8hL_WzHX8rJU_cTXoko4,2549
+truefoundry/ml/autogen/client/models/sklearn_serialization_format.py,sha256=jdJg9MPqVGGDDOn4X7KqxepfGQfGmxpn6SFdcKPjxZo,854
 truefoundry/ml/autogen/client/models/source.py,sha256=ZG2-3oIs_M8mZNPjoIRPzhHqlm63vLCVQC9idUxQ0XU,5180
 truefoundry/ml/autogen/client/models/source1.py,sha256=io35b9uVzqSvOHw6RlCXxM30thN-vHm1upaSiN2i0II,5183
 truefoundry/ml/autogen/client/models/spa_cy_framework.py,sha256=Kyrvhdalnl1uQLoPMGMry2PuFGCcQ8AGp4Cx4B9AGXQ,2150
@@ -300,13 +300,16 @@ truefoundry/ml/autogen/client/models/update_run_request_dto.py,sha256=fn4VB4Zlll
 truefoundry/ml/autogen/client/models/update_run_response_dto.py,sha256=wolIFiihGZCBBkuj7mJEYCmGAOqDLFo_bNLcH4LeXSo,2231
 truefoundry/ml/autogen/client/models/url.py,sha256=zMyOmdVkp1ANnQnc9GrHt42xlVwES7FTjrpOp_OtHqo,4953
 truefoundry/ml/autogen/client/models/user_message.py,sha256=3TEM8qH_zT3dgM197bl44uV7m20IWn6sWeLGfStsN90,2784
+truefoundry/ml/autogen/client/models/validate_external_storage_root_request_dto.py,sha256=vIvIIabschL01XXA1WH33Nx8GTMGC0SJzIEq5HXIKN4,2172
+truefoundry/ml/autogen/client/models/validate_external_storage_root_response_dto.py,sha256=ulOfhhYtuxPoBxTusCcwbu668gW1hew4vq7dGk0itKc,2122
 truefoundry/ml/autogen/client/models/validation_error.py,sha256=mFjwoc8g2-Usu1HXZhOQKQ4TGvLy4lwCzk8dHrJ69aA,2597
 truefoundry/ml/autogen/client/models/validation_error_loc_inner.py,sha256=nThJ5Gmy8W2Wok-ZOI4sK7uRe1BAkLS0qzq-XZbq8zs,4915
-truefoundry/ml/autogen/client/models/xg_boost_framework.py,sha256=ZTvdmfX5eB_ZbA7Yy62NYAzbYKOQkPtA-zF5xxRnVfE,3429
-truefoundry/ml/autogen/client/models/xg_boost_model_schema.py,sha256=5bzi5NEeR9HetiHFJwwkmeC4LE2WvGF9hp9dTLPfEJ4,2526
+truefoundry/ml/autogen/client/models/xg_boost_framework.py,sha256=a0EHzncxe_hSFUCAk3nFOMuh8ztQQQBLbTfg0IbwTE8,3324
+truefoundry/ml/autogen/client/models/xg_boost_model_schema.py,sha256=QKJkiERNqF8UiolX2FjvhkUkUlUgDoo3CpTYPAdCxV8,2735
+truefoundry/ml/autogen/client/models/xg_boost_serialization_format.py,sha256=wFFSH0Z9MP8eT_fysHfUu4iqdjEZART-HxTH4e4gNI8,872
 truefoundry/ml/autogen/client/rest.py,sha256=9goba8qHjQuVx5O_yRaTKu7PvBnb7r7swfy3dwuTEgk,14281
-truefoundry/ml/autogen/client_README.md,sha256=8M-8txeNhpwAiXx7z5UEjPYSq2SoqA1EqtCkkKZDCMc,36276
-truefoundry/ml/autogen/entities/artifacts.py,sha256=x3DxHXPS5OpNb1vmuy6p44g2hymq6BXnivFqv6mQ-ds,20113
+truefoundry/ml/autogen/client_README.md,sha256=9g4j_zKjNJleHNZu_N15KaFRsCAspWI_xMzuHkgq6iQ,36930
+truefoundry/ml/autogen/entities/artifacts.py,sha256=tH_gCB-w2A4Yq8kJck9vLVCQ1CxnIKVzBhDWMdH3Uho,20294
 truefoundry/ml/autogen/models/__init__.py,sha256=--TGRea9SQBWFfwtcl3ekb1XGfFTdEkQGSG8a2SJ60I,187
 truefoundry/ml/autogen/models/exceptions.py,sha256=q3n7FGBrg_hUy1UyoefhMnhcXUAaqXlLIGHoOVzn_d8,1390
 truefoundry/ml/autogen/models/schema.py,sha256=IhpO9qbygLqEamP3NIt3m90SseJXCOm1ZTqNbNbW-M0,55772
@@ -346,7 +349,7 @@ truefoundry/ml/log_types/utils.py,sha256=xjJ21jdPScvFmw3TbVh5NCzbzJwaqiXJyiiT4xx
 truefoundry/ml/logger.py,sha256=VT-BF3BnBYTWVq87O58F0c8uXMu94gYzsiFlGY3_7Ao,458
 truefoundry/ml/mlfoundry_api.py,sha256=PefznXwBPBv8TD4RYid9aiE8WWEkZQAHoaAbvklFoRs,62364
 truefoundry/ml/mlfoundry_run.py,sha256=RLTZGz3htyxNmG1Xat9WevawEIjDoUCbWrSNql8WynI,44597
-truefoundry/ml/model_framework.py,sha256=I04pbczUh6bBmxWI-fzWgpUSY_ab7Js2zDB6fX4zvYM,16458
+truefoundry/ml/model_framework.py,sha256=trqqzy1-BpW3bF4V7wETiOlRBdlcQi2Lk7MXORngmoU,17585
 truefoundry/ml/run_utils.py,sha256=0W208wSLUrbdfk2pjNcZlkUi9bNxG2JORqoe-5rVqHI,2423
 truefoundry/ml/session.py,sha256=F83GTC5WwGBjnJ69Ct8MqMnlutYc56JCc6YhEY1Wl-A,5394
 truefoundry/ml/validation_utils.py,sha256=J5atNhcJLvKj64ralSV9Y5Fv1Rt4SE237ICdP9-7sP4,12149
@@ -367,7 +370,7 @@ truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=5mBCIc-ON
 truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=Hf6Dk6Fu6P7DqsK5ULgraf9DStjgigf-kjaRAMBW-RU,8680
 truefoundry/workflow/task.py,sha256=ToitYiKcNzFCtOVQwz1W8sRjbR97eVS7vQBdbgUQtKg,1779
 truefoundry/workflow/workflow.py,sha256=WaTqUjhwfAXDWu4E5ehuwAxrCbDJkoAf1oWmR2E9Qy0,4575
-truefoundry-0.5.1rc6.dist-info/METADATA,sha256=sCfEogpg-mXkzRc_S6z6T3zFA7fvzj9k6rL83R538w8,2887
-truefoundry-0.5.1rc6.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-truefoundry-0.5.1rc6.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.5.1rc6.dist-info/RECORD,,
+truefoundry-0.5.1rc7.dist-info/METADATA,sha256=p7N_7W876IqUUAqgllMzYDrKisvxTAPUmfib0G30Uwo,2887
+truefoundry-0.5.1rc7.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+truefoundry-0.5.1rc7.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
+truefoundry-0.5.1rc7.dist-info/RECORD,,