truefoundry 0.5.0rc2__py3-none-any.whl → 0.5.0rc4__py3-none-any.whl

truefoundry/ml/log_types/artifacts/model.py

@@ -4,6 +4,7 @@ import json
 import logging
 import os.path
 import tempfile
+import typing
 import uuid
 import warnings
 from pathlib import Path
@@ -15,25 +16,23 @@ from truefoundry.ml.artifact.truefoundry_artifact_repo import (
 )
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
-    ArtifactVersionSerializationFormat,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
-    ExternalModelSource,
+    ExternalArtifactSource,
     FinalizeArtifactVersionRequestDto,
     Framework,
+    Manifest,
     MlfoundryArtifactsApi,
     ModelDto,
     ModelVersionDto,
     ModelVersionManifest,
     NotifyArtifactVersionFailureDto,
-    TruefoundryModelSource,
+    TrueFoundryArtifactSource,
     UpdateModelVersionRequestDto,
 )
-from truefoundry.ml.autogen.client import (
-    Source as ModelVersionSource,
-)
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.exceptions import MlFoundryException
+from truefoundry.ml.log_types.artifacts.artifact import BlobStorageDirectory
 from truefoundry.ml.log_types.artifacts.constants import (
     INTERNAL_METADATA_PATH,
 )
@@ -46,7 +45,7 @@ from truefoundry.ml.log_types.artifacts.utils import (
 )
 from truefoundry.ml.model_framework import ModelFrameworkType, _ModelFramework
 from truefoundry.ml.session import _get_api_client
-from truefoundry.pydantic_v1 import BaseModel, Extra, StrictStr
+from truefoundry.pydantic_v1 import BaseModel, Extra
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
@@ -75,10 +74,6 @@ class ModelVersionInternalMetadata(BaseModel):
         return dct
 
 
-class BlobStorageModelDirectory(BaseModel):
-    uri: StrictStr
-
-
 class ModelVersionDownloadInfo(BaseModel):
     download_dir: str
     model_dir: str
@@ -101,7 +96,7 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._framework: Optional["Framework"] = None
+        self._framework: Optional[ModelFrameworkType] = None
         self._set_mutable_attrs()
 
     @classmethod
@@ -140,20 +135,29 @@ class ModelVersion:
             )
 
     def _set_mutable_attrs(self):
-        self._description = self._model_version.description or ""
-        self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
-        self._framework = (
-            copy.deepcopy(self._model_version.manifest.framework)
-            if self._model_version.manifest
-            else None
-        )
+        if self._model_version.manifest:
+            self._description = self._model_version.manifest.description or ""
+            self._metadata = copy.deepcopy(self._model_version.manifest.metadata)
+            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._framework = _ModelFramework.to_model_framework_type(
+                self._model_version.model_framework
+            )
 
-    def _refetch_model_version(self):
+    def _refetch_model_version(self, reset_mutable_attrs: bool = True):
         _model_version = self._mlfoundry_artifacts_api.get_model_version_get(
             id=self._model_version.id
         )
         self._model_version = _model_version.model_version
-        self._set_mutable_attrs()
+        if reset_mutable_attrs:
+            self._set_mutable_attrs()
 
     def __repr__(self):
         return f"{self.__class__.__name__}(fqn={self.fqn!r})"
@@ -187,12 +191,14 @@ class ModelVersion:
         return self._model_version.fqn
 
     @property
-    def step(self) -> int:
+    def step(self) -> Optional[int]:
         """Get the step in which model was created"""
+        if self._model_version.manifest:
+            return self._model_version.manifest.step
         return self._model_version.step
 
     @property
-    def description(self) -> Optional[str]:
+    def description(self) -> str:
         """Get description of the model"""
         return self._description
 
@@ -201,8 +207,6 @@ class ModelVersion:
         """set the description of the model"""
         _validate_description(value)
         self._description = value
-        if self._model_version.manifest:
-            self._model_version.manifest.description = value
 
     @property
     def metadata(self) -> Dict[str, Any]:
@@ -213,52 +217,34 @@ class ModelVersion:
     def metadata(self, value: Dict[str, Any]):
         """set the metadata for current model"""
         _validate_artifact_metadata(value)
-        self._metadata = value
-        if self._model_version.manifest:
-            self._model_version.manifest.metadata = value
+        self._metadata = copy.deepcopy(value)
 
     @property
     def framework(self) -> Optional["ModelFrameworkType"]:
         """Get the framework of the model"""
-        return (
-            _ModelFramework.from_dict(self._framework.actual_instance.dict())
-            if self._framework and self._framework.actual_instance
-            else None
-        )
+        return self._framework
 
     @framework.setter
     def framework(
         self, value: Optional[Union[str, ModelFramework, "ModelFrameworkType"]]
     ):
-        """set the framework of the model"""
-        model_framework_type: Optional[ModelFrameworkType] = None
-
-        if value is not None:
-            model_framework_type = _ModelFramework.to_model_framework_type(value)
-
-        if (
-            self._model_version.serialization_format
-            == ArtifactVersionSerializationFormat.V2
-        ):
-            self._model_version.manifest.framework = self._framework = (
-                Framework.from_dict(model_framework_type.dict())
-                if model_framework_type
-                else None
-            )
-        else:
+        """Set the framework of the model"""
+        if not self._model_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. Framework will not be updated",
                 category=DeprecationWarning,
                 stacklevel=2,
             )
+            return
+        self._framework = _ModelFramework.to_model_framework_type(value)
 
     @property
-    def created_at(self) -> datetime.datetime:
+    def created_at(self) -> Optional[datetime.datetime]:
         """Get the time at which model version was created"""
         return self._model_version.created_at
 
     @property
-    def updated_at(self) -> datetime.datetime:
+    def updated_at(self) -> Optional[datetime.datetime]:
         """Get the information about when the model version was updated"""
         return self._model_version.updated_at
 
@@ -311,13 +297,11 @@ class ModelVersion:
             path=path, overwrite=overwrite, progress=progress
         )
 
-        if (
-            self._model_version.serialization_format
-            == ArtifactVersionSerializationFormat.V2
-        ):
+        if self._model_version.manifest:
+            _framework = self._model_version.manifest.framework
             model_framework = (
-                ModelFramework(self.framework.type)
-                if self.framework
+                ModelFramework(_framework.actual_instance.type)
+                if _framework
                 else ModelFramework.UNKNOWN
             )
             download_info = ModelVersionDownloadInfo(
@@ -429,22 +413,34 @@ class ModelVersion:
             ```
         """
         self._ensure_not_deleted()
-        _model_version = self._mlfoundry_artifacts_api.update_model_version_post(
-            update_model_version_request_dto=UpdateModelVersionRequestDto(
-                id=self._model_version.id,
-                description=self.description,
-                artifact_metadata=self.metadata,
-                manifest=self._model_version.manifest,
+        if self._model_version.manifest:
+            self._model_version.manifest.description = self.description
+            self._model_version.manifest.metadata = self.metadata
+            self._model_version.manifest.framework = (
+                Framework.from_dict(self.framework.dict()) if self.framework else None
             )
-        )
-        self._model_version = _model_version.model_version
-        self._set_mutable_attrs()
+        try:
+            _model_version = self._mlfoundry_artifacts_api.update_model_version_post(
+                update_model_version_request_dto=UpdateModelVersionRequestDto(
+                    id=self._model_version.id,
+                    description=self.description,
+                    artifact_metadata=self.metadata,
+                    manifest=self._model_version.manifest,
+                )
+            )
+        except Exception:
+            # rollback edits to internal object
+            self._refetch_model_version(reset_mutable_attrs=False)
+            raise
+        else:
+            self._model_version = _model_version.model_version
+            self._set_mutable_attrs()
 
 
 def _log_model_version(  # noqa: C901
     run: Optional["MlFoundryRun"],
     name: str,
-    model_file_or_folder: Union[str, BlobStorageModelDirectory],
+    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]]] = (),
@@ -459,9 +455,9 @@ def _log_model_version(  # noqa: C901
             "Exactly one of run, mlfoundry_artifacts_api should be passed"
         )
 
-    if not isinstance(model_file_or_folder, (str, BlobStorageModelDirectory)):
+    if not isinstance(model_file_or_folder, (str, BlobStorageDirectory)):
         raise MlFoundryException(
-            "model_file_or_folder should be of type str or BlobStorageModelDirectory"
+            "model_file_or_folder should be of type str or BlobStorageDirectory"
         )
 
     if mlfoundry_artifacts_api and not ml_repo_id:
@@ -522,11 +518,9 @@ def _log_model_version(  # noqa: C901
     version_id = _create_artifact_version_response.id
     artifact_storage_root = _create_artifact_version_response.artifact_storage_root
     if isinstance(model_file_or_folder, str):
-        # Source is of type TruefoundryModelSource
-        source = ModelVersionSource.from_json(
-            TruefoundryModelSource(
-                type="truefoundry", uri=artifact_storage_root
-            ).to_json()
+        # Source is of type TrueFoundryArtifactSource
+        source = TrueFoundryArtifactSource(
+            type="truefoundry", uri=artifact_storage_root
         )
         artifacts_repo = MlFoundryArtifactsRepository(
             artifact_identifier=ArtifactIdentifier(
@@ -555,20 +549,19 @@ def _log_model_version(  # noqa: C901
             raise MlFoundryException("Failed to log model") from e
         finally:
             temp_dir.cleanup()
-    elif isinstance(model_file_or_folder, BlobStorageModelDirectory):
-        source = ModelVersionSource.from_json(
-            ExternalModelSource(type="external", uri=model_file_or_folder.uri).to_json()
-        )
+    elif isinstance(model_file_or_folder, BlobStorageDirectory):
+        source = ExternalArtifactSource(type="external", uri=model_file_or_folder.uri)
     else:
         raise MlFoundryException("Invalid model_file_or_folder provided")
 
     _framework = _ModelFramework.to_model_framework_type(framework)
+    _source_cls = typing.get_type_hints(ModelVersionManifest)["source"]
     model_manifest = ModelVersionManifest(
         description=description,
         metadata=metadata,
-        source=source,
+        source=_source_cls.from_dict(source.dict()),
         framework=Framework.from_dict(_framework.dict()) if _framework else None,
-        step=step if run else 0,
+        step=step,
     )
     artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
@@ -578,7 +571,7 @@ def _log_model_version(  # noqa: C901
             artifact_metadata=metadata,
             internal_metadata=None,
             step=model_manifest.step,
-            manifest=model_manifest,
+            manifest=Manifest.from_dict(model_manifest.to_dict()),
         )
     )
     return ModelVersion.from_fqn(fqn=artifact_version_response.artifact_version.fqn)

truefoundry/ml/mlfoundry_api.py

@@ -48,12 +48,12 @@ 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 (
-    BlobStorageModelDirectory,
     ModelVersion,
     _log_model_version,
 )
@@ -1226,7 +1226,7 @@ class MlFoundry:
         *,
         ml_repo: str,
         name: str,
-        model_file_or_folder: Union[str, BlobStorageModelDirectory],
+        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,
@@ -1244,12 +1244,12 @@ class MlFoundry:
                 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 (Union[str, BlobStorageModelDirectory]):
+            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(...)`.
-                BlobStorageModelDirectory:
+                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`.

truefoundry/ml/mlfoundry_run.py

@@ -44,10 +44,13 @@ 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 (
-    BlobStorageModelDirectory,
     ModelVersion,
     _log_model_version,
 )
@@ -926,7 +929,7 @@ class MlFoundryRun:
         self,
         *,
         name: str,
-        model_file_or_folder: Union[str, BlobStorageModelDirectory],
+        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,
@@ -945,12 +948,12 @@ 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 (Union[str, BlobStorageModelDirectory]):
+            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(...)`.
-                BlobStorageModelDirectory:
+                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`.

truefoundry/ml/model_framework.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Literal, Optional, Union, get_args
+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
@@ -161,7 +161,7 @@ class _ModelFramework(BaseModel):
         )
 
     @classmethod
-    def from_dict(cls, obj: dict) -> ModelFrameworkType:
+    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

truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py

@@ -1,4 +1,6 @@
 # file: client.py
+import io
+import os
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlencode, urljoin
@@ -125,7 +127,12 @@ class SignedURLClient:
         )
 
     @log_time(prefix=LOG_PREFIX)
-    def _upload_data(self, signed_url: str, data: Any) -> None:
+    def _upload_data(
+        self,
+        signed_url: str,
+        data: Union[bytes, io.BufferedReader],
+        headers: Optional[Dict] = None,
+    ) -> None:
         """
         Upload data to the specified storage path using a signed URL.
 
@@ -133,10 +140,15 @@ class SignedURLClient:
             signed_url: str: The signed URL to upload the data to.
             data: Bytes or IO: The data to upload.
         """
+        if isinstance(data, io.BufferedReader):
+            if os.fstat(data.fileno()).st_size == 0:
+                data = b""
         try:
+            headers = headers or {}
+            headers["Content-Type"] = "application/octet-stream"
             response = self.session.put(
                 url=signed_url,
-                headers={"Content-Type": "application/octet-stream"},
+                headers=headers,
                 data=data,
                 timeout=REQUEST_TIMEOUT,
             )
@@ -153,7 +165,11 @@ class SignedURLClient:
             headers=self.signed_url_server_headers,
         )
         pre_signed_object_dto = SignedURLAPIResponseDto.parse_obj(signed_object)
-        self._upload_data(pre_signed_object_dto.signed_url, data)
+        self._upload_data(
+            signed_url=pre_signed_object_dto.signed_url,
+            data=data,
+            headers=pre_signed_object_dto.headers,
+        )
         return storage_uri
 
     @log_time(prefix=LOG_PREFIX)
@@ -167,19 +183,30 @@ class SignedURLClient:
         )
         pre_signed_object_dto = SignedURLAPIResponseDto.parse_obj(response)
         with open(file_path, "rb") as file:
-            self._upload_data(pre_signed_object_dto.signed_url, file)
+            self._upload_data(
+                signed_url=pre_signed_object_dto.signed_url,
+                data=file,
+                headers=pre_signed_object_dto.headers,
+            )
         return storage_uri
 
     @log_time(prefix=LOG_PREFIX)
     def _download_file(
-        self, signed_url: str, local_path: Optional[str] = None
+        self,
+        signed_url: str,
+        local_path: Optional[str] = None,
+        headers: Optional[Dict] = None,
     ) -> Optional[bytes]:
         """Common method to download a file using a signed URL."""
         try:
+            if headers is None:
+                headers = {"Content-Type": "application/octet-stream"}
+            else:
+                headers["Content-Type"] = "application/octet-stream"
             response = self.session.get(
                 signed_url,
                 stream=True,
-                headers={"Content-Type": "application/octet-stream"},
+                headers=headers,
                 timeout=REQUEST_TIMEOUT,
             )
             response.raise_for_status()
@@ -202,7 +229,11 @@ class SignedURLClient:
             headers=self.signed_url_server_headers,
         )
         presigned_object = SignedURLAPIResponseDto.parse_obj(response)
-        self._download_file(presigned_object.signed_url, local_path)
+        self._download_file(
+            signed_url=presigned_object.signed_url,
+            local_path=local_path,
+            headers=presigned_object.headers,
+        )
         return local_path
 
     @log_time(prefix=LOG_PREFIX)
@@ -214,7 +245,9 @@ class SignedURLClient:
             headers=self.signed_url_server_headers,
         )
         presigned_object = SignedURLAPIResponseDto.parse_obj(response)
-        return self._download_file(presigned_object.signed_url)
+        return self._download_file(
+            signed_url=presigned_object.signed_url, headers=presigned_object.headers
+        )
 
     @log_time(prefix=LOG_PREFIX)
     def exists(self, uri: str) -> bool:

{truefoundry-0.5.0rc2.dist-info → truefoundry-0.5.0rc4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.5.0rc2
+Version: 0.5.0rc4
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
@@ -21,7 +21,7 @@ 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.12.2) ; extra == "workflow"
+Requires-Dist: flytekit (==1.13.13) ; extra == "workflow"
 Requires-Dist: gitignorefile (>=1.1.2,<2.0.0)
 Requires-Dist: importlib-metadata (>=4.11.3,<9.0.0)
 Requires-Dist: importlib-resources (>=5.2.0,<7.0.0)