truefoundry 0.4.10__py3-none-any.whl → 0.5.0__py3-none-any.whl

truefoundry/ml/log_types/artifacts/model.py

@@ -4,40 +4,39 @@ import json
 import logging
 import os.path
 import tempfile
+import typing
 import uuid
+import warnings
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
     MlFoundryArtifactsRepository,
 )
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    AddCustomMetricsToModelVersionRequestDto,
     ArtifactType,
     CreateArtifactVersionRequestDto,
-    CreateModelVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
+    ExternalArtifactSource,
     FinalizeArtifactVersionRequestDto,
-    MetricDto,
+    Framework,
+    Manifest,
     MlfoundryArtifactsApi,
     ModelDto,
     ModelVersionDto,
+    ModelVersionEnvironment,
+    ModelVersionManifest,
     NotifyArtifactVersionFailureDto,
+    TrueFoundryArtifactSource,
     UpdateModelVersionRequestDto,
 )
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    InternalMetadata as InternalMetadataDto,
-)
 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 (
-    FILES_DIR,
     INTERNAL_METADATA_PATH,
-    MODEL_DIR_NAME,
-    MODEL_SCHEMA_UPDATE_FAILURE_HELP,
 )
-from truefoundry.ml.log_types.artifacts.model_extras import CustomMetric, ModelSchema
 from truefoundry.ml.log_types.artifacts.utils import (
     _copy_additional_files,
     _get_src_dest_pairs,
@@ -45,19 +44,18 @@ from truefoundry.ml.log_types.artifacts.utils import (
     _validate_description,
     calculate_total_size,
 )
+from truefoundry.ml.model_framework import ModelFrameworkType, _ModelFramework
 from truefoundry.ml.session import _get_api_client
 from truefoundry.pydantic_v1 import BaseModel, Extra
-from truefoundry.version import __version__
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
+
 logger = logging.getLogger(__name__)
 
 
 # TODO: Support async download and upload
-
-
 class ModelVersionInternalMetadata(BaseModel):
     class Config:
         extra = Extra.allow
@@ -99,8 +97,8 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._metrics: List[MetricDto] = []
-        self._set_metrics_attr()
+        self._model_schema: Optional[Dict[str, Any]] = None
+        self._framework: Optional[ModelFrameworkType] = None
         self._set_mutable_attrs()
 
     @classmethod
@@ -138,22 +136,34 @@ class ModelVersion:
                 "Model Version was deleted, cannot perform updates on a deleted version"
             )
 
-    def _set_metrics_attr(self):
-        self._metrics = sorted(
-            self._model_version.metrics or [], key=lambda m: m.timestamp
-        )
-
     def _set_mutable_attrs(self):
-        self._description = self._model_version.description or ""
-        self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
+        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
+            )
+            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
+            )
+            self._model_schema = None
 
-    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_metrics_attr()
-        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 +197,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
 
@@ -211,23 +223,51 @@ class ModelVersion:
     def metadata(self, value: Dict[str, Any]):
         """set the metadata for current model"""
         _validate_artifact_metadata(value)
-        self._metadata = value
+        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
+
+    @model_schema.setter
+    def model_schema(self, value: Optional[Dict[str, Any]]):
+        """set the model_schema for current 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",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+            return
+        self._model_schema = copy.deepcopy(value)
 
     @property
-    def metrics(self) -> Dict[str, Union[float, int]]:
-        """get the metrics for the current version of the model"""
-        metrics_as_kv: Dict[str, Union[float, int]] = {}
-        for metric in self._metrics:
-            metrics_as_kv[metric.key] = metric.value
-        return metrics_as_kv
+    def framework(self) -> Optional["ModelFrameworkType"]:
+        """Get the framework of the model"""
+        return self._framework
+
+    @framework.setter
+    def framework(
+        self, value: Optional[Union[str, ModelFramework, "ModelFrameworkType"]]
+    ):
+        """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
 
@@ -274,11 +314,27 @@ class ModelVersion:
         path: Optional[Union[str, Path]],
         overwrite: bool = False,
         progress: Optional[bool] = None,
-    ) -> Tuple[ModelVersionInternalMetadata, ModelVersionDownloadInfo]:
+    ) -> ModelVersionDownloadInfo:
         self._ensure_not_deleted()
         download_dir = self.raw_download(
             path=path, overwrite=overwrite, progress=progress
         )
+
+        if self._model_version.manifest:
+            _framework = self._model_version.manifest.framework
+            model_framework = (
+                ModelFramework(_framework.actual_instance.type)
+                if _framework
+                else ModelFramework.UNKNOWN
+            )
+            download_info = ModelVersionDownloadInfo(
+                download_dir=download_dir,
+                model_dir=download_dir,
+                model_framework=model_framework,
+                model_filename=None,
+            )
+            return download_info
+
         internal_metadata_path = os.path.join(download_dir, INTERNAL_METADATA_PATH)
         if not os.path.exists(internal_metadata_path):
             raise MlFoundryException(
@@ -290,12 +346,14 @@ class ModelVersion:
         download_info = ModelVersionDownloadInfo(
             download_dir=os.path.join(download_dir, internal_metadata.files_dir),
             model_dir=os.path.join(
-                download_dir, internal_metadata.files_dir, internal_metadata.model_dir
+                download_dir,
+                internal_metadata.files_dir,
+                internal_metadata.model_dir,
             ),
             model_framework=internal_metadata.framework,
             model_filename=internal_metadata.model_filename,
         )
-        return internal_metadata, download_info
+        return download_info
 
     def download(
         self,
@@ -331,7 +389,7 @@ class ModelVersion:
             print(download_info.model_dir)
             ```
         """
-        _, download_info = self._download(
+        download_info = self._download(
             path=path, overwrite=overwrite, progress=progress
         )
         return download_info
@@ -378,37 +436,55 @@ 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,
+        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.framework = (
+                Framework.from_dict(self.framework.dict()) if self.framework else None
             )
-        )
-        self._model_version = _model_version.model_version
-        self._set_metrics_attr()
-        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: str,
-    framework: Optional[Union[ModelFramework, str]],
+    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,
-    model_schema: Optional[Union[Dict[str, Any], ModelSchema]] = None,
-    custom_metrics: Optional[List[Union[CustomMetric, Dict[str, Any]]]] = None,
     step: Optional[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:
     if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
         raise MlFoundryException(
             "Exactly one of run, mlfoundry_artifacts_api should be passed"
         )
+
+    if not isinstance(model_file_or_folder, (str, BlobStorageDirectory)):
+        raise MlFoundryException(
+            "model_file_or_folder should be of type str or BlobStorageDirectory"
+        )
+
     if mlfoundry_artifacts_api and not ml_repo_id:
         raise MlFoundryException(
             "If mlfoundry_artifacts_api is passed, ml_repo_id must also be passed"
@@ -418,89 +494,29 @@ def _log_model_version(  # noqa: C901
 
     assert mlfoundry_artifacts_api is not None
 
-    custom_metrics = custom_metrics or []
-    metadata = metadata or {}
-    additional_files = additional_files or {}
     step = step or 0
-
-    # validations
-    if framework is None:
-        framework = ModelFramework.UNKNOWN
-    elif not isinstance(framework, ModelFramework):
-        framework = ModelFramework(framework)
+    total_size = None
+    metadata = metadata or {}
 
     _validate_description(description)
     _validate_artifact_metadata(metadata)
 
-    if model_schema is not None and not isinstance(model_schema, ModelSchema):
-        model_schema = ModelSchema.parse_obj(model_schema)
-
-    if custom_metrics and not model_schema:
-        raise MlFoundryException(
-            "Custom Metrics defined without adding the Model Schema"
-        )
-    custom_metrics = [
-        CustomMetric.parse_obj(cm) if not isinstance(cm, CustomMetric) else cm
-        for cm in custom_metrics
-    ]
-
-    logger.info("Logging model and additional files, this might take a while ...")
-    temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
-
-    internal_metadata = ModelVersionInternalMetadata(
-        framework=framework,
-        files_dir=FILES_DIR,
-        model_dir=MODEL_DIR_NAME,
-        model_filename=(
-            os.path.basename(model_file_or_folder)
-            if model_file_or_folder and os.path.isfile(model_file_or_folder)
-            else None
-        ),
-        mlfoundry_version=__version__,
-        truefoundry_version=__version__,
-    )
-
-    try:
-        local_files_dir = os.path.join(temp_dir.name, internal_metadata.files_dir)
-        os.makedirs(local_files_dir, exist_ok=True)
-        # in case model was None, we still create an empty dir
-        local_model_dir = os.path.join(local_files_dir, internal_metadata.model_dir)
-        os.makedirs(local_model_dir, exist_ok=True)
-
-        logger.info("Adding model file/folder to model version content")
-        _model_file_or_folder: Sequence[Tuple[str, str]] = [
-            (model_file_or_folder, MODEL_DIR_NAME.rstrip(os.sep) + os.sep),
-        ]
-
-        temp_dest_to_src_map = _copy_additional_files(
-            root_dir=temp_dir.name,
-            files_dir=internal_metadata.files_dir,
-            model_dir=internal_metadata.model_dir,
-            additional_files=_model_file_or_folder,
-            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")
+    if isinstance(model_file_or_folder, str):
+        logger.info("Logging model and additional files, this might take a while ...")
+        temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
+        try:
+            logger.info("Adding model file/folder to model version content")
             temp_dest_to_src_map = _copy_additional_files(
                 root_dir=temp_dir.name,
-                files_dir=internal_metadata.files_dir,
-                model_dir=internal_metadata.model_dir,
-                additional_files=additional_files,
-                ignore_model_dir_dest_conflict=False,
-                existing_dest_to_src_map=temp_dest_to_src_map,
+                files_dir="",
+                model_dir=None,
+                additional_files=[(model_file_or_folder, "")],
+                ignore_model_dir_dest_conflict=True,
             )
-    except Exception as e:
-        temp_dir.cleanup()
-        raise MlFoundryException("Failed to log model") from e
 
-    # save internal metadata
-    local_internal_metadata_path = os.path.join(temp_dir.name, INTERNAL_METADATA_PATH)
-    os.makedirs(os.path.dirname(local_internal_metadata_path), exist_ok=True)
-    with open(local_internal_metadata_path, "w") as f:
-        json.dump(internal_metadata.dict(), f)
-    temp_dest_to_src_map[local_internal_metadata_path] = local_internal_metadata_path
+        except Exception as e:
+            temp_dir.cleanup()
+            raise MlFoundryException("Failed to log model") from e
 
     # create entry
     _create_artifact_version_response = (
@@ -513,82 +529,64 @@ def _log_model_version(  # noqa: C901
         )
     )
     version_id = _create_artifact_version_response.id
-    artifacts_repo = MlFoundryArtifactsRepository(
-        artifact_identifier=ArtifactIdentifier(
-            artifact_version_id=uuid.UUID(version_id)
-        ),
-        api_client=mlfoundry_artifacts_api.api_client,
-    )
-
-    total_size = calculate_total_size(list(temp_dest_to_src_map.values()))
-    try:
-        logger.info(
-            "Packaging and uploading files to remote with size: %.6f MB",
-            total_size / 1000000.0,
+    artifact_storage_root = _create_artifact_version_response.artifact_storage_root
+    if isinstance(model_file_or_folder, str):
+        # Source is of type TrueFoundryArtifactSource
+        source = TrueFoundryArtifactSource(
+            type="truefoundry", uri=artifact_storage_root
         )
-        src_dest_pairs = _get_src_dest_pairs(
-            root_dir=temp_dir.name, dest_to_src_map=temp_dest_to_src_map
+        artifacts_repo = MlFoundryArtifactsRepository(
+            artifact_identifier=ArtifactIdentifier(
+                artifact_version_id=uuid.UUID(version_id)
+            ),
+            api_client=mlfoundry_artifacts_api.api_client,
         )
-        artifacts_repo.log_artifacts(src_dest_pairs=src_dest_pairs, progress=progress)
-    except Exception as e:
-        mlfoundry_artifacts_api.notify_failure_post(
-            notify_artifact_version_failure_dto=NotifyArtifactVersionFailureDto(
-                id=version_id
+        total_size = calculate_total_size(list(temp_dest_to_src_map.values()))
+        try:
+            logger.info(
+                "Packaging and uploading files to remote with Total Size: %.6f MB",
+                total_size / 1000000.0,
             )
-        )
-        raise MlFoundryException("Failed to log model") from e
-    finally:
-        temp_dir.cleanup()
-
-    # Note: Here we call from_dict instead of directly passing in init and relying on it
-    # to convert because the complicated union of types generates a custom type to handle casting
-    # Check the source of `InternalMetadataDto` to see the generated code
-    internal_metadata_dto = InternalMetadataDto.from_dict(
-        internal_metadata.dict() if internal_metadata is not None else {}
+            src_dest_pairs = _get_src_dest_pairs(
+                root_dir=temp_dir.name, dest_to_src_map=temp_dest_to_src_map
+            )
+            artifacts_repo.log_artifacts(
+                src_dest_pairs=src_dest_pairs, progress=progress
+            )
+        except Exception as e:
+            mlfoundry_artifacts_api.notify_failure_post(
+                notify_artifact_version_failure_dto=NotifyArtifactVersionFailureDto(
+                    id=version_id
+                )
+            )
+            raise MlFoundryException("Failed to log model") from e
+        finally:
+            temp_dir.cleanup()
+    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_cls.from_dict(source.dict()),
+        framework=Framework.from_dict(_framework.dict()) if _framework else None,
+        environment=environment,
+        step=step,
+        model_schema=model_schema,
     )
-    mlfoundry_artifacts_api.finalize_artifact_version_post(
+    artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
             id=version_id,
             run_uuid=run.run_id if run else None,
             artifact_size=total_size,
-            internal_metadata=internal_metadata_dto,
-            step=step if run else None,
-        )
-    )
-    _model_version = mlfoundry_artifacts_api.create_model_version_post(
-        create_model_version_request_dto=CreateModelVersionRequestDto(
-            artifact_version_id=version_id,
-            description=description,
             artifact_metadata=metadata,
-            internal_metadata=internal_metadata_dto,
-            data_path=INTERNAL_METADATA_PATH,
-            step=step if run else None,
+            internal_metadata=None,
+            step=model_manifest.step,
+            manifest=Manifest.from_dict(model_manifest.to_dict()),
         )
     )
-    model_version = _model_version.model_version
-
-    # update model schema at end
-    update_args: Dict[str, Any] = {
-        "id": version_id,
-        "model_framework": framework.value,
-    }
-    if model_schema:
-        update_args["model_schema"] = model_schema
-
-    try:
-        _model_version = mlfoundry_artifacts_api.update_model_version_post(
-            update_model_version_request_dto=UpdateModelVersionRequestDto(**update_args)
-        )
-        model_version = _model_version.model_version
-        if model_schema:
-            _model_version = mlfoundry_artifacts_api.add_custom_metrics_to_model_version_post(
-                add_custom_metrics_to_model_version_request_dto=AddCustomMetricsToModelVersionRequestDto(
-                    id=version_id, custom_metrics=custom_metrics
-                )
-            )
-            model_version = _model_version.model_version
-    except Exception:
-        # TODO (chiragjn): what is the best exception to catch here?
-        logger.error(MODEL_SCHEMA_UPDATE_FAILURE_HELP.format(fqn=model_version.fqn))
-
-    return ModelVersion.from_fqn(fqn=model_version.fqn)
+    return ModelVersion.from_fqn(fqn=artifact_version_response.artifact_version.fqn)