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

truefoundry/common/constants.py

@@ -1,3 +1,4 @@
+import enum
 import os
 from pathlib import Path
 from typing import Optional
@@ -15,6 +16,11 @@ TFY_INTERNAL_SIGNED_URL_SERVER_HOST_ENV_KEY = "TFY_INTERNAL_SIGNED_URL_SERVER_HO
 TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY = "TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN"
 
 
+class PythonPackageManager(str, enum.Enum):
+    PIP = "pip"
+    UV = "uv"
+
+
 class TrueFoundrySdkEnv(BaseSettings):
     # Note: Every field in this class should have a default value
     # Never expect the user to set these values
@@ -46,6 +52,9 @@ class TrueFoundrySdkEnv(BaseSettings):
     # For local development, this enables futher configuration via _TFYServersConfig
     TFY_CLI_LOCAL_DEV_MODE: bool = False
 
+    TFY_PYTHON_BUILD_PACKAGE_MANAGER: PythonPackageManager = PythonPackageManager.PIP
+    TFY_PYTHON_BUILD_UV_IMAGE_URI: str = "ghcr.io/astral-sh/uv:latest"
+
 
 ENV_VARS = TrueFoundrySdkEnv()
 API_SERVER_RELATIVE_PATH = "api/svc"

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py

@@ -8,7 +8,7 @@ from truefoundry.deploy.builder.builders.tfy_notebook_buildpack.dockerfile_templ
     NotebookImageBuild,
     generate_dockerfile_content,
 )
-from truefoundry.deploy.builder.utils import has_pip_conf_secret
+from truefoundry.deploy.builder.utils import has_python_package_manager_conf_secret
 
 __all__ = ["generate_dockerfile_content", "build"]
 
@@ -38,7 +38,9 @@ def build(
     build_configuration: NotebookImageBuild,
     extra_opts: Optional[List[str]] = None,
 ):
-    mount_pip_conf_secret = has_pip_conf_secret(extra_opts) if extra_opts else False
+    mount_pip_conf_secret = (
+        has_python_package_manager_conf_secret(extra_opts) if extra_opts else False
+    )
     with TemporaryDirectory() as local_dir:
         docker_build_configuration = _convert_to_dockerfile_build_config(
             build_configuration,

truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py

@@ -7,7 +7,7 @@ from truefoundry.deploy.builder.builders import dockerfile
 from truefoundry.deploy.builder.builders.tfy_python_buildpack.dockerfile_template import (
     generate_dockerfile_content,
 )
-from truefoundry.deploy.builder.utils import has_pip_conf_secret
+from truefoundry.deploy.builder.utils import has_python_package_manager_conf_secret
 
 __all__ = ["generate_dockerfile_content", "build"]
 
@@ -15,11 +15,11 @@ __all__ = ["generate_dockerfile_content", "build"]
 def _convert_to_dockerfile_build_config(
     build_configuration: PythonBuild,
     dockerfile_path: str,
-    mount_pip_conf_secret: bool = False,
+    mount_python_package_manager_conf_secret: bool = False,
 ) -> DockerFileBuild:
     dockerfile_content = generate_dockerfile_content(
         build_configuration=build_configuration,
-        mount_pip_conf_secret=mount_pip_conf_secret,
+        mount_python_package_manager_conf_secret=mount_python_package_manager_conf_secret,
     )
     with open(dockerfile_path, "w", encoding="utf8") as fp:
         fp.write(dockerfile_content)
@@ -36,12 +36,14 @@ def build(
     build_configuration: PythonBuild,
     extra_opts: Optional[List[str]] = None,
 ):
-    mount_pip_conf_secret = has_pip_conf_secret(extra_opts) if extra_opts else False
+    mount_python_package_manager_conf_secret = (
+        has_python_package_manager_conf_secret(extra_opts) if extra_opts else False
+    )
     with TemporaryDirectory() as local_dir:
         docker_build_configuration = _convert_to_dockerfile_build_config(
             build_configuration,
             dockerfile_path=os.path.join(local_dir, "Dockerfile"),
-            mount_pip_conf_secret=mount_pip_conf_secret,
+            mount_python_package_manager_conf_secret=mount_python_package_manager_conf_secret,
         )
         dockerfile.build(
             tag=tag,

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -1,13 +1,16 @@
 import os
+import shlex
 from typing import Dict, List, Optional
 
 from mako.template import Template
 
-from truefoundry.common.constants import ENV_VARS
+from truefoundry.common.constants import ENV_VARS, PythonPackageManager
 from truefoundry.deploy.auto_gen.models import PythonBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
     PIP_CONF_SECRET_MOUNT_AS_ENV,
+    UV_CONF_BUILDKIT_SECRET_MOUNT,
+    UV_CONF_SECRET_MOUNT_AS_ENV,
 )
 from truefoundry.deploy.v2.lib.patched_models import CUDAVersion
 
@@ -22,8 +25,8 @@ RUN ${apt_install_command}
 % if requirements_path is not None:
 COPY ${requirements_path} ${requirements_destination_path}
 % endif
-% if pip_install_command is not None:
-RUN ${pip_config_secret_mount} ${pip_install_command}
+% if python_packages_install_command is not None:
+RUN ${package_manager_config_secret_mount} ${python_packages_install_command}
 % endif
 COPY . /app
 WORKDIR /app
@@ -114,44 +117,91 @@ def generate_pip_install_command(
     mount_pip_conf_secret: bool = False,
 ) -> Optional[str]:
     upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
-    final_pip_install_command = None
-    pip_install_base_command = "python -m pip install --use-pep517 --no-cache-dir"
+    envs = []
+    if mount_pip_conf_secret:
+        envs.append(PIP_CONF_SECRET_MOUNT_AS_ENV)
+
+    command = ["python", "-m", "pip", "install", "--use-pep517", "--no-cache-dir"]
+    args = []
     if requirements_path:
-        final_pip_install_command = f"{pip_install_base_command} -r {requirements_path}"
+        args.append("-r")
+        args.append(requirements_path)
 
     if pip_packages:
-        final_pip_install_command = (
-            final_pip_install_command or pip_install_base_command
-        )
-        final_pip_install_command += " " + " ".join(
-            f"'{package}'" for package in pip_packages
-        )
+        args.extend(pip_packages)
 
-    if not final_pip_install_command:
+    if not args:
         return None
 
-    if mount_pip_conf_secret:
-        final_pip_install_command = (
-            f"{PIP_CONF_SECRET_MOUNT_AS_ENV} {final_pip_install_command}"
-        )
+    final_pip_install_command = shlex.join(envs + command + args)
+    final_docker_run_command = " && ".join(
+        [upgrade_pip_command, final_pip_install_command]
+    )
+    return final_docker_run_command
 
-    return " && ".join([upgrade_pip_command, final_pip_install_command])
+
+def generate_uv_pip_install_command(
+    requirements_path: Optional[str],
+    pip_packages: Optional[List[str]],
+    mount_uv_conf_secret: bool = False,
+) -> Optional[str]:
+    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
+    uv_mount = f"--mount=from={ENV_VARS.TFY_PYTHON_BUILD_UV_IMAGE_URI},source=/uv,target=/usr/local/bin/uv"
+    envs = [
+        "UV_SYSTEM_PYTHON=true",
+        "UV_LINK_MODE=copy",
+        "UV_PYTHON_DOWNLOADS=never",
+        "UV_INDEX_STRATEGY=unsafe-best-match",
+    ]
+    if mount_uv_conf_secret:
+        envs.append(UV_CONF_SECRET_MOUNT_AS_ENV)
+
+    command = ["uv", "pip", "install", "--no-cache-dir"]
+
+    args = []
+
+    if requirements_path:
+        args.append("-r")
+        args.append(requirements_path)
+
+    if pip_packages:
+        args.extend(pip_packages)
+
+    if not args:
+        return None
+
+    uv_pip_install_command = shlex.join(envs + command + args)
+    shell_commands = " && ".join([upgrade_pip_command, uv_pip_install_command])
+    final_docker_run_command = " ".join([uv_mount, shell_commands])
+
+    return final_docker_run_command
 
 
 def generate_dockerfile_content(
     build_configuration: PythonBuild,
-    mount_pip_conf_secret: bool = False,
+    package_manager: str = ENV_VARS.TFY_PYTHON_BUILD_PACKAGE_MANAGER,
+    mount_python_package_manager_conf_secret: bool = False,
 ) -> str:
     # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
     requirements_path = resolve_requirements_txt_path(build_configuration)
     requirements_destination_path = (
         "/tmp/requirements.txt" if requirements_path else None
     )
-    pip_install_command = generate_pip_install_command(
-        requirements_path=requirements_destination_path,
-        pip_packages=build_configuration.pip_packages,
-        mount_pip_conf_secret=mount_pip_conf_secret,
-    )
+    if package_manager == PythonPackageManager.PIP.value:
+        python_packages_install_command = generate_pip_install_command(
+            requirements_path=requirements_destination_path,
+            pip_packages=build_configuration.pip_packages,
+            mount_pip_conf_secret=mount_python_package_manager_conf_secret,
+        )
+    elif package_manager == PythonPackageManager.UV.value:
+        python_packages_install_command = generate_uv_pip_install_command(
+            requirements_path=requirements_destination_path,
+            pip_packages=build_configuration.pip_packages,
+            mount_uv_conf_secret=mount_python_package_manager_conf_secret,
+        )
+    else:
+        raise ValueError(f"Unsupported package manager: {package_manager}")
+
     apt_install_command = generate_apt_install_command(
         apt_packages=build_configuration.apt_packages
     )
@@ -161,13 +211,22 @@ def generate_dockerfile_content(
         "apt_install_command": apt_install_command,
         "requirements_path": requirements_path,
         "requirements_destination_path": requirements_destination_path,
-        "pip_install_command": pip_install_command,
+        "python_packages_install_command": python_packages_install_command,
     }
 
-    if mount_pip_conf_secret:
-        template_args["pip_config_secret_mount"] = PIP_CONF_BUILDKIT_SECRET_MOUNT
+    if mount_python_package_manager_conf_secret:
+        if package_manager == PythonPackageManager.PIP.value:
+            template_args["package_manager_config_secret_mount"] = (
+                PIP_CONF_BUILDKIT_SECRET_MOUNT
+            )
+        elif package_manager == PythonPackageManager.UV.value:
+            template_args["package_manager_config_secret_mount"] = (
+                UV_CONF_BUILDKIT_SECRET_MOUNT
+            )
+        else:
+            raise ValueError(f"Unsupported package manager: {package_manager}")
     else:
-        template_args["pip_config_secret_mount"] = ""
+        template_args["package_manager_config_secret_mount"] = ""
 
     if build_configuration.cuda_version:
         template = CUDA_DOCKERFILE_TEMPLATE

truefoundry/deploy/builder/constants.py

@@ -5,3 +5,11 @@ PIP_CONF_BUILDKIT_SECRET_MOUNT = (
 PIP_CONF_SECRET_MOUNT_AS_ENV = (
     f"PIP_CONFIG_FILE=/run/secrets/{BUILDKIT_SECRET_MOUNT_PIP_CONF_ID}"
 )
+
+BUILDKIT_SECRET_MOUNT_UV_CONF_ID = "uv.toml"
+UV_CONF_BUILDKIT_SECRET_MOUNT = (
+    f"--mount=type=secret,id={BUILDKIT_SECRET_MOUNT_UV_CONF_ID}"
+)
+UV_CONF_SECRET_MOUNT_AS_ENV = (
+    f"UV_CONFIG_FILE=/run/secrets/{BUILDKIT_SECRET_MOUNT_UV_CONF_ID}"
+)

truefoundry/deploy/builder/utils.py

@@ -1,6 +1,9 @@
 from typing import List, Optional
 
-from truefoundry.deploy.builder.constants import BUILDKIT_SECRET_MOUNT_PIP_CONF_ID
+from truefoundry.deploy.builder.constants import (
+    BUILDKIT_SECRET_MOUNT_PIP_CONF_ID,
+    BUILDKIT_SECRET_MOUNT_UV_CONF_ID,
+)
 
 
 def _get_id_from_buildkit_secret_value(value: str) -> Optional[str]:
@@ -19,14 +22,16 @@ def _get_id_from_buildkit_secret_value(value: str) -> Optional[str]:
     return None
 
 
-def has_pip_conf_secret(docker_build_extra_args: List[str]) -> bool:
+def has_python_package_manager_conf_secret(docker_build_extra_args: List[str]) -> bool:
     args = [arg.strip() for arg in docker_build_extra_args]
     for i, arg in enumerate(docker_build_extra_args):
         if (
             arg == "--secret"
             and i + 1 < len(args)
-            and _get_id_from_buildkit_secret_value(args[i + 1])
-            == BUILDKIT_SECRET_MOUNT_PIP_CONF_ID
+            and (
+                _get_id_from_buildkit_secret_value(args[i + 1])
+                in (BUILDKIT_SECRET_MOUNT_PIP_CONF_ID, BUILDKIT_SECRET_MOUNT_UV_CONF_ID)
+            )
         ):
             return True
     return False

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))