truefoundry 0.4.9rc1__py3-none-any.whl → 0.5.0__py3-none-any.whl

truefoundry/ml/mlfoundry_api.py

@@ -1,14 +1,23 @@
 import os
 import time
 import uuid
-from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import coolname
 import pandas as pd
 
 from truefoundry.common.utils import relogin_error_message
-from truefoundry.ml import constants
+from truefoundry.ml import ModelVersionEnvironment, constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
     ArtifactType,
@@ -38,11 +47,15 @@ from truefoundry.ml.internal_namespace import NAMESPACE
 from truefoundry.ml.log_types.artifacts.artifact import (
     ArtifactPath,
     ArtifactVersion,
+    BlobStorageDirectory,
     ChatPromptVersion,
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
-from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
+from truefoundry.ml.log_types.artifacts.model import (
+    ModelVersion,
+    _log_model_version,
+)
 from truefoundry.ml.logger import logger
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
 from truefoundry.ml.session import (
@@ -57,6 +70,9 @@ from truefoundry.ml.validation_utils import (
     _validate_run_name,
 )
 
+if TYPE_CHECKING:
+    from truefoundry.ml import ModelFrameworkType
+
 _SEARCH_MAX_RESULTS_DEFAULT = 1000
 
 _INTERNAL_ENV_VARS = [
@@ -1209,12 +1225,13 @@ class MlFoundry:
         *,
         ml_repo: str,
         name: str,
-        model_file_or_folder: str,
-        framework: Optional[Union[ModelFramework, str]],
-        additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
+        model_file_or_folder: Union[str, BlobStorageDirectory],
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         progress: Optional[bool] = None,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
+        environment: Optional[ModelVersionEnvironment] = None,
+        model_schema: Optional[Dict[str, Any]] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1226,45 +1243,24 @@ class MlFoundry:
             name (str): Name of the model. If a model with this name already exists under the current ML Repo,
                 the logged model will be added as a new version under that `name`. If no models exist with the given
                 `name`, the given model will be logged as version 1.
-            model_file_or_folder (str): Path to either a single file or a folder containing model files. This folder
-                is usually created using serialization methods of libraries or frameworks e.g. `joblib.dump`,
-                `model.save_pretrained(...)`, `torch.save(...)`, `model.save(...)`
-            framework (Union[enums.ModelFramework, str]): Model Framework. Ex:- pytorch, sklearn, tensorflow etc.
-                The full list of supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
-                Can also be `None` when `model` is `None`.
-            additional_files (Sequence[Tuple[Union[str, Path], Optional[str]]], optional): A list of pairs
-                of (source path, destination path) to add additional files and folders
-                to the model version contents. The first member of the pair should be a file or directory path
-                and the second member should be the path inside the model versions contents to upload to.
-                The model version contents are arranged like follows
-                .
-                └── model/
-                    └── # model files are serialized here
-                └── # any additional files and folders can be added here.
-
-                You can also add additional files to model/ subdirectory by specifying the destination path as model/
-
-                ```python
-                run.log_model(
-                    name="xyz",
-                    model_file_or_folder="clf.joblib",
-                    framework="sklearn",
-                    additional_files=[("foo.txt", "foo/bar/foo.txt"), ("tokenizer/", "foo/tokenizer/")]
-                )
-                ```
 
-                would result in
+            model_file_or_folder (Union[str, BlobStorageDirectory]):
+                str:
+                    Path to either a single file or a folder containing model files.
+                    This folder is typically created using serialization methods from libraries or frameworks,
+                    e.g., `joblib.dump`, `model.save_pretrained(...)`, `torch.save(...)`, or `model.save(...)`.
+                BlobStorageDirectory:
+                    uri (str): URI to the model file or folder in a storage integration associated with the specified ML Repo.
+                        The model files or folder must reside within the same storage integration as the specified ML Repo.
+                        Accepted URI formats include `s3://integration-bucket-name/prefix/path/to/model` or `gs://integration-bucket-name/prefix/path/to/model`.
+                        If the URI points to a model in a different storage integration, an error will be raised indicating "Invalid source URI."
+
+            framework (Optional[Union[ModelFramework, ModelFrameworkType]]): Framework used for model serialization.
+                Supported frameworks values (ModelFrameworkType) can be imported from `from truefoundry.ml import *`.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                Can also be `None` if the framework is not known or not supported.
+                **Deprecated**: Prefer `ModelFrameworkType` over `enums.ModelFramework`.
 
-                ```
-                .
-                ├── model/
-                │   └── clf.joblib # if `model_file_or_folder` is a folder, contents will be added here
-                └── foo/
-                    ├── bar/
-                    │   └── foo.txt
-                    └── tokenizer/
-                        └── # contents of tokenizer/ directory will be uploaded here
-                ```
             description (Optional[str], optional): arbitrary text upto 1024 characters to store as description.
                 This field can be updated at any time after logging. Defaults to `None`
             metadata (Optional[Dict[str, Any]], optional): arbitrary json serializable dictionary to store metadata.
@@ -1281,8 +1277,7 @@ class MlFoundry:
             ### Sklearn
 
             ```python
-            from truefoundry.ml import get_client
-            from truefoundry.ml.enums import ModelFramework
+            from truefoundry.ml import get_client, SklearnFramework
 
             import joblib
             import numpy as np
@@ -1307,7 +1302,7 @@ class MlFoundry:
                 ml_repo="my-classification-project",
                 name="my-sklearn-model",
                 model_file_or_folder="sklearn-pipeline.joblib",
-                framework=ModelFramework.SKLEARN,
+                framework=SklearnFramework(),
                 metadata={"accuracy": 0.99, "f1": 0.80},
                 step=1,  # step number, useful when using iterative algorithms like SGD
             )
@@ -1317,8 +1312,7 @@ class MlFoundry:
             ### Huggingface Transformers
 
             ```python
-            from truefoundry.ml import get_client
-            from truefoundry.ml.enums import ModelFramework
+            from truefoundry.ml import get_client, TransformersFramework, LibraryName
 
             import torch
             from transformers import AutoTokenizer, AutoConfig, pipeline, AutoModelForCausalLM
@@ -1342,7 +1336,7 @@ class MlFoundry:
                 ml_repo="my-llm-project",
                 name="my-transformers-model",
                 model_file_or_folder="my-transformers-model/",
-                framework=ModelFramework.TRANSFORMERS
+                framework=TransformersFramework(library_name=LibraryName.TRANSFORMERS, pipeline_tag='text-generation')
             )
             print(model_version.fqn)
             ```
@@ -1356,12 +1350,13 @@ class MlFoundry:
             ml_repo_id=ml_repo_id,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            framework=framework,
-            additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=None,
             progress=progress,
+            framework=framework,
+            environment=environment,
+            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/mlfoundry_run.py

@@ -3,7 +3,6 @@ import os
 import platform
 import re
 import time
-from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -12,7 +11,6 @@ from typing import (
     Iterator,
     List,
     Optional,
-    Sequence,
     Tuple,
     Union,
 )
@@ -20,7 +18,7 @@ from urllib.parse import urljoin, urlsplit
 
 from truefoundry import version
 from truefoundry.common.utils import relogin_error_message
-from truefoundry.ml import constants, enums
+from truefoundry.ml import ModelVersionEnvironment, constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     DeleteRunRequest,
@@ -40,13 +38,20 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     UpdateRunRequestDto,
 )
 from truefoundry.ml.entities import Metric
-from truefoundry.ml.enums import RunStatus
+from truefoundry.ml.enums import ModelFramework, RunStatus
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
 from truefoundry.ml.log_types import Image, Plot
-from truefoundry.ml.log_types.artifacts.artifact import ArtifactPath, ArtifactVersion
+from truefoundry.ml.log_types.artifacts.artifact import (
+    ArtifactPath,
+    ArtifactVersion,
+    BlobStorageDirectory,
+)
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
-from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
+from truefoundry.ml.log_types.artifacts.model import (
+    ModelVersion,
+    _log_model_version,
+)
 from truefoundry.ml.logger import logger
 from truefoundry.ml.run_utils import ParamsType, flatten_dict, process_params
 from truefoundry.ml.session import ACTIVE_RUNS, _get_api_client, get_active_session
@@ -61,6 +66,8 @@ if TYPE_CHECKING:
     import matplotlib
     import plotly
 
+    from truefoundry.ml import ModelFrameworkType
+
 
 def _ensure_not_deleted(method):
     @functools.wraps(method)
@@ -920,13 +927,14 @@ class MlFoundryRun:
         self,
         *,
         name: str,
-        model_file_or_folder: str,
-        framework: Optional[Union[enums.ModelFramework, str]],
-        additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
+        model_file_or_folder: Union[str, BlobStorageDirectory],
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         step: int = 0,
         progress: Optional[bool] = None,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
+        environment: Optional[ModelVersionEnvironment] = None,
+        model_schema: Optional[Dict[str, Any]] = None,
     ) -> ModelVersion:
         # TODO (chiragjn): Document mapping of framework to list of valid model save kwargs
         # TODO (chiragjn): Add more examples
@@ -939,39 +947,22 @@ class MlFoundryRun:
             name (str): Name of the model. If a model with this name already exists under the current ML Repo,
                 the logged model will be added as a new version under that `name`. If no models exist with the given
                 `name`, the given model will be logged as version 1.
-            model_file_or_folder (str): Path to either a single file or a folder containing model files. This folder
-                is usually created using serialization methods of libraries or frameworks e.g. `joblib.dump`,
-                `model.save_pretrained(...)`, `torch.save(...)`, `model.save(...)`
-            framework (Union[enums.ModelFramework, str]): Model Framework. Ex:- pytorch, sklearn, tensorflow etc.
-                The full list of supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
-                Can also be `None` when `model` is `None`.
-            additional_files (Sequence[Tuple[Union[str, Path], Optional[str]]], optional): A list of pairs
-                of (source path, destination path) to add additional files and folders
-                to the model version contents. The first member of the pair should be a file or directory path
-                and the second member should be the path inside the model versions contents to upload to.
-                The model version contents are arranged like follows
-                .
-                └── model/
-                    └── # model files are serialized here
-                └── # any additional files and folders can be added here.
+            model_file_or_folder (Union[str, BlobStorageDirectory]):
+                str:
+                    Path to either a single file or a folder containing model files.
+                    This folder is typically created using serialization methods from libraries or frameworks,
+                    e.g., `joblib.dump`, `model.save_pretrained(...)`, `torch.save(...)`, or `model.save(...)`.
+                BlobStorageDirectory:
+                    uri (str): URI to the model file or folder in a storage integration associated with the specified ML Repo.
+                        The model files or folder must reside within the same storage integration as the specified ML Repo.
+                        Accepted URI formats include `s3://integration-bucket-name/prefix/path/to/model` or `gs://integration-bucket-name/prefix/path/to/model`.
+                        If the URI points to a model in a different storage integration, an error will be raised indicating "Invalid source URI."
+            framework (Optional[Union[ModelFramework, ModelFrameworkType]]): Framework used for model serialization.
+                Supported frameworks values (ModelFrameworkType) can be imported from `from truefoundry.ml import *`.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                Can also be `None` if the framework is not known or not supported.
+                **Deprecated**: Prefer `ModelFrameworkType` over `enums.ModelFramework`.
 
-                You can also add additional files to model/ subdirectory by specifying the destination path as model/
-
-                ```
-                E.g. >>> run.log_model(
-                     ...     name="xyz", model_file_or_folder="clf.joblib", framework="sklearn",
-                     ...     additional_files=[("foo.txt", "foo/bar/foo.txt"), ("tokenizer/", "foo/tokenizer/")]
-                     ... )
-                would result in
-                .
-                ├── model/
-                │   └── clf.joblib  # if `model_file_or_folder` is a folder, contents will be added here
-                └── foo/
-                    ├── bar/
-                    │   └── foo.txt
-                    └── tokenizer/
-                        └── # contents of tokenizer/ directory will be uploaded here
-                ```
             description (Optional[str], optional): arbitrary text upto 1024 characters to store as description.
                 This field can be updated at any time after logging. Defaults to `None`
             metadata (Optional[Dict[str, Any]], optional): arbitrary json serializable dictionary to store metadata.
@@ -1064,12 +1055,13 @@ class MlFoundryRun:
             run=self,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            framework=framework,
-            additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=step,
             progress=progress,
+            framework=framework,
+            environment=environment,
+            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/model_framework.py

@@ -0,0 +1,169 @@
+import warnings
+from typing import Any, Dict, Literal, Optional, Union, get_args
+
+from truefoundry.ml import ModelFramework
+from truefoundry.ml.autogen.entities import artifacts as autogen_artifacts
+from truefoundry.pydantic_v1 import BaseModel, Field
+
+
+class FastAIFramework(autogen_artifacts.FastAIFramework):
+    """FastAI model Framework"""
+
+    type: Literal["fastai"] = "fastai"
+
+
+class GluonFramework(autogen_artifacts.GluonFramework):
+    """Gluon model Framework"""
+
+    type: Literal["gluon"] = "gluon"
+
+
+class H2OFramework(autogen_artifacts.H2OFramework):
+    """H2O model Framework"""
+
+    type: Literal["h2o"] = "h2o"
+
+
+class KerasFramework(autogen_artifacts.KerasFramework):
+    """Keras model Framework"""
+
+    type: Literal["keras"] = "keras"
+
+
+class LightGBMFramework(autogen_artifacts.LightGBMFramework):
+    """LightGBM model Framework"""
+
+    type: Literal["lightgbm"] = "lightgbm"
+
+
+class ONNXFramework(autogen_artifacts.ONNXFramework):
+    """ONNX model Framework"""
+
+    type: Literal["onnx"] = "onnx"
+
+
+class PaddleFramework(autogen_artifacts.PaddleFramework):
+    """Paddle model Framework"""
+
+    type: Literal["paddle"] = "paddle"
+
+
+class PyTorchFramework(autogen_artifacts.PyTorchFramework):
+    """PyTorch model Framework"""
+
+    type: Literal["pytorch"] = "pytorch"
+
+
+class SklearnFramework(autogen_artifacts.SklearnFramework):
+    """Sklearn model Framework"""
+
+    type: Literal["sklearn"] = "sklearn"
+
+
+class SpaCyFramework(autogen_artifacts.SpaCyFramework):
+    """SpaCy model Framework"""
+
+    type: Literal["spacy"] = "spacy"
+
+
+class StatsModelsFramework(autogen_artifacts.StatsModelsFramework):
+    """StatsModels model Framework"""
+
+    type: Literal["statsmodels"] = "statsmodels"
+
+
+class TensorFlowFramework(autogen_artifacts.TensorFlowFramework):
+    """TensorFlow model Framework"""
+
+    type: Literal["tensorflow"] = "tensorflow"
+
+
+class TransformersFramework(autogen_artifacts.TransformersFramework):
+    """Transformers model Framework"""
+
+    type: Literal["transformers"] = "transformers"
+
+
+class XGBoostFramework(autogen_artifacts.XGBoostFramework):
+    """XGBoost model Framework"""
+
+    type: Literal["xgboost"] = "xgboost"
+
+
+# Union of all the model frameworks
+
+
+ModelFrameworkType = Union[
+    FastAIFramework,
+    GluonFramework,
+    H2OFramework,
+    KerasFramework,
+    LightGBMFramework,
+    ONNXFramework,
+    PaddleFramework,
+    PyTorchFramework,
+    SklearnFramework,
+    SpaCyFramework,
+    StatsModelsFramework,
+    TensorFlowFramework,
+    TransformersFramework,
+    XGBoostFramework,
+]
+
+
+class _ModelFramework(BaseModel):
+    __root__: ModelFrameworkType = Field(discriminator="type")
+
+    @classmethod
+    def to_model_framework_type(
+        cls,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
+    ) -> Optional["ModelFrameworkType"]:
+        """
+        Converts a ModelFramework or string representation to a ModelFrameworkType object.
+
+        Args:
+            framework (Optional[Union[str, ModelFramework, ModelFrameworkType]]): ModelFrameworkType or equivalent input.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                May be `None` if the framework is unknown or unsupported.
+                **Deprecated**: Prefer passing a `ModelFrameworkType` instance.
+
+        Returns:
+            ModelFrameworkType corresponding to the input, or None if the input is None.
+        """
+        if framework is None:
+            return None
+
+        # Issue a deprecation warning for str and ModelFramework types
+        if isinstance(framework, (str, ModelFramework)):
+            warnings.warn(
+                "Passing a string or ModelFramework Enum is deprecated. Please use a ModelFrameworkType object.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        # Convert string to ModelFramework
+        if isinstance(framework, str):
+            framework = ModelFramework(framework)
+
+        # Convert ModelFramework to ModelFrameworkType
+        if isinstance(framework, ModelFramework):
+            if framework == ModelFramework.UNKNOWN:
+                return None
+            return cls.parse_obj({"type": framework.value}).__root__
+
+        # Directly return if already a ModelFrameworkType
+        if isinstance(framework, get_args(ModelFrameworkType)):
+            return framework
+
+        raise ValueError(
+            "framework must be a string, ModelFramework enum, or ModelFrameworkType object"
+        )
+
+    @classmethod
+    def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[ModelFrameworkType]:
+        """Create an instance of ModelFramework from a dict"""
+        if obj is None:
+            return None
+
+        return cls.parse_obj(obj).__root__

{truefoundry-0.4.9rc1.dist-info → truefoundry-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.4.9rc1
+Version: 0.5.0
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com