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

truefoundry/ml/log_types/artifacts/model.py

@@ -8,7 +8,7 @@ import typing
 import uuid
 import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
@@ -25,11 +25,13 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     MlfoundryArtifactsApi,
     ModelDto,
     ModelVersionDto,
+    ModelVersionEnvironment,
     ModelVersionManifest,
     NotifyArtifactVersionFailureDto,
     TrueFoundryArtifactSource,
     UpdateModelVersionRequestDto,
 )
+from truefoundry.ml.autogen.models import infer_signature as _infer_signature
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.artifact import BlobStorageDirectory
@@ -43,11 +45,19 @@ from truefoundry.ml.log_types.artifacts.utils import (
     _validate_description,
     calculate_total_size,
 )
-from truefoundry.ml.model_framework import ModelFrameworkType, _ModelFramework
+from truefoundry.ml.model_framework import (
+    ModelFrameworkType,
+    _ModelFramework,
+    auto_update_environment_details,
+    auto_update_model_framework_details,
+)
 from truefoundry.ml.session import _get_api_client
 from truefoundry.pydantic_v1 import BaseModel, Extra
 
 if TYPE_CHECKING:
+    import numpy as np
+    import pandas as pd
+
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
 
@@ -96,7 +106,7 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._model_schema: Optional[Dict[str, Any]] = None
+        self._environment: Optional[ModelVersionEnvironment] = None
         self._framework: Optional[ModelFrameworkType] = None
         self._set_mutable_attrs()
 
@@ -139,22 +149,19 @@ class ModelVersion:
         if self._model_version.manifest:
             self._description = self._model_version.manifest.description or ""
             self._metadata = copy.deepcopy(self._model_version.manifest.metadata)
-            self._model_schema = copy.deepcopy(
-                self._model_version.manifest.model_schema
+            self._environment = copy.deepcopy(self._model_version.manifest.environment)
+            self._framework = (
+                copy.deepcopy(self._model_version.manifest.framework.actual_instance)
+                if self._model_version.manifest.framework
+                else None
             )
-            if self._model_version.manifest.framework:
-                self._framework = copy.deepcopy(
-                    self._model_version.manifest.framework.actual_instance
-                )
-            else:
-                self._framework = None
         else:
             self._description = self._model_version.description or ""
             self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
+            self._environment = None
             self._framework = _ModelFramework.to_model_framework_type(
                 self._model_version.model_framework
             )
-            self._model_schema = None
 
     def _refetch_model_version(self, reset_mutable_attrs: bool = True):
         _model_version = self._mlfoundry_artifacts_api.get_model_version_get(
@@ -225,21 +232,21 @@ class ModelVersion:
         self._metadata = copy.deepcopy(value)
 
     @property
-    def model_schema(self) -> Optional[Dict[str, Any]]:
-        """Get model_schema for the current model"""
-        return self._model_schema
+    def environment(self) -> Optional[ModelVersionEnvironment]:
+        """Get the environment details for the model"""
+        return self._environment
 
-    @model_schema.setter
-    def model_schema(self, value: Optional[Dict[str, Any]]):
-        """set the model_schema for current model"""
+    @environment.setter
+    def environment(self, value: Optional[Dict[str, Any]]):
+        """set the environment details for the model"""
         if not self._model_version.manifest:
             warnings.warn(
-                message="This model version was created using an older serialization format. model_schema will not be updated",
+                message="This model version was created using an older serialization format. Environment will not be updated",
                 category=DeprecationWarning,
                 stacklevel=2,
             )
             return
-        self._model_schema = copy.deepcopy(value)
+        self._environment = copy.deepcopy(value)
 
     @property
     def framework(self) -> Optional["ModelFrameworkType"]:
@@ -438,7 +445,7 @@ class ModelVersion:
         if self._model_version.manifest:
             self._model_version.manifest.description = self.description
             self._model_version.manifest.metadata = self.metadata
-            self._model_version.manifest.model_schema = self.model_schema
+            self._model_version.manifest.environment = self.environment
             self._model_version.manifest.framework = (
                 Framework.from_dict(self.framework.dict()) if self.framework else None
             )
@@ -466,13 +473,12 @@ def _log_model_version(  # noqa: C901
     model_file_or_folder: Union[str, BlobStorageDirectory],
     mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
     ml_repo_id: Optional[str] = None,
-    additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
     description: Optional[str] = None,
     metadata: Optional[Dict[str, Any]] = None,
     step: Optional[int] = 0,
     progress: Optional[bool] = None,
     framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
-    model_schema: Optional[Dict[str, Any]] = None,
+    environment: Optional[ModelVersionEnvironment] = None,
 ) -> ModelVersion:
     if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
         raise MlFoundryException(
@@ -496,7 +502,6 @@ def _log_model_version(  # noqa: C901
     step = step or 0
     total_size = None
     metadata = metadata or {}
-    additional_files = additional_files or {}
 
     _validate_description(description)
     _validate_artifact_metadata(metadata)
@@ -514,17 +519,6 @@ def _log_model_version(  # noqa: C901
                 ignore_model_dir_dest_conflict=True,
             )
 
-            # verify additional files and paths, copy additional files
-            if additional_files:
-                logger.info("Adding `additional_files` to model version contents")
-                temp_dest_to_src_map = _copy_additional_files(
-                    root_dir=temp_dir.name,
-                    files_dir="",
-                    model_dir=None,
-                    additional_files=additional_files,
-                    ignore_model_dir_dest_conflict=False,
-                    existing_dest_to_src_map=temp_dest_to_src_map,
-                )
         except Exception as e:
             temp_dir.cleanup()
             raise MlFoundryException("Failed to log model") from e
@@ -578,15 +572,24 @@ def _log_model_version(  # noqa: C901
     else:
         raise MlFoundryException("Invalid model_file_or_folder provided")
 
-    _framework = _ModelFramework.to_model_framework_type(framework)
     _source_cls = typing.get_type_hints(ModelVersionManifest)["source"]
+
+    # Auto fetch the framework & environment details if not provided
+    framework = _ModelFramework.to_model_framework_type(framework)
+    if framework and isinstance(model_file_or_folder, str):
+        auto_update_model_framework_details(
+            framework=framework, model_file_or_folder=model_file_or_folder
+        )
+    environment = environment or ModelVersionEnvironment()
+    auto_update_environment_details(environment=environment, framework=framework)
+
     model_manifest = ModelVersionManifest(
         description=description,
         metadata=metadata,
         source=_source_cls.from_dict(source.dict()),
-        framework=Framework.from_dict(_framework.dict()) if _framework else None,
+        framework=Framework.from_dict(framework.dict()) if framework else None,
+        environment=environment,
         step=step,
-        model_schema=model_schema,
     )
     artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
@@ -600,3 +603,15 @@ def _log_model_version(  # noqa: C901
         )
     )
     return ModelVersion.from_fqn(fqn=artifact_version_response.artifact_version.fqn)
+
+
+def infer_signature(
+    model_input: Any = None,
+    model_output: Optional[
+        Union["pd.DataFrame", "np.ndarray", Dict[str, "np.ndarray"]]
+    ] = None,
+    params: Optional[Dict[str, Any]] = None,
+):
+    return _infer_signature(
+        model_input=model_input, model_output=model_output, params=params
+    )

truefoundry/ml/log_types/artifacts/utils.py

@@ -2,7 +2,7 @@ import json
 import logging
 import os
 import posixpath
-from pathlib import Path
+from pathlib import Path, PureWindowsPath
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
 from truefoundry.ml.exceptions import MlFoundryException
@@ -11,6 +11,13 @@ from truefoundry.ml.log_types.artifacts.constants import DESCRIPTION_MAX_LENGTH
 logger = logging.getLogger(__name__)
 
 
+def to_unix_path(path):
+    path = os.path.normpath(path)
+    if os.path.sep == "\\":
+        path = PureWindowsPath(path).as_posix()
+    return path
+
+
 def _copy_tree(
     root_dir: str, src_path: str, dest_path: str, dest_to_src: Dict[str, str]
 ):
@@ -42,6 +49,35 @@ def is_destination_path_dirlike(dest_path) -> bool:
     return False
 
 
+def get_single_file_path_if_only_one_in_directory(path: str) -> Optional[str]:
+    """
+    Get the filename from a path, or return a filename from a directory if a single file exists.
+    Args:
+        path: The file or folder path.
+    Returns:
+        Optional[str]: The filename or None if no files are found or multiple files are found.
+    """
+    # If it's already a file, return it as-is
+    if os.path.isfile(path):
+        return path
+
+    # If it's a directory, check if it contains a single file
+    if is_destination_path_dirlike(path):
+        all_files = []
+        for root, _, files in os.walk(path):
+            # Collect all files found in any subdirectory
+            all_files.extend(os.path.join(root, f) for f in files)
+            # If more than one file is found, stop early
+            if len(all_files) > 1:
+                return None
+
+        # If only one file is found, return it
+        if len(all_files) == 1:
+            return all_files[0]
+
+    return None  # No file found or Multiple files found
+
+
 def _copy_additional_files(
     root_dir: str,
     files_dir: str,  # relative to root dir e.g. "files/"
@@ -136,7 +172,7 @@ def _get_src_dest_pairs(
     dest_to_src_map: Dict[str, str],
 ) -> Sequence[Tuple[str, str]]:
     src_dest_pairs = [
-        (src_path, os.path.relpath(dest_abs_path, root_dir))
+        (src_path, to_unix_path(os.path.relpath(dest_abs_path, root_dir)))
         for dest_abs_path, src_path in dest_to_src_map.items()
     ]
     return src_dest_pairs

truefoundry/ml/mlfoundry_api.py

@@ -1,7 +1,8 @@
 import os
 import time
 import uuid
-from pathlib import Path
+import zipfile
+from io import BytesIO
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -15,10 +16,9 @@ from typing import (
 )
 
 import coolname
-import pandas as pd
 
-from truefoundry.common.utils import relogin_error_message
-from truefoundry.ml import constants
+from truefoundry.common.utils import ContextualDirectoryManager, relogin_error_message
+from truefoundry.ml import ModelVersionEnvironment, constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
     ArtifactType,
@@ -27,12 +27,14 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     CreateRunRequestDto,
     DatasetDto,
     ExperimentsApi,
+    ExportDeploymentFilesRequestDto,
     ListArtifactsRequestDto,
     ListArtifactVersionsRequestDto,
     ListDatasetsRequestDto,
     ListModelVersionsRequestDto,
     MlfoundryArtifactsApi,
     ModelDto,
+    ModelServer,
     RunsApi,
     RunTagDto,
     SearchRunsRequestDto,
@@ -75,7 +77,7 @@ if TYPE_CHECKING:
     from truefoundry.ml import ModelFrameworkType
 
 _SEARCH_MAX_RESULTS_DEFAULT = 1000
-
+_ML_FOUNDRY_API_REQUEST_TIMEOUT = 10
 _INTERNAL_ENV_VARS = [
     "TFY_INTERNAL_APPLICATION_ID",
     "TFY_INTERNAL_JOB_RUN_NAME",
@@ -476,43 +478,6 @@ class MlFoundry:
         )
         return mlfoundry_run
 
-    def get_all_runs(
-        self,
-        ml_repo: str,
-    ) -> pd.DataFrame:
-        """Returns all the run name and id present under a ML Repo.
-
-        The user must have `READ` access to the ML Repo.
-
-        Args:
-            ml_repo (str): Name of the ML Repo.
-        Returns:
-            pd.DataFrame: dataframe with two columns- run_id and run_name
-
-        Examples:
-
-            ### get all the runs from a ml_repo
-            ```python
-            from truefoundry.ml import get_client
-
-            client = get_client()
-
-            run = client.get_all_runs(ml_repo='my-repo')
-            ```
-        """
-        runs = []
-        for run in self.search_runs(ml_repo=ml_repo):
-            runs.append((run.run_id, run.run_name))
-
-        if len(runs) == 0:
-            return pd.DataFrame(
-                columns=[constants.RUN_ID_COL_NAME, constants.RUN_NAME_COL_NAME]
-            )
-
-        return pd.DataFrame(
-            runs, columns=[constants.RUN_ID_COL_NAME, constants.RUN_NAME_COL_NAME]
-        )
-
     def search_runs(
         self,
         ml_repo: str,
@@ -658,6 +623,74 @@ class MlFoundry:
         """
         return self._tracking_uri
 
+    def _initialize_model_server(
+        self,
+        name: str,
+        model_version_fqn: str,
+        workspace_fqn: str,
+        model_server: ModelServer,
+        output_dir: Optional[str] = None,
+    ) -> str:
+        """
+        Initialize the model server for deployment.
+
+        Args:
+            name (str): Name of the application.
+            model_version_fqn (str): Fully Qualified Name of the model version.
+            workspace_fqn (str): Fully Qualified Name of the workspace.
+            model_server (ModelServer): Server type for deployment (e.g., TRITON).
+            output_dir (Optional[str]): Directory where the model files will be extracted.
+                                        Defaults to the current working directory.
+
+        Returns:
+            str: Path to the directory where the model files are extracted.
+
+        Raises:
+            MlFoundryException: If an error occurs during API calls, directory creation, or file extraction.
+        """
+
+        # Using get_with_http_info to get the response object and extract the raw data
+        try:
+            export_deployment_files_request_dto = ExportDeploymentFilesRequestDto(
+                model_version_fqn=model_version_fqn,
+                workspace_fqn=workspace_fqn,
+                service_name=name,
+                model_server=model_server,
+            )
+            response = self._mlfoundry_artifacts_api.export_deployment_files_by_fqn_post_with_http_info(
+                export_deployment_files_request_dto=export_deployment_files_request_dto,
+                _preload_content=False,
+                _request_timeout=_ML_FOUNDRY_API_REQUEST_TIMEOUT,
+            )
+        except ApiException as e:
+            err_msg = (
+                f"Failed to fetch deployment files for name={name}, "
+                f"model_version_fqn={model_version_fqn}, workspace_fqn={workspace_fqn}. "
+                f"Error: {e}"
+            )
+            raise MlFoundryException(err_msg) from e
+
+        output_dir = os.path.abspath(output_dir or os.getcwd())
+        codegen_dir = os.path.join(output_dir, name)
+        if codegen_dir == output_dir:
+            raise ValueError("Name cannot be empty, please provide a valid name")
+
+        try:
+            with ContextualDirectoryManager(dir_path=codegen_dir) as dir_path:
+                with zipfile.ZipFile(BytesIO(response.raw_data), mode="r") as zip_file:
+                    zip_file.extractall(dir_path)
+        except FileExistsError as e:
+            err_msg = (
+                f"Deployment directory {codegen_dir!r} already exists. "
+                "Please choose a different deployment name or delete the existing directory."
+            )
+            raise MlFoundryException(err_msg) from e
+        except zipfile.BadZipFile as e:
+            raise MlFoundryException(
+                f"Failed to extract model files. Error: {e}"
+            ) from e
+        return codegen_dir
+
     def get_model_version(
         self,
         ml_repo: str,
@@ -1227,12 +1260,11 @@ class MlFoundry:
         ml_repo: str,
         name: str,
         model_file_or_folder: Union[str, BlobStorageDirectory],
-        additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         progress: Optional[bool] = None,
         framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
-        model_schema: Optional[Dict[str, Any]] = None,
+        environment: Optional[ModelVersionEnvironment] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1262,41 +1294,6 @@ class MlFoundry:
                 Can also be `None` if the framework is not known or not supported.
                 **Deprecated**: Prefer `ModelFrameworkType` over `enums.ModelFramework`.
 
-            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
-                from truefoundry.ml import TensorFlowFramework
-
-                run.log_model(
-                    name="xyz",
-                    model_file_or_folder="clf.joblib",
-                    framework=TensorFlowFramework(),
-                    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.
@@ -1386,13 +1383,12 @@ class MlFoundry:
             ml_repo_id=ml_repo_id,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=None,
             progress=progress,
             framework=framework,
-            model_schema=model_schema,
+            environment=environment,
         )
         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
+from truefoundry.ml import ModelVersionEnvironment, constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     DeleteRunRequest,
@@ -930,13 +928,12 @@ class MlFoundryRun:
         *,
         name: str,
         model_file_or_folder: Union[str, BlobStorageDirectory],
-        additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         step: int = 0,
         progress: Optional[bool] = None,
         framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
-        model_schema: Optional[Dict[str, Any]] = None,
+        environment: Optional[ModelVersionEnvironment] = None,
     ) -> ModelVersion:
         # TODO (chiragjn): Document mapping of framework to list of valid model save kwargs
         # TODO (chiragjn): Add more examples
@@ -964,33 +961,7 @@ class MlFoundryRun:
                 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`.
-            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/
 
-                ```
-                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.
@@ -1083,13 +1054,12 @@ class MlFoundryRun:
             run=self,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=step,
             progress=progress,
             framework=framework,
-            model_schema=model_schema,
+            environment=environment,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version