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

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,9 +16,8 @@ from typing import (
 )
 
 import coolname
-import pandas as pd
 
-from truefoundry.common.utils import relogin_error_message
+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,
@@ -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,13 +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,
         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
@@ -1263,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.
@@ -1387,14 +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,
             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,
 )
@@ -930,14 +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,
         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
@@ -965,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.
@@ -1084,14 +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,
             environment=environment,
-            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version