truefoundry 0.4.4rc12__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

truefoundry/ml/log_types/artifacts/artifact.py

@@ -29,9 +29,10 @@ from truefoundry.ml.autogen.entities.artifacts import ChatPrompt
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import INTERNAL_METADATA_PATH
 from truefoundry.ml.log_types.artifacts.utils import (
+    _get_src_dest_pairs,
     _validate_artifact_metadata,
     _validate_description,
-    calculate_local_directory_size,
+    calculate_total_size,
 )
 from truefoundry.ml.logger import logger
 from truefoundry.ml.session import _get_api_client
@@ -397,7 +398,7 @@ class ChatPromptVersion(ArtifactVersion):
     @property
     def extra_parameters(self) -> Dict[str, Any]:
         _extra_parameters = self._chat_prompt.model_configuration.extra_parameters
-        return _extra_parameters.dict(exclude_unset=True) if _extra_parameters else {}
+        return _extra_parameters if _extra_parameters else {}
 
     @property
     def variables(self) -> Dict[str, Any]:
@@ -409,6 +410,7 @@ def _log_artifact_version_helper(
     name: str,
     artifact_type: ArtifactType,
     artifact_dir: tempfile.TemporaryDirectory,
+    dest_to_src_map: Dict[str, str],
     mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
     ml_repo_id: Optional[str] = None,
     description: Optional[str] = None,
@@ -443,15 +445,17 @@ def _log_artifact_version_helper(
         ),
         api_client=mlfoundry_artifacts_api.api_client,
     )
-    total_size = calculate_local_directory_size(artifact_dir)
+
+    total_size = calculate_total_size(list(dest_to_src_map.values()))
     try:
         logger.info(
-            "Packaging and uploading files to remote with Artifact Size: %.6f MB",
+            "Packaging and uploading files to remote with size: %.6f MB",
             total_size / 1000000.0,
         )
-        artifacts_repo.log_artifacts(
-            local_dir=artifact_dir.name, artifact_path=None, progress=progress
+        src_dest_pairs = _get_src_dest_pairs(
+            root_dir=artifact_dir.name, dest_to_src_map=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(

truefoundry/ml/log_types/artifacts/dataset.py

@@ -19,9 +19,10 @@ from truefoundry.ml.entities import FileInfo
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.utils import (
     _copy_additional_files,
+    _get_src_dest_pairs,
     _validate_artifact_metadata,
     _validate_description,
-    calculate_local_directory_size,
+    calculate_total_size,
 )
 from truefoundry.ml.logger import logger
 from truefoundry.ml.session import _get_api_client
@@ -139,6 +140,7 @@ class DataDirectory:
         file_paths: List[
             Union[Tuple[str], Tuple[str, Optional[str]], DataDirectoryPath]
         ],
+        progress: Optional[bool] = None,
     ) -> None:
         """Logs File in the `DataDirectory`.
 
@@ -207,7 +209,7 @@ class DataDirectory:
 
         try:
             logger.info("Copying the files to add")
-            _copy_additional_files(
+            temp_dest_to_src_map = _copy_additional_files(
                 root_dir=temp_dir.name,
                 files_dir="",
                 model_dir=None,
@@ -219,13 +221,19 @@ class DataDirectory:
             raise MlFoundryException("Failed to Add Files to DataDirectory") from e
 
         artifacts_repo = self._get_artifacts_repo()
-        total_size = calculate_local_directory_size(temp_dir)
+        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",
+                "Packaging and uploading files to remote with size: %.6f MB",
                 total_size / 1000000.0,
             )
-            artifacts_repo.log_artifacts(local_dir=temp_dir.name, artifact_path=None)
+            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:
             raise MlFoundryException("Failed to Add Files to DataDirectory") from e
         finally:

truefoundry/ml/log_types/artifacts/general_artifact.py

@@ -77,7 +77,7 @@ def _log_artifact_version(
         os.makedirs(local_files_dir, exist_ok=True)
 
         logger.info("Copying the files to log")
-        _copy_additional_files(
+        temp_dest_to_src_map = _copy_additional_files(
             root_dir=temp_dir.name,
             files_dir=internal_metadata.files_dir,
             model_dir=None,
@@ -94,6 +94,7 @@ def _log_artifact_version(
     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
 
     return _log_artifact_version_helper(
         run=run,
@@ -101,6 +102,7 @@ def _log_artifact_version(
         name=name,
         artifact_type=ArtifactType.ARTIFACT,
         artifact_dir=temp_dir,
+        dest_to_src_map=temp_dest_to_src_map,
         mlfoundry_artifacts_api=mlfoundry_artifacts_api,
         description=description,
         internal_metadata=internal_metadata,

truefoundry/ml/log_types/artifacts/model.py

@@ -5,57 +5,57 @@ import logging
 import os.path
 import tempfile
 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, Sequence, Tuple, Union
 
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
     MlFoundryArtifactsRepository,
 )
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    AddCustomMetricsToModelVersionRequestDto,
     ArtifactType,
+    ArtifactVersionSerializationFormat,
     CreateArtifactVersionRequestDto,
-    CreateModelVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
+    ExternalModelSource,
     FinalizeArtifactVersionRequestDto,
-    MetricDto,
+    Framework,
     MlfoundryArtifactsApi,
     ModelDto,
     ModelVersionDto,
+    ModelVersionManifest,
     NotifyArtifactVersionFailureDto,
+    TruefoundryModelSource,
     UpdateModelVersionRequestDto,
 )
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    InternalMetadata as InternalMetadataDto,
+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.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,
     _validate_artifact_metadata,
     _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__
+from truefoundry.pydantic_v1 import BaseModel, Extra, StrictStr
 
 if TYPE_CHECKING:
     from truefoundry.ml.mlfoundry_run import MlFoundryRun
 
-logger = logging.getLogger("truefoundry.ml")
-
 
-# TODO: Support async download and upload
+logger = logging.getLogger(__name__)
 
 
+# TODO: Support async download and upload
 class ModelVersionInternalMetadata(BaseModel):
     class Config:
         extra = Extra.allow
@@ -75,6 +75,10 @@ class ModelVersionInternalMetadata(BaseModel):
         return dct
 
 
+class BlobStorageModelDirectory(BaseModel):
+    uri: StrictStr
+
+
 class ModelVersionDownloadInfo(BaseModel):
     download_dir: str
     model_dir: str
@@ -97,8 +101,7 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._metrics: List[MetricDto] = []
-        self._set_metrics_attr()
+        self._framework: Optional["Framework"] = None
         self._set_mutable_attrs()
 
     @classmethod
@@ -136,21 +139,20 @@ 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)
+        self._framework = (
+            copy.deepcopy(self._model_version.manifest.framework)
+            if self._model_version.manifest
+            else None
+        )
 
     def _refetch_model_version(self):
         _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()
 
     def __repr__(self):
@@ -199,6 +201,8 @@ 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]:
@@ -210,14 +214,43 @@ class ModelVersion:
         """set the metadata for current model"""
         _validate_artifact_metadata(value)
         self._metadata = value
+        if self._model_version.manifest:
+            self._model_version.manifest.metadata = 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 (
+            _ModelFramework.from_dict(self._framework.actual_instance.dict())
+            if self._framework and self._framework.actual_instance
+            else None
+        )
+
+    @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:
+            warnings.warn(
+                message="This model version was created using an older serialization format. Framework will not be updated",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
 
     @property
     def created_at(self) -> datetime.datetime:
@@ -272,11 +305,29 @@ 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.serialization_format
+            == ArtifactVersionSerializationFormat.V2
+        ):
+            model_framework = (
+                ModelFramework(self.framework.type)
+                if self.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(
@@ -288,12 +339,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,
@@ -329,7 +382,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
@@ -381,50 +434,36 @@ class ModelVersion:
                 id=self._model_version.id,
                 description=self.description,
                 artifact_metadata=self.metadata,
+                manifest=self._model_version.manifest,
             )
         )
         self._model_version = _model_version.model_version
-        self._set_metrics_attr()
         self._set_mutable_attrs()
 
 
-def calculate_model_size(artifact_dir: tempfile.TemporaryDirectory):
-    """
-    Tells about the size of the model
-
-    Args:
-        artifact_dir (str): directory in which model is present.
-
-    Returns:
-        total size of the model
-    """
-    total_size = 0
-    for path, _dirs, files in os.walk(artifact_dir.name):
-        for f in files:
-            file_path = os.path.join(path, f)
-            total_size += os.stat(file_path).st_size
-    return total_size
-
-
 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, BlobStorageModelDirectory],
     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,
 ) -> 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, BlobStorageModelDirectory)):
+        raise MlFoundryException(
+            "model_file_or_folder should be of type str or BlobStorageModelDirectory"
+        )
+
     if mlfoundry_artifacts_api and not ml_repo_id:
         raise MlFoundryException(
             "If mlfoundry_artifacts_api is passed, ml_repo_id must also be passed"
@@ -434,87 +473,41 @@ def _log_model_version(  # noqa: C901
 
     assert mlfoundry_artifacts_api is not None
 
-    custom_metrics = custom_metrics or []
+    step = step or 0
+    total_size = None
     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)
 
     _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 = [
-            (model_file_or_folder, MODEL_DIR_NAME.rstrip(os.sep) + os.sep)
-        ]
-        _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")
-            _copy_additional_files(
+    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,
+                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)
+            # 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
 
     # create entry
     _create_artifact_version_response = (
@@ -527,80 +520,65 @@ 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,
-    )
-    model_size = calculate_model_size(temp_dir)
-    try:
-        logger.info(
-            "Packaging and uploading files to remote with Total Size: %.6f MB",
-            model_size / 1000000.0,
+    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()
         )
-        artifacts_repo.log_artifacts(
-            local_dir=temp_dir.name, artifact_path=None, progress=progress
+        artifacts_repo = MlFoundryArtifactsRepository(
+            artifact_identifier=ArtifactIdentifier(
+                artifact_version_id=uuid.UUID(version_id)
+            ),
+            api_client=mlfoundry_artifacts_api.api_client,
         )
-    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,
+            )
+            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, BlobStorageModelDirectory):
+        source = ModelVersionSource.from_json(
+            ExternalModelSource(type="external", uri=model_file_or_folder.uri).to_json()
         )
-        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 {}
+    else:
+        raise MlFoundryException("Invalid model_file_or_folder provided")
+
+    _framework = _ModelFramework.to_model_framework_type(framework)
+    model_manifest = ModelVersionManifest(
+        description=description,
+        metadata=metadata,
+        source=source,
+        framework=Framework.from_dict(_framework.dict()) if _framework else None,
+        step=step if run else 0,
     )
-    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=model_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_size=total_size,
             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=model_manifest,
         )
     )
-    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)