truefoundry 0.5.0rc6__py3-none-any.whl → 0.5.1__py3-none-any.whl

truefoundry/ml/model_framework.py

@@ -1,10 +1,50 @@
+import json
+import os
 import warnings
-from typing import Any, Dict, Literal, Optional, Union, get_args
-
-from truefoundry.ml import ModelFramework
+from collections import OrderedDict
+from pickle import load as pickle_load
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Type,
+    Union,
+    get_args,
+)
+
+from truefoundry.common.utils import (
+    get_python_version_major_minor,
+    list_pip_packages_installed,
+)
+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
+from truefoundry.ml.log_types.artifacts.utils import (
+    get_single_file_path_if_only_one_in_directory,
+    to_unix_path,
+)
 from truefoundry.pydantic_v1 import BaseModel, Field
 
+if TYPE_CHECKING:
+    from sklearn.base import BaseEstimator
+    from xgboost import Booster, XGBModel
+
+# Map serialization format to corresponding pip packages
+SERIALIZATION_FORMAT_TO_PACKAGES_NAME_MAP = {
+    SklearnSerializationFormat.JOBLIB: ["joblib"],
+    SklearnSerializationFormat.CLOUDPICKLE: ["cloudpickle"],
+    XGBoostSerializationFormat.JOBLIB: ["joblib"],
+    XGBoostSerializationFormat.CLOUDPICKLE: ["cloudpickle"],
+}
+
 
 class FastAIFramework(autogen_artifacts.FastAIFramework):
     """FastAI model Framework"""
@@ -111,6 +151,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")
 
@@ -167,3 +288,251 @@ class _ModelFramework(BaseModel):
             return None
 
         return cls.parse_obj(obj).__root__
+
+
+# Mapping of model frameworks to pip packages
+_MODEL_FRAMEWORK_TO_PIP_PACKAGES: Dict[Type[ModelFrameworkType], List[str]] = {
+    SklearnFramework: ["scikit-learn", "numpy", "pandas"],
+    XGBoostFramework: ["xgboost", "numpy", "pandas"],
+}
+
+
+def _get_required_framework_pip_packages(framework: "ModelFrameworkType") -> List[str]:
+    """
+    Fetches the pip packages required for a given model framework.
+
+    Args:
+        framework ("ModelFrameworkType"): The model framework for which to fetch the pip packages.
+
+    Returns:
+        List[str]: The list of pip packages required for the given model framework.
+                  If no packages are found for the framework type, returns an empty list.
+    """
+    return _MODEL_FRAMEWORK_TO_PIP_PACKAGES.get(framework.__class__, [])
+
+
+def _fetch_framework_specific_pip_packages(
+    framework: "ModelFrameworkType",
+) -> List[str]:
+    """
+    Fetch the pip packages required for the given framework, including any dependencies
+    related to the framework's serialization format.
+
+    Args:
+        framework: The framework object (e.g., SklearnFramework, XGBoostFramework).
+
+    Returns:
+        List[str]: A list of pip packages for the given framework and environment,
+                   including any dependencies based on the serialization format
+                   (e.g., ['numpy==1.19.5', ...]).
+    """
+    framework_package_names = _get_required_framework_pip_packages(framework=framework)
+
+    # Add serialization format dependencies if applicable
+    if isinstance(framework, (SklearnFramework, XGBoostFramework)):
+        framework_package_names.extend(
+            SERIALIZATION_FORMAT_TO_PACKAGES_NAME_MAP.get(
+                framework.serialization_format, []
+            )
+        )
+    return [
+        f"{package.name}=={package.version}"
+        for package in list_pip_packages_installed(
+            filter_package_names=framework_package_names
+        )
+    ]
+
+
+def auto_update_environment_details(
+    environment: autogen_artifacts.ModelVersionEnvironment,
+    framework: Optional[ModelFrameworkType],
+):
+    """
+    Auto fetch the environment details if not provided, based on the provided environment and framework.
+
+    Args:
+        environment: The environment object that holds environment details like python_version and pip_packages.
+        framework: The framework object (e.g., SklearnFramework, XGBoostFramework) that may affect pip_package fetching.
+    """
+    # Auto fetch python_version if not provided
+    if not environment.python_version:
+        environment.python_version = get_python_version_major_minor()
+
+    # Framework-specific pip_package handling
+    if framework and not environment.pip_packages:
+        environment.pip_packages = _fetch_framework_specific_pip_packages(framework)
+
+
+def _validate_and_get_absolute_model_filepath(
+    model_file_or_folder: str,
+    model_filepath: Optional[str] = None,
+) -> Optional[str]:
+    # If no model_filepath is set, resolve it from the directory
+    if not model_filepath:
+        # If model_filepath is not set, resolve it based on these cases:
+        #    - Case 1: model_file_or_folder/model.joblib -> model.joblib
+        #    - Case 2: model_file_or_folder/folder/model.joblib -> folder/model.joblib
+        #    - Case 3: model_file_or_folder/folder/model.joblib, model_file_or_folder/config.json -> None
+        return get_single_file_path_if_only_one_in_directory(model_file_or_folder)
+
+    # If model_filepath is already set, validate and resolve it:
+    #    - Case 1: Resolve the absolute file path of the model file relative to the provided directory.
+    #        Example: If model_file_or_folder is '/root/models' and model_filepath is 'model.joblib',
+    #                 the resolved model file path would be '/root/models/model.joblib'. Validate it.
+    #
+    #    - Case 2: If model_filepath is a relative path, resolve it to an absolute path based on the provided directory.
+    #        Example: If model_file_or_folder is '/root/models' and model_filepath is 'subfolder/model.joblib',
+    #                 the resolved path would be '/root/models/subfolder/model.joblib'. Validate it.
+    #
+    #    - Case 3: Verify that the resolved model file exists and is a valid file.
+    #        Example: If the resolved path is '/root/models/model.joblib', check if the file exists.
+    #                 If it does not exist, raise a FileNotFoundError.
+    #
+    #    - Case 4: Ensure the resolved model file is located within the specified directory or is the directory itself.
+    #        Example: If the resolved path is '/root/models/model.joblib' and model_file_or_folder is '/root/models',
+    #                 the resolved path is valid. If the file lies outside '/root/models', raise a ValueError.
+    #
+
+    # If model_filepath is set, Resolve the absolute path of the model file (It can be a relative path or absolute path)
+    model_dir = (
+        os.path.dirname(model_file_or_folder)
+        if os.path.isfile(model_file_or_folder)
+        else model_file_or_folder
+    )
+    absolute_model_filepath = os.path.abspath(os.path.join(model_dir, model_filepath))
+
+    # Validate if resolve valid is within the provided directory or is the same as it
+    if not (
+        absolute_model_filepath == model_file_or_folder
+        or absolute_model_filepath.startswith(model_file_or_folder + os.sep)
+    ):
+        raise ValueError(
+            f"model_filepath '{model_filepath}' must be relative to "
+            f"{model_file_or_folder}. Resolved path '{absolute_model_filepath}' is invalid."
+        )
+
+    if not os.path.isfile(absolute_model_filepath):
+        raise FileNotFoundError(f"Model file not found: {absolute_model_filepath}")
+
+    return absolute_model_filepath
+
+
+def _validate_and_resolve_model_filepath(
+    model_file_or_folder: str,
+    model_filepath: Optional[str] = None,
+) -> Optional[str]:
+    absolute_model_filepath = _validate_and_get_absolute_model_filepath(
+        model_file_or_folder=model_file_or_folder, model_filepath=model_filepath
+    )
+    if absolute_model_filepath:
+        return to_unix_path(
+            os.path.relpath(absolute_model_filepath, model_file_or_folder)
+            if os.path.isdir(model_file_or_folder)
+            else os.path.basename(absolute_model_filepath)
+        )
+
+
+def auto_update_model_framework_details(
+    framework: "ModelFrameworkType", model_file_or_folder: str
+):
+    """
+    Auto update the model framework details based on the provided model file or folder path.
+
+    Args:
+        framework: The framework object (e.g., SklearnFramework, XGBoostFramework) to update.
+        model_file_or_folder: The path to the model file or folder.
+    """
+
+    # Ensure the model file or folder path is an absolute path
+    model_file_or_folder = os.path.abspath(model_file_or_folder)
+
+    if isinstance(framework, (SklearnFramework, XGBoostFramework)):
+        framework.model_filepath = _validate_and_resolve_model_filepath(
+            model_file_or_folder=model_file_or_folder,
+            model_filepath=framework.model_filepath,
+        )
+        if framework.model_filepath:
+            absolute_model_filepath = (
+                model_file_or_folder
+                if os.path.isfile(model_file_or_folder)
+                else os.path.join(model_file_or_folder, framework.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 _infer_schema(
+    model_input: Any,
+    model: Union["BaseEstimator", "Booster", "XGBModel"],
+    infer_method_name: str = "predict",
+) -> Dict[str, Any]:
+    if not hasattr(model, infer_method_name):
+        raise ValueError(
+            f"Model does not have the method '{infer_method_name}' to infer the schema."
+        )
+    model_infer_method = getattr(model, infer_method_name)
+    model_output = model_infer_method(model_input)
+
+    model_signature = infer_signature(
+        model_input=model_input, model_output=model_output
+    )
+    return model_signature.to_dict()
+
+
+def sklearn_infer_schema(
+    model_input: Any,
+    model: "BaseEstimator",
+    infer_method_name: str = "predict",
+) -> autogen_artifacts.SklearnModelSchema:
+    """
+    Infer the schema of a Sklearn model.
+
+    Args:
+        model_input (Any): The input data to be used for schema inference.
+        model (Any): The Sklearn model instance.
+        infer_method_name (str): The name of the method to be used for schema inference.
+            Eg: predict (default), predict_proba
+    Returns:
+        SklearnModelSchema: The inferred schema of the Sklearn model.
+    """
+    model_signature_json = _infer_schema(
+        model_input=model_input, model=model, infer_method_name=infer_method_name
+    )
+    return autogen_artifacts.SklearnModelSchema(
+        infer_method_name=infer_method_name,
+        inputs=json.loads(model_signature_json["inputs"]),
+        outputs=json.loads(model_signature_json["outputs"]),
+    )
+
+
+def xgboost_infer_schema(
+    model_input: Any,
+    model: Union["Booster", "XGBModel"],
+    infer_method_name: str = "predict",
+) -> autogen_artifacts.XGBoostModelSchema:
+    """
+    Infer the schema of an XGBoost model.
+
+    Args:
+        model_input (Any): The input data to be used for schema inference.
+        model (Any): The XGBoost model instance.
+        infer_method_name (str): The name of the method to be used for schema inference.
+            Eg: predict (default), predict_proba
+    Returns:
+        XGBoostModelSchema: The inferred schema of the XGBoost model.
+    """
+    model_signature_json = _infer_schema(
+        model_input=model_input, model=model, infer_method_name=infer_method_name
+    )
+    return autogen_artifacts.XGBoostModelSchema(
+        infer_method_name=infer_method_name,
+        inputs=json.loads(model_signature_json["inputs"]),
+        outputs=json.loads(model_signature_json["outputs"]),
+    )

truefoundry/ml/validation_utils.py

@@ -29,6 +29,8 @@ _ML_REPO_NAME_REGEX = re.compile(r"^[a-zA-Z][a-zA-Z0-9\-]{1,98}[a-zA-Z0-9]$")
 _RUN_NAME_REGEX = re.compile(r"^[a-zA-Z0-9-]*$")
 _RUN_LOG_LOG_TYPE_REGEX = re.compile(r"^[a-zA-Z0-9-/]*$")
 _RUN_LOG_KEY_REGEX = re.compile(r"^[a-zA-Z0-9-_]*$")
+_APP_NAME_REGEX = re.compile(r"^[a-z][a-z0-9\\-]{1,30}[a-z0-9]$")
+
 
 MAX_PARAMS_TAGS_PER_BATCH = 100
 MAX_METRICS_PER_BATCH = 1000

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.5.0rc6
+Version: 0.5.1
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
@@ -19,7 +19,6 @@ Requires-Dist: aenum (>=3.0.0,<4.0.0)
 Requires-Dist: click (>=7.0.0,<9.0.0)
 Requires-Dist: coolname (>=1.1.0,<2.0.0)
 Requires-Dist: docker (>=6.1.2,<8.0.0)
-Requires-Dist: fastapi (>=0.56.0,<0.200.0)
 Requires-Dist: filelock (>=3.8.0,<4.0.0)
 Requires-Dist: flytekit (==1.13.13) ; extra == "workflow"
 Requires-Dist: gitignorefile (>=1.1.2,<2.0.0)
@@ -29,15 +28,13 @@ Requires-Dist: numpy (>=1.23.0,<2.0.0) ; python_version < "3.12"
 Requires-Dist: numpy (>=1.26.0,<2.0.0) ; python_version >= "3.12"
 Requires-Dist: openai (>=1.16.2,<2.0.0)
 Requires-Dist: packaging (>=20.0,<25.0)
-Requires-Dist: pandas (>=1.0.0,<3.0.0) ; python_version < "3.10"
-Requires-Dist: pandas (>=1.4.0,<3.0.0) ; python_version >= "3.10"
 Requires-Dist: pydantic (>=1.8.2,<3.0.0)
 Requires-Dist: pygments (>=2.12.0,<3.0.0)
 Requires-Dist: python-dateutil (>=2.8.2,<3.0.0)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: python-socketio[client] (>=5.5.2,<6.0.0)
 Requires-Dist: questionary (>=1.10.0,<2.0.0)
-Requires-Dist: requests (>=2.31.0,<3.0.0)
+Requires-Dist: requests (>=2.18.0,<3.0.0)
 Requires-Dist: requirements-parser (>=0.11.0,<0.12.0)
 Requires-Dist: rich (>=13.7.1,<14.0.0)
 Requires-Dist: rich-click (>=1.2.1,<2.0.0)
@@ -46,7 +43,6 @@ Requires-Dist: scipy (>=1.5.0,<2.0.0) ; python_version < "3.12"
 Requires-Dist: tqdm (>=4.0.0,<5.0.0)
 Requires-Dist: typing-extensions (>=4.0)
 Requires-Dist: urllib3 (>=1.26.18,<3)
-Requires-Dist: uvicorn (>=0.13.0,<1.0.0)
 Requires-Dist: yq (>=3.1.0,<4.0.0)
 Description-Content-Type: text/markdown