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

truefoundry/ml/log_types/artifacts/artifact.py

@@ -3,9 +3,12 @@ import datetime
 import json
 import os
 import tempfile
+import typing
 import uuid
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Union
+
+from pydantic import StrictStr
 
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
@@ -15,11 +18,15 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
     ArtifactType,
     ArtifactVersionDto,
+    ArtifactVersionManifest,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
+    ExternalArtifactSource,
     FinalizeArtifactVersionRequestDto,
+    Manifest,
     MlfoundryArtifactsApi,
     NotifyArtifactVersionFailureDto,
+    TrueFoundryArtifactSource,
     UpdateArtifactVersionRequestDto,
 )
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
@@ -47,6 +54,10 @@ class ArtifactPath(NamedTuple):
     dest: Optional[str] = None
 
 
+class BlobStorageDirectory(BaseModel):
+    uri: StrictStr
+
+
 class ArtifactVersionInternalMetadata(BaseModel):
     class Config:
         extra = Extra.allow
@@ -56,7 +67,6 @@ class ArtifactVersionInternalMetadata(BaseModel):
 
 class ArtifactVersionDownloadInfo(BaseModel):
     download_dir: str
-    content_dir: str
 
 
 class ArtifactVersion:
@@ -116,16 +126,24 @@ class ArtifactVersion:
                 "Artifact Version was deleted, cannot access a deleted version"
             )
 
-    def _set_mutable_attrs(self, refetch=False):
-        if refetch:
-            _artifact_version = (
-                self._mlfoundry_artifacts_api.get_artifact_version_by_id_get(
-                    id=self._artifact_version.id
-                )
+    def _set_mutable_attrs(self):
+        if self._artifact_version.manifest:
+            manifest = self._artifact_version.manifest.actual_instance
+            self._description = manifest.description or ""
+            self._metadata = copy.deepcopy(manifest.metadata)
+        else:
+            self._description = self._artifact_version.description or ""
+            self._metadata = copy.deepcopy(self._artifact_version.artifact_metadata)
+
+    def _refetch_artifact_version(self, reset_mutable_attrs: bool = True):
+        _artifact_version = (
+            self._mlfoundry_artifacts_api.get_artifact_version_by_id_get(
+                id=self._artifact_version.id
             )
-            self._artifact_version = _artifact_version.artifact_version
-        self._description = self._artifact_version.description or ""
-        self._metadata = copy.deepcopy(self._artifact_version.artifact_metadata)
+        )
+        self._artifact_version = _artifact_version.artifact_version
+        if reset_mutable_attrs:
+            self._set_mutable_attrs()
 
     def __repr__(self):
         return f"{self.__class__.__name__}(fqn={self.fqn!r})"
@@ -159,12 +177,14 @@ class ArtifactVersion:
         return self._artifact_version.fqn
 
     @property
-    def step(self) -> int:
+    def step(self) -> Optional[int]:
         """Get the step in which artifact was created"""
+        if self._artifact_version.manifest:
+            return self._artifact_version.manifest.actual_instance.step
         return self._artifact_version.step
 
     @property
-    def description(self) -> Optional[str]:
+    def description(self) -> str:
         """Get description of the artifact"""
         return self._description
 
@@ -183,15 +203,15 @@ class ArtifactVersion:
     def metadata(self, value: Dict[str, Any]):
         """set the metadata for current artifact"""
         _validate_artifact_metadata(value)
-        self._metadata = value
+        self._metadata = copy.deepcopy(value)
 
     @property
-    def created_at(self) -> datetime.datetime:
+    def created_at(self) -> Optional[datetime.datetime]:
         """Get the time at which artifact was created"""
         return self._artifact_version.created_at
 
     @property
-    def updated_at(self) -> datetime.datetime:
+    def updated_at(self) -> Optional[datetime.datetime]:
         """Get the information about when the artifact was updated"""
         return self._artifact_version.updated_at
 
@@ -240,11 +260,17 @@ class ArtifactVersion:
         path: Optional[Union[str, Path]],
         overwrite: bool = False,
         progress: Optional[bool] = None,
-    ) -> Tuple[ArtifactVersionInternalMetadata, str]:
+    ) -> ArtifactVersionDownloadInfo:
         self._ensure_not_deleted()
         download_dir = self.raw_download(
             path=path, overwrite=overwrite, progress=progress
         )
+        if self._artifact_version.manifest:
+            download_info = ArtifactVersionDownloadInfo(
+                download_dir=download_dir,
+            )
+            return download_info
+
         internal_metadata_path = os.path.join(download_dir, INTERNAL_METADATA_PATH)
         if not os.path.exists(internal_metadata_path):
             raise MlFoundryException(
@@ -254,7 +280,7 @@ class ArtifactVersion:
         with open(internal_metadata_path) as f:
             internal_metadata = ArtifactVersionInternalMetadata.parse_obj(json.load(f))
         download_path = os.path.join(download_dir, internal_metadata.files_dir)
-        return internal_metadata, download_path
+        return ArtifactVersionDownloadInfo(download_dir=download_path)
 
     def download(
         self,
@@ -289,10 +315,10 @@ class ArtifactVersion:
             artifact_version.download(path="<your-desired-download-path>")
             ```
         """
-        _, download_path = self._download(
+        download_path = self._download(
             path=path, overwrite=overwrite, progress=progress
         )
-        return download_path
+        return download_path.download_dir
 
     def delete(self) -> bool:
         """
@@ -336,16 +362,31 @@ class ArtifactVersion:
             ```
         """
         self._ensure_not_deleted()
-
-        _artifact_version = self._mlfoundry_artifacts_api.update_artifact_version_post(
-            update_artifact_version_request_dto=UpdateArtifactVersionRequestDto(
-                id=self._artifact_version.id,
-                description=self.description,
-                artifact_metadata=self.metadata,
+        if self._artifact_version.manifest:
+            manifest = self._artifact_version.manifest.actual_instance
+            assert isinstance(manifest, ArtifactVersionManifest)
+            manifest.description = self.description
+            manifest.metadata = self.metadata
+        else:
+            manifest = None
+        try:
+            _artifact_version = (
+                self._mlfoundry_artifacts_api.update_artifact_version_post(
+                    update_artifact_version_request_dto=UpdateArtifactVersionRequestDto(
+                        id=self._artifact_version.id,
+                        description=self.description,
+                        artifact_metadata=self.metadata,
+                        manifest=manifest,
+                    )
+                )
             )
-        )
-        self._artifact_version = _artifact_version.artifact_version
-        self._set_mutable_attrs()
+        except Exception:
+            # rollback edits to internal object
+            self._refetch_artifact_version(reset_mutable_attrs=False)
+            raise
+        else:
+            self._artifact_version = _artifact_version.artifact_version
+            self._set_mutable_attrs()
 
 
 class ChatPromptVersion(ArtifactVersion):
@@ -404,12 +445,15 @@ class ChatPromptVersion(ArtifactVersion):
     def variables(self) -> Dict[str, Any]:
         return self._chat_prompt.variables or {}
 
+    def update(self):
+        raise NotImplementedError("ChatPromptVersion does not support update operation")
+
 
 def _log_artifact_version_helper(
-    run: "MlFoundryRun",
+    run: Optional["MlFoundryRun"],
     name: str,
     artifact_type: ArtifactType,
-    artifact_dir: tempfile.TemporaryDirectory,
+    artifact_dir: Union[tempfile.TemporaryDirectory, BlobStorageDirectory],
     dest_to_src_map: Dict[str, str],
     mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
     ml_repo_id: Optional[str] = None,
@@ -439,48 +483,72 @@ def _log_artifact_version_helper(
         )
     )
     version_id = _create_artifact_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(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_response.artifact_storage_root
+    if isinstance(artifact_dir, tempfile.TemporaryDirectory):
+        # Source is of type TrueFoundryArtifactSource
+        source = TrueFoundryArtifactSource(
+            type="truefoundry", uri=artifact_storage_root
         )
-        src_dest_pairs = _get_src_dest_pairs(
-            root_dir=artifact_dir.name, dest_to_src_map=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(dest_to_src_map.values()))
+        try:
+            logger.info(
+                "Packaging and uploading files to remote with size: %.6f MB",
+                total_size / 1000000.0,
+            )
+            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(
+                    id=version_id
+                )
+            )
+            raise MlFoundryException("Failed to log Artifact") from e
+        finally:
+            artifact_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 {}
+        )
+    elif isinstance(artifact_dir, BlobStorageDirectory):
+        source = ExternalArtifactSource(type="external", uri=artifact_dir.uri)
+    else:
+        raise MlFoundryException("Invalid artifact_dir provided")
+
+    artifact_manifest = None
+    if artifact_type == ArtifactType.ARTIFACT:
+        _source_cls = typing.get_type_hints(ArtifactVersionManifest)["source"]
+        artifact_manifest = ArtifactVersionManifest(
+            description=description,
+            metadata=metadata or {},
+            source=_source_cls.from_dict(source.dict()),
+            step=step,
         )
-        raise MlFoundryException("Failed to log Artifact") from e
-    finally:
-        artifact_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 {}
-    )
     finalize_artifact_version_request_dto = FinalizeArtifactVersionRequestDto(
         id=version_id,
         run_uuid=run.run_id if run else None,
-        description=description,
-        internal_metadata=internal_metadata_dto,
+        internal_metadata=internal_metadata_dto if internal_metadata else None,
         artifact_metadata=metadata,
-        data_path=INTERNAL_METADATA_PATH,
+        data_path=INTERNAL_METADATA_PATH if internal_metadata else None,
         step=step,
         artifact_size=total_size,
+        manifest=Manifest.from_dict(artifact_manifest.to_dict())
+        if artifact_manifest
+        else None,
     )
     _artifact_version = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=finalize_artifact_version_request_dto

truefoundry/ml/log_types/artifacts/general_artifact.py

@@ -1,8 +1,6 @@
 import collections
-import json
-import os.path
 import tempfile
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
@@ -12,13 +10,8 @@ from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.artifact import (
     ArtifactPath,
     ArtifactVersion,
-    ArtifactVersionInternalMetadata,
     _log_artifact_version_helper,
 )
-from truefoundry.ml.log_types.artifacts.constants import (
-    FILES_DIR,
-    INTERNAL_METADATA_PATH,
-)
 from truefoundry.ml.log_types.artifacts.utils import (
     _copy_additional_files,
     _validate_artifact_metadata,
@@ -26,9 +19,12 @@ from truefoundry.ml.log_types.artifacts.utils import (
 )
 from truefoundry.ml.logger import logger
 
+if TYPE_CHECKING:
+    from truefoundry.ml.mlfoundry_run import MlFoundryRun
+
 
 def _log_artifact_version(
-    run,
+    run: Optional["MlFoundryRun"],
     name: str,
     artifact_paths: List[Union[ArtifactPath, Tuple[str, Optional[str]], Tuple[str]]],
     mlfoundry_artifacts_api: Optional[MlfoundryArtifactsApi] = None,
@@ -68,18 +64,10 @@ def _log_artifact_version(
     logger.info("Logging the artifact, this might take a while ...")
     temp_dir = tempfile.TemporaryDirectory(prefix="truefoundry-")
 
-    internal_metadata = ArtifactVersionInternalMetadata(
-        files_dir=FILES_DIR,
-    )
-
     try:
-        local_files_dir = os.path.join(temp_dir.name, internal_metadata.files_dir)
-        os.makedirs(local_files_dir, exist_ok=True)
-
-        logger.info("Copying the files to log")
         temp_dest_to_src_map = _copy_additional_files(
             root_dir=temp_dir.name,
-            files_dir=internal_metadata.files_dir,
+            files_dir="",
             model_dir=None,
             additional_files=artifact_paths,
         )
@@ -89,13 +77,6 @@ def _log_artifact_version(
         temp_dir.cleanup()
         raise MlFoundryException("Failed to log artifact") 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
-
     return _log_artifact_version_helper(
         run=run,
         ml_repo_id=ml_repo_id,
@@ -105,7 +86,7 @@ def _log_artifact_version(
         dest_to_src_map=temp_dest_to_src_map,
         mlfoundry_artifacts_api=mlfoundry_artifacts_api,
         description=description,
-        internal_metadata=internal_metadata,
+        internal_metadata=None,
         metadata=metadata,
         step=step,
         progress=progress,