truefoundry 0.4.3__py3-none-any.whl → 0.4.4__py3-none-any.whl

truefoundry/ml/autogen/client_README.md

@@ -78,7 +78,6 @@ Class | Method | HTTP request | Description
 *ExperimentsApi* | [**list_experiments_get**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#list_experiments_get) | **GET** /api/2.0/mlflow/experiments/list | List Experiments
 *ExperimentsApi* | [**put_privacy_type_put**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#put_privacy_type_put) | **PUT** /api/2.0/mlflow/experiments/privacy-type | Put Privacy Type
 *ExperimentsApi* | [**restore_experiment_post**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#restore_experiment_post) | **POST** /api/2.0/mlflow/experiments/restore | Restore Experiment
-*ExperimentsApi* | [**seed_list_experiment_get**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#seed_list_experiment_get) | **GET** /api/2.0/mlflow/experiments/seed/list | Seed List Experiment
 *ExperimentsApi* | [**set_experiment_tag_post**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#set_experiment_tag_post) | **POST** /api/2.0/mlflow/experiments/set-experiment-tag | Set Experiment Tag
 *ExperimentsApi* | [**update_experiment_post**](truefoundry/ml/autogen/client/docs/ExperimentsApi.md#update_experiment_post) | **POST** /api/2.0/mlflow/experiments/update | Update Experiment
 *HealthApi* | [**health_get**](truefoundry/ml/autogen/client/docs/HealthApi.md#health_get) | **GET** /health | Health
@@ -166,6 +165,7 @@ Class | Method | HTTP request | Description
  - [ArtifactType](truefoundry/ml/autogen/client/docs/ArtifactType.md)
  - [ArtifactVersionDto](truefoundry/ml/autogen/client/docs/ArtifactVersionDto.md)
  - [ArtifactVersionResponseDto](truefoundry/ml/autogen/client/docs/ArtifactVersionResponseDto.md)
+ - [ArtifactVersionSerializationFormat](truefoundry/ml/autogen/client/docs/ArtifactVersionSerializationFormat.md)
  - [ArtifactVersionStatus](truefoundry/ml/autogen/client/docs/ArtifactVersionStatus.md)
  - [AssistantMessage](truefoundry/ml/autogen/client/docs/AssistantMessage.md)
  - [AuthorizeUserForModelRequestDto](truefoundry/ml/autogen/client/docs/AuthorizeUserForModelRequestDto.md)
@@ -206,10 +206,12 @@ Class | Method | HTTP request | Description
  - [ExperimentIdRequestDto](truefoundry/ml/autogen/client/docs/ExperimentIdRequestDto.md)
  - [ExperimentResponseDto](truefoundry/ml/autogen/client/docs/ExperimentResponseDto.md)
  - [ExperimentTagDto](truefoundry/ml/autogen/client/docs/ExperimentTagDto.md)
+ - [ExternalModelSource](truefoundry/ml/autogen/client/docs/ExternalModelSource.md)
  - [FeatureDto](truefoundry/ml/autogen/client/docs/FeatureDto.md)
  - [FeatureValueType](truefoundry/ml/autogen/client/docs/FeatureValueType.md)
  - [FileInfoDto](truefoundry/ml/autogen/client/docs/FileInfoDto.md)
  - [FinalizeArtifactVersionRequestDto](truefoundry/ml/autogen/client/docs/FinalizeArtifactVersionRequestDto.md)
+ - [Framework](truefoundry/ml/autogen/client/docs/Framework.md)
  - [GetExperimentResponseDto](truefoundry/ml/autogen/client/docs/GetExperimentResponseDto.md)
  - [GetLatestRunLogResponseDto](truefoundry/ml/autogen/client/docs/GetLatestRunLogResponseDto.md)
  - [GetMetricHistoryResponse](truefoundry/ml/autogen/client/docs/GetMetricHistoryResponse.md)
@@ -227,6 +229,7 @@ Class | Method | HTTP request | Description
  - [ImageUrl](truefoundry/ml/autogen/client/docs/ImageUrl.md)
  - [InternalMetadata](truefoundry/ml/autogen/client/docs/InternalMetadata.md)
  - [LatestRunLogDto](truefoundry/ml/autogen/client/docs/LatestRunLogDto.md)
+ - [LibraryName](truefoundry/ml/autogen/client/docs/LibraryName.md)
  - [ListArtifactVersionsRequestDto](truefoundry/ml/autogen/client/docs/ListArtifactVersionsRequestDto.md)
  - [ListArtifactVersionsResponseDto](truefoundry/ml/autogen/client/docs/ListArtifactVersionsResponseDto.md)
  - [ListArtifactsRequestDto](truefoundry/ml/autogen/client/docs/ListArtifactsRequestDto.md)
@@ -248,7 +251,6 @@ Class | Method | HTTP request | Description
  - [ListModelsResponseDto](truefoundry/ml/autogen/client/docs/ListModelsResponseDto.md)
  - [ListRunArtifactsResponseDto](truefoundry/ml/autogen/client/docs/ListRunArtifactsResponseDto.md)
  - [ListRunLogsResponseDto](truefoundry/ml/autogen/client/docs/ListRunLogsResponseDto.md)
- - [ListSeedExperimentsResponseDto](truefoundry/ml/autogen/client/docs/ListSeedExperimentsResponseDto.md)
  - [LogBatchRequestDto](truefoundry/ml/autogen/client/docs/LogBatchRequestDto.md)
  - [LogMetricRequestDto](truefoundry/ml/autogen/client/docs/LogMetricRequestDto.md)
  - [LogParamRequestDto](truefoundry/ml/autogen/client/docs/LogParamRequestDto.md)
@@ -261,6 +263,7 @@ Class | Method | HTTP request | Description
  - [ModelResponseDto](truefoundry/ml/autogen/client/docs/ModelResponseDto.md)
  - [ModelSchemaDto](truefoundry/ml/autogen/client/docs/ModelSchemaDto.md)
  - [ModelVersionDto](truefoundry/ml/autogen/client/docs/ModelVersionDto.md)
+ - [ModelVersionManifest](truefoundry/ml/autogen/client/docs/ModelVersionManifest.md)
  - [ModelVersionResponseDto](truefoundry/ml/autogen/client/docs/ModelVersionResponseDto.md)
  - [MultiPartUploadDto](truefoundry/ml/autogen/client/docs/MultiPartUploadDto.md)
  - [MultiPartUploadResponseDto](truefoundry/ml/autogen/client/docs/MultiPartUploadResponseDto.md)
@@ -284,17 +287,24 @@ Class | Method | HTTP request | Description
  - [SetExperimentTagRequestDto](truefoundry/ml/autogen/client/docs/SetExperimentTagRequestDto.md)
  - [SetTagRequestDto](truefoundry/ml/autogen/client/docs/SetTagRequestDto.md)
  - [SignedURLDto](truefoundry/ml/autogen/client/docs/SignedURLDto.md)
+ - [Source](truefoundry/ml/autogen/client/docs/Source.md)
  - [Stop](truefoundry/ml/autogen/client/docs/Stop.md)
  - [StoreRunLogsRequestDto](truefoundry/ml/autogen/client/docs/StoreRunLogsRequestDto.md)
+ - [Subject](truefoundry/ml/autogen/client/docs/Subject.md)
+ - [SubjectType](truefoundry/ml/autogen/client/docs/SubjectType.md)
  - [SystemMessage](truefoundry/ml/autogen/client/docs/SystemMessage.md)
+ - [TensorflowFramework](truefoundry/ml/autogen/client/docs/TensorflowFramework.md)
  - [Text](truefoundry/ml/autogen/client/docs/Text.md)
  - [TextContentPart](truefoundry/ml/autogen/client/docs/TextContentPart.md)
+ - [TransformersFramework](truefoundry/ml/autogen/client/docs/TransformersFramework.md)
+ - [TruefoundryModelSource](truefoundry/ml/autogen/client/docs/TruefoundryModelSource.md)
  - [UpdateArtifactVersionRequestDto](truefoundry/ml/autogen/client/docs/UpdateArtifactVersionRequestDto.md)
  - [UpdateDatasetRequestDto](truefoundry/ml/autogen/client/docs/UpdateDatasetRequestDto.md)
  - [UpdateExperimentRequestDto](truefoundry/ml/autogen/client/docs/UpdateExperimentRequestDto.md)
  - [UpdateModelVersionRequestDto](truefoundry/ml/autogen/client/docs/UpdateModelVersionRequestDto.md)
  - [UpdateRunRequestDto](truefoundry/ml/autogen/client/docs/UpdateRunRequestDto.md)
  - [UpdateRunResponseDto](truefoundry/ml/autogen/client/docs/UpdateRunResponseDto.md)
+ - [UploadModelSource](truefoundry/ml/autogen/client/docs/UploadModelSource.md)
  - [Url](truefoundry/ml/autogen/client/docs/Url.md)
  - [UserMessage](truefoundry/ml/autogen/client/docs/UserMessage.md)
  - [ValidationError](truefoundry/ml/autogen/client/docs/ValidationError.md)

truefoundry/ml/autogen/entities/artifacts.py

@@ -1,13 +1,13 @@
 # generated by datamodel-codegen:
 #   filename:  artifacts.json
-#   timestamp: 2024-09-12T13:40:36+00:00
+#   timestamp: 2024-10-25T06:52:16+00:00
 
 from __future__ import annotations
 
 from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
-from truefoundry.pydantic_v1 import BaseModel, Field, constr
+from truefoundry.pydantic_v1 import BaseModel, Field, conint, constr
 
 
 class Agent(BaseModel):
@@ -49,6 +49,17 @@ class AgentWithFQN(Agent):
     fqn: str
 
 
+class BaseArtifactVersion(BaseModel):
+    description: Optional[constr(max_length=512)] = Field(
+        None,
+        description="+label=Description\n+docs=Description of the artifact version",
+    )
+    metadata: Dict[str, Any] = Field(
+        ...,
+        description="+label=Metadata\n+docs=Metadata for the model version\n+usage=Metadata for the model version\n+uiType=JsonInput",
+    )
+
+
 class MimeType(str, Enum):
     """
     +label=MIME Type
@@ -78,6 +89,51 @@ class BlobStorageReference(BaseModel):
     )
 
 
+class ExternalModelSource(BaseModel):
+    """
+    +label=External
+    """
+
+    type: Literal["external"] = Field(
+        ...,
+        description="+label=Type\n+usage=Type of the source\n+value=external\n+type=External",
+    )
+    uri: str = Field(..., description="+label=URI\n+usage=URI of the model source")
+
+
+class FastAIFramework(BaseModel):
+    """
+    +docs=FastAI framework for the model version
+    +label=FastAI
+    """
+
+    type: Literal["fastai"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=fastai"
+    )
+
+
+class GluonFramework(BaseModel):
+    """
+    +docs=Gluon framework for the model version
+    +label=Gluon
+    """
+
+    type: Literal["gluon"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=gluon"
+    )
+
+
+class H2OFramework(BaseModel):
+    """
+    +docs=H2O framework for the model version
+    +label=H2O
+    """
+
+    type: Literal["h2o"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=h2o"
+    )
+
+
 class ImageUrl(BaseModel):
     """
     +label=URL for the image
@@ -108,6 +164,28 @@ class ImageContentPart(BaseModel):
     )
 
 
+class KerasFramework(BaseModel):
+    """
+    +docs=Keras framework for the model version
+    +label=Keras
+    """
+
+    type: Literal["keras"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=keras"
+    )
+
+
+class LightGBMFramework(BaseModel):
+    """
+    +docs=LightGBM framework for the model version
+    +label=LightGBM
+    """
+
+    type: Literal["lightgbm"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=lightgbm"
+    )
+
+
 class Parameters(BaseModel):
     """
     +usage=Parameters for the provider
@@ -148,6 +226,61 @@ class ModelConfiguration(BaseModel):
     )
 
 
+class PaddleFramework(BaseModel):
+    """
+    +docs=PaddlePaddle framework for the model version
+    +label=Paddle
+    """
+
+    type: Literal["paddle"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=paddle"
+    )
+
+
+class PyTorchFramework(BaseModel):
+    """
+    +docs=PyTorch framework for the model version
+    +label=PyTorch
+    """
+
+    type: Literal["pytorch"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=pytorch"
+    )
+
+
+class SklearnFramework(BaseModel):
+    """
+    +docs=Scikit-learn framework for the model version
+    +label=Sklearn
+    """
+
+    type: Literal["sklearn"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
+    )
+
+
+class SpaCyFramework(BaseModel):
+    """
+    +docs=spaCy framework for the model version
+    +label=SpaCy
+    """
+
+    type: Literal["spacy"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=spacy"
+    )
+
+
+class StatsModelsFramework(BaseModel):
+    """
+    +docs=StatsModels framework for the model version
+    +label=StatsModels
+    """
+
+    type: Literal["statsmodels"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=statsmodels"
+    )
+
+
 class SystemMessage(BaseModel):
     """
     +usage=System message
@@ -164,6 +297,17 @@ class SystemMessage(BaseModel):
     )
 
 
+class TensorFlowFramework(BaseModel):
+    """
+    +docs=TensorFlow framework for the model version
+    +label=TensorFlow
+    """
+
+    type: Literal["tensorflow"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=tensorflow"
+    )
+
+
 class TextContentPart(BaseModel):
     """
     +usage=Text content
@@ -174,6 +318,53 @@ class TextContentPart(BaseModel):
     text: Union[constr(regex=r"^.[\s\S]*$"), BlobStorageReference]
 
 
+class LibraryName(str, Enum):
+    """
+    +label=Library Name
+    +usage=Name of the library for the framework
+    """
+
+    transformers = "transformers"
+    sentence_transformers = "sentence-transformers"
+    diffusers = "diffusers"
+
+
+class TransformersFramework(BaseModel):
+    """
+    +docs=Transformers framework for the model version
+    +label=Transformers
+    +value=transformers
+    """
+
+    type: Literal["transformers"] = Field(
+        ...,
+        description="+label=Type\n+usage=Type of the framework\n+value=transformers",
+    )
+    library_name: Optional[LibraryName] = Field(
+        "transformers",
+        description="+label=Library Name\n+usage=Name of the library for the framework",
+    )
+    pipeline_tag: Optional[str] = Field(
+        None,
+        description="+label=Pipeline Tag\n+usage=Pipeline tag\n+docs=Pipeline tag for the framework",
+    )
+
+
+class TruefoundryModelSource(BaseModel):
+    """
+    +usage=Source for the Model
+    +label=Upload
+    """
+
+    type: Literal["truefoundry"] = Field(
+        ...,
+        description="+usage=Type of the source\n+type=Upload\n+value=truefoundry\n+label=Upload",
+    )
+    uri: str = Field(
+        ..., description='+label=URI\n+uiType=UploadInput\n+uiProps={"hideClear":true}'
+    )
+
+
 class UserMessage(BaseModel):
     """
     +usage=User message
@@ -194,6 +385,17 @@ class UserMessage(BaseModel):
     )
 
 
+class XGBoostFramework(BaseModel):
+    """
+    +docs=XGBoost framework for the model version
+    +label=XGBoost
+    """
+
+    type: Literal["xgboost"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
+    )
+
+
 class AgentOpenAPITool(BaseModel):
     type: Literal["openapi-tool"] = Field(..., description="+value=openapi-tool")
     name: constr(regex=r"^[a-zA-Z][a-zA-Z0-9\-]{1,30}[a-zA-Z0-9]$") = Field(
@@ -275,6 +477,36 @@ class ChatPrompt(BasePrompt):
     )
 
 
+class ModelVersion(BaseArtifactVersion):
+    type: Literal["model-version"] = Field(
+        ..., description='+label=Type\n+usage=Model Version\n+value="model-version"'
+    )
+    source: Union[TruefoundryModelSource, ExternalModelSource] = Field(
+        ...,
+        description="+label=Model Source\n+usage=Source for the model version\n+uiType=Group",
+    )
+    framework: Optional[
+        Union[
+            TransformersFramework,
+            TensorFlowFramework,
+            SklearnFramework,
+            PyTorchFramework,
+            KerasFramework,
+            XGBoostFramework,
+            LightGBMFramework,
+            FastAIFramework,
+            H2OFramework,
+            SpaCyFramework,
+            StatsModelsFramework,
+            GluonFramework,
+            PaddleFramework,
+        ]
+    ] = Field(
+        None, description="+label=Framework\n+usage=Framework for the model version"
+    )
+    step: conint(ge=0) = Field(0, description="+label=Step")
+
+
 class AgentApp(BaseModel):
     type: Literal["agent-app"] = Field(..., description="+value=agent-app")
     tools: List[AgentOpenAPIToolWithFQN]
@@ -282,5 +514,5 @@ class AgentApp(BaseModel):
     root_agent: constr(min_length=1)
 
 
-class Artifacts(BaseModel):
-    __root__: Union[ChatPrompt, AgentOpenAPITool, Agent, AgentApp]
+class ArtifactsVersion(BaseModel):
+    __root__: Union[ChatPrompt, AgentOpenAPITool, Agent, AgentApp, ModelVersion]

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
@@ -184,11 +185,6 @@ class ArtifactVersion:
         _validate_artifact_metadata(value)
         self._metadata = value
 
-    @property
-    def created_by(self) -> str:
-        """Get the information about who created the artifact"""
-        return self._artifact_version.created_by
-
     @property
     def created_at(self) -> datetime.datetime:
         """Get the time at which artifact was created"""
@@ -402,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]:
@@ -414,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,
@@ -448,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
@@ -124,11 +125,6 @@ class DataDirectory:
         _validate_artifact_metadata(value)
         self._metadata = value
 
-    @property
-    def created_by(self) -> str:
-        """Get the information about who created the DataDirectory"""
-        return self._dataset.created_by
-
     @property
     def created_at(self) -> datetime.datetime:
         """Get the time at which DataDirectory was created"""
@@ -144,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`.
 
@@ -212,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,
@@ -224,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

@@ -40,8 +40,10 @@ from truefoundry.ml.log_types.artifacts.constants import (
 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.session import _get_api_client
 from truefoundry.pydantic_v1 import BaseModel, Extra
@@ -219,11 +221,6 @@ class ModelVersion:
             metrics_as_kv[metric.key] = metric.value
         return metrics_as_kv
 
-    @property
-    def created_by(self) -> str:
-        """Get the information about who created the model version"""
-        return self._model_version.created_by
-
     @property
     def created_at(self) -> datetime.datetime:
         """Get the time at which model version was created"""
@@ -393,24 +390,6 @@ class ModelVersion:
         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,
@@ -489,28 +468,29 @@ def _log_model_version(  # noqa: C901
         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)
+        _model_file_or_folder: Sequence[Tuple[str, str]] = [
+            (model_file_or_folder, MODEL_DIR_NAME.rstrip(os.sep) + os.sep),
         ]
-        _copy_additional_files(
+
+        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,
+            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(
+            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,
             )
-
     except Exception as e:
         temp_dir.cleanup()
         raise MlFoundryException("Failed to log model") from e
@@ -520,6 +500,7 @@ def _log_model_version(  # noqa: C901
     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
 
     # create entry
     _create_artifact_version_response = (
@@ -538,15 +519,17 @@ def _log_model_version(  # noqa: C901
         ),
         api_client=mlfoundry_artifacts_api.api_client,
     )
-    model_size = calculate_model_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 Total Size: %.6f MB",
-            model_size / 1000000.0,
+            "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, progress=progress
+        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(
@@ -567,7 +550,7 @@ def _log_model_version(  # noqa: C901
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
             id=version_id,
             run_uuid=run.run_id if run else None,
-            artifact_size=model_size,
+            artifact_size=total_size,
             internal_metadata=internal_metadata_dto,
             step=step if run else None,
         )