truefoundry 0.5.3rc3__py3-none-any.whl → 0.5.3rc5__py3-none-any.whl

truefoundry/deploy/v2/lib/patched_models.py

@@ -1,9 +1,11 @@
 import enum
+import os
 import re
 import warnings
 from collections.abc import Mapping
 from typing import Literal, Optional, Union
 
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.deploy.auto_gen import models
 from truefoundry.pydantic_v1 import (
     BaseModel,
@@ -71,6 +73,66 @@ resources:
 """
 
 
+AUTO_DISCOVERED_REQUIREMENTS_TXT_WARNING_MESSAGE_TEMPLATE = """\
+Using automatically discovered {requirements_txt_path} as the requirements file.
+This auto discovery behavior is deprecated and will be removed in a future release.
+Please specify the relative path of requirements file explicitly as
+
+```python
+build=Build(
+    ...
+    build_spec=PythonBuild(
+        ...
+        requirements_path={requirements_txt_path!r}
+    )
+)
+```
+
+OR
+
+```yaml
+build:
+  type: build
+  build_spec:
+    type: tfy-python-buildpack
+    requirements_path: {requirements_txt_path!r}
+    ...
+...
+```
+
+or set it to None if you don't want use any requirements file.
+"""
+
+
+def _resolve_requirements_path(
+    build_context_path: str,
+    requirements_path: Optional[str],
+) -> Optional[str]:
+    if requirements_path:
+        return requirements_path
+
+    # TODO: Deprecated behavior, phase out auto discovery in future release
+    possible_requirements_txt_filename = "requirements.txt"
+    possible_requirements_txt_path = os.path.join(
+        build_context_path, possible_requirements_txt_filename
+    )
+
+    if os.path.isfile(possible_requirements_txt_path):
+        requirements_txt_path = os.path.relpath(
+            possible_requirements_txt_path, start=build_context_path
+        )
+        warnings.warn(
+            AUTO_DISCOVERED_REQUIREMENTS_TXT_WARNING_MESSAGE_TEMPLATE.format(
+                requirements_txt_path=requirements_txt_path
+            ),
+            category=TrueFoundryDeprecationWarning,
+            stacklevel=2,
+        )
+        return requirements_txt_path
+
+    return None
+
+
 class CUDAVersion(str, enum.Enum):
     CUDA_11_0_CUDNN8 = "11.0-cudnn8"
     CUDA_11_1_CUDNN8 = "11.1-cudnn8"
@@ -139,7 +201,7 @@ class PythonBuild(models.PythonBuild, PatchedModelBase):
     type: Literal["tfy-python-buildpack"] = "tfy-python-buildpack"
 
     @root_validator
-    def validate_python_version_when_cuda_version(cls, values):
+    def validate_values(cls, values):
         if values.get("cuda_version"):
             python_version = values.get("python_version")
             if python_version and not re.match(r"^3\.\d+$", python_version):
@@ -148,6 +210,10 @@ class PythonBuild(models.PythonBuild, PatchedModelBase):
                     f"provided but got {python_version!r}. If you are adding a "
                     f'patch version, please remove it (e.g. "3.9.2" should be "3.9")'
                 )
+        _resolve_requirements_path(
+            build_context_path=values.get("build_context_path") or "./",
+            requirements_path=values.get("requirements_path"),
+        )
         return values
 
 
@@ -251,7 +317,7 @@ class Resources(models.Resources, PatchedModelBase):
                     gpu_type=gpu_type,
                     gpu_count=gpu_count,
                 ),
-                category=FutureWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
         elif gpu_count:
@@ -259,7 +325,7 @@ class Resources(models.Resources, PatchedModelBase):
                 LEGACY_GPU_COUNT_WARNING_MESSAGE_TEMPLATE.format(
                     gpu_count=gpu_count,
                 ),
-                category=FutureWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
         return values
@@ -295,7 +361,7 @@ class Autoscaling(ServiceAutoscaling):
             "`truefoundry.deploy.Autoscaling` is deprecated and will be removed in a future version. "
             "Please use `truefoundry.deploy.ServiceAutoscaling` instead. "
             "You can rename `Autoscaling` to `ServiceAutoscaling` in your script.",
-            category=DeprecationWarning,
+            category=TrueFoundryDeprecationWarning,
             stacklevel=2,
         )
         super().__init__(**kwargs)

truefoundry/deploy/v2/lib/source.py

@@ -8,6 +8,7 @@ from typing import Callable, List, Optional
 import gitignorefile
 from tqdm import tqdm
 
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.deploy import builder
 from truefoundry.deploy.auto_gen import models
 from truefoundry.deploy.builder.docker_service import (
@@ -102,7 +103,7 @@ def _get_callback_handler_to_ignore_file_path(
         warnings.warn(
             "`.sfyignore` is deprecated and will be ignored in future versions. "
             "Please rename the file to `.tfyignore`",
-            category=DeprecationWarning,
+            category=TrueFoundryDeprecationWarning,
             stacklevel=2,
         )
         return gitignorefile.parse(path=ignorefile_path, base_path=source_dir)

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -1,5 +1,3 @@
-import math
-import mmap
 import os
 import posixpath
 import sys
@@ -14,7 +12,6 @@ from typing import (
     Dict,
     Iterator,
     List,
-    NamedTuple,
     Optional,
     Sequence,
     Tuple,
@@ -23,7 +20,6 @@ from typing import (
 from urllib.parse import unquote
 from urllib.request import pathname2url
 
-import requests
 from rich.console import _is_jupyter
 from rich.progress import (
     BarColumn,
@@ -36,6 +32,17 @@ from rich.progress import (
 from tqdm.utils import CallbackIOWrapper
 
 from truefoundry.common.constants import ENV_VARS
+from truefoundry.common.request_utils import (
+    augmented_raise_for_status,
+    cloud_storage_http_request,
+)
+from truefoundry.common.storage_provider_utils import (
+    MultiPartUpload,
+    _FileMultiPartInfo,
+    azure_multi_part_upload,
+    decide_file_parts,
+    s3_compatible_multipart_upload,
+)
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
     CreateMultiPartUploadForDatasetRequestDto,
@@ -56,34 +63,11 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     RunArtifactsApi,
     SignedURLDto,
 )
-from truefoundry.ml.clients.utils import (
-    augmented_raise_for_status,
-    cloud_storage_http_request,
-)
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
 from truefoundry.ml.session import _get_api_client
 from truefoundry.pydantic_v1 import BaseModel, root_validator
 
-_MIN_BYTES_REQUIRED_FOR_MULTIPART = 100 * 1024 * 1024
-# GCP/S3 Maximum number of parts per upload	10,000
-# Maximum number of blocks in a block blob 50,000 blocks
-# TODO: This number is artificially limited now. Later
-# we will ask for parts signed URI in batches rather than in a single
-# API Calls:
-# Create Multipart Upload (Returns maximum number of parts, size limit of
-#                            a single part, upload id for s3 etc )
-#   Get me signed uris for first 500 parts
-#     Upload 500 parts
-#   Get me signed uris for the next 500 parts
-#     Upload 500 parts
-#   ...
-# Finalize the Multipart upload using the finalize signed url returned
-# by Create Multipart Upload or get a new one.
-_MAX_NUM_PARTS_FOR_MULTIPART = 1000
-# Azure Maximum size of a block in a block blob	4000 MiB
-# GCP/S3 Maximum size of an individual part in a multipart upload 5 GiB
-_MAX_PART_SIZE_BYTES_FOR_MULTIPART = 4 * 1024 * 1024 * 1000
 _LIST_FILES_PAGE_SIZE = 500
 _GENERATE_SIGNED_URL_BATCH_SIZE = 50
 DEFAULT_PRESIGNED_URL_EXPIRY_TIME = 3600
@@ -123,21 +107,6 @@ def relative_path_to_artifact_path(path):
     return unquote(pathname2url(path))
 
 
-def _align_part_size_with_mmap_allocation_granularity(part_size: int) -> int:
-    modulo = part_size % mmap.ALLOCATIONGRANULARITY
-    if modulo == 0:
-        return part_size
-
-    part_size += mmap.ALLOCATIONGRANULARITY - modulo
-    return part_size
-
-
-# Can not be less than 5 * 1024 * 1024
-_PART_SIZE_BYTES_FOR_MULTIPART = _align_part_size_with_mmap_allocation_granularity(
-    10 * 1024 * 1024
-)
-
-
 def bad_path_message(name):
     return (
         "Names may be treated as files in certain cases, and must not resolve to other names"
@@ -157,68 +126,6 @@ def verify_artifact_path(artifact_path):
         )
 
 
-class _PartNumberEtag(NamedTuple):
-    part_number: int
-    etag: str
-
-
-def _get_s3_compatible_completion_body(multi_parts: List[_PartNumberEtag]) -> str:
-    body = "<CompleteMultipartUpload>\n"
-    for part in multi_parts:
-        body += "  <Part>\n"
-        body += f"    <PartNumber>{part.part_number}</PartNumber>\n"
-        body += f"    <ETag>{part.etag}</ETag>\n"
-        body += "  </Part>\n"
-    body += "</CompleteMultipartUpload>"
-    return body
-
-
-def _get_azure_blob_completion_body(block_ids: List[str]) -> str:
-    body = "<BlockList>\n"
-    for block_id in block_ids:
-        body += f"<Uncommitted>{block_id}</Uncommitted> "
-    body += "</BlockList>"
-    return body
-
-
-class _FileMultiPartInfo(NamedTuple):
-    num_parts: int
-    part_size: int
-    file_size: int
-
-
-def _decide_file_parts(file_path: str) -> _FileMultiPartInfo:
-    file_size = os.path.getsize(file_path)
-    if (
-        file_size < _MIN_BYTES_REQUIRED_FOR_MULTIPART
-        or ENV_VARS.TFY_ARTIFACTS_DISABLE_MULTIPART_UPLOAD
-    ):
-        return _FileMultiPartInfo(1, part_size=file_size, file_size=file_size)
-
-    ideal_num_parts = math.ceil(file_size / _PART_SIZE_BYTES_FOR_MULTIPART)
-    if ideal_num_parts <= _MAX_NUM_PARTS_FOR_MULTIPART:
-        return _FileMultiPartInfo(
-            ideal_num_parts,
-            part_size=_PART_SIZE_BYTES_FOR_MULTIPART,
-            file_size=file_size,
-        )
-
-    part_size_when_using_max_parts = math.ceil(file_size / _MAX_NUM_PARTS_FOR_MULTIPART)
-    part_size_when_using_max_parts = _align_part_size_with_mmap_allocation_granularity(
-        part_size_when_using_max_parts
-    )
-    if part_size_when_using_max_parts > _MAX_PART_SIZE_BYTES_FOR_MULTIPART:
-        raise ValueError(
-            f"file {file_path!r} is too big for upload. Multipart chunk"
-            f" size {part_size_when_using_max_parts} is higher"
-            f" than {_MAX_PART_SIZE_BYTES_FOR_MULTIPART}"
-        )
-    num_parts = math.ceil(file_size / part_size_when_using_max_parts)
-    return _FileMultiPartInfo(
-        num_parts, part_size=part_size_when_using_max_parts, file_size=file_size
-    )
-
-
 def _signed_url_upload_file(
     signed_url: SignedURLDto,
     local_file: str,
@@ -227,9 +134,12 @@ def _signed_url_upload_file(
 ):
     if os.stat(local_file).st_size == 0:
         with cloud_storage_http_request(
-            method="put", url=signed_url.signed_url, data=""
+            method="put",
+            url=signed_url.signed_url,
+            data="",
+            exception_class=MlFoundryException,  # type: ignore
         ) as response:
-            augmented_raise_for_status(response)
+            augmented_raise_for_status(response, exception_class=MlFoundryException)  # type: ignore
         return
 
     task_progress_bar = progress_bar.add_task(
@@ -247,9 +157,12 @@ def _signed_url_upload_file(
         # NOTE: Azure Put Blob does not support Transfer Encoding header.
         wrapped_file = CallbackIOWrapper(callback, file, "read")
         with cloud_storage_http_request(
-            method="put", url=signed_url.signed_url, data=wrapped_file
+            method="put",
+            url=signed_url.signed_url,
+            data=wrapped_file,
+            exception_class=MlFoundryException,  # type: ignore
         ) as response:
-            augmented_raise_for_status(response)
+            augmented_raise_for_status(response, exception_class=MlFoundryException)  # type: ignore
 
 
 def _download_file_using_http_uri(
@@ -265,9 +178,12 @@ def _download_file_using_http_uri(
             providers.
     """
     with cloud_storage_http_request(
-        method="get", url=http_uri, stream=True
+        method="get",
+        url=http_uri,
+        stream=True,
+        exception_class=MlFoundryException,  # type: ignore
     ) as response:
-        augmented_raise_for_status(response)
+        augmented_raise_for_status(response, exception_class=MlFoundryException)  # type: ignore
         file_size = int(response.headers.get("Content-Length", 0))
         with open(download_path, "wb") as output_file:
             for chunk in response.iter_content(chunk_size=chunk_size):
@@ -281,188 +197,6 @@ def _download_file_using_http_uri(
                     os.fsync(output_file.fileno())
 
 
-class _CallbackIOWrapperForMultiPartUpload(CallbackIOWrapper):
-    def __init__(self, callback, stream, method, length: int):
-        self.wrapper_setattr("_length", length)
-        super().__init__(callback, stream, method)
-
-    def __len__(self):
-        return self.wrapper_getattr("_length")
-
-
-def _file_part_upload(
-    url: str,
-    file_path: str,
-    seek: int,
-    length: int,
-    file_size: int,
-    abort_event: Optional[Event] = None,
-    method: str = "put",
-):
-    def callback(*_, **__):
-        if abort_event and abort_event.is_set():
-            raise Exception("aborting upload")
-
-    with open(file_path, "rb") as file:
-        with mmap.mmap(
-            file.fileno(),
-            length=min(file_size - seek, length),
-            offset=seek,
-            access=mmap.ACCESS_READ,
-        ) as mapped_file:
-            wrapped_file = _CallbackIOWrapperForMultiPartUpload(
-                callback, mapped_file, "read", len(mapped_file)
-            )
-            with cloud_storage_http_request(
-                method=method,
-                url=url,
-                data=wrapped_file,
-            ) as response:
-                augmented_raise_for_status(response)
-                return response
-
-
-def _s3_compatible_multipart_upload(
-    multipart_upload: MultiPartUploadDto,
-    local_file: str,
-    multipart_info: _FileMultiPartInfo,
-    executor: ThreadPoolExecutor,
-    progress_bar: Progress,
-    abort_event: Optional[Event] = None,
-):
-    abort_event = abort_event or Event()
-    parts = []
-
-    multi_part_upload_progress = progress_bar.add_task(
-        f"[green]Uploading {local_file}:", start=True
-    )
-
-    def upload(part_number: int, seek: int) -> None:
-        logger.debug(
-            "Uploading part %d/%d of %s",
-            part_number,
-            multipart_info.num_parts,
-            local_file,
-        )
-        response = _file_part_upload(
-            url=multipart_upload.part_signed_urls[part_number].signed_url,
-            file_path=local_file,
-            seek=seek,
-            length=multipart_info.part_size,
-            file_size=multipart_info.file_size,
-            abort_event=abort_event,
-        )
-        logger.debug(
-            "Uploaded part %d/%d of %s",
-            part_number,
-            multipart_info.num_parts,
-            local_file,
-        )
-        progress_bar.update(
-            multi_part_upload_progress,
-            advance=multipart_info.part_size,
-            total=multipart_info.file_size,
-        )
-        etag = response.headers["ETag"]
-        parts.append(_PartNumberEtag(etag=etag, part_number=part_number + 1))
-
-    futures: List[Future] = []
-    for part_number, seek in enumerate(
-        range(0, multipart_info.file_size, multipart_info.part_size)
-    ):
-        future = executor.submit(upload, part_number=part_number, seek=seek)
-        futures.append(future)
-
-    done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
-    if len(not_done) > 0:
-        abort_event.set()
-    for future in not_done:
-        future.cancel()
-    for future in done:
-        if future.exception() is not None:
-            raise future.exception()
-
-    logger.debug("Finalizing multipart upload of %s", local_file)
-    parts = sorted(parts, key=lambda part: part.part_number)
-    response = requests.post(
-        multipart_upload.finalize_signed_url.signed_url,
-        data=_get_s3_compatible_completion_body(parts),
-        timeout=2 * 60,
-    )
-    response.raise_for_status()
-    logger.debug("Multipart upload of %s completed", local_file)
-
-
-def _azure_multi_part_upload(
-    multipart_upload: MultiPartUploadDto,
-    local_file: str,
-    multipart_info: _FileMultiPartInfo,
-    executor: ThreadPoolExecutor,
-    progress_bar: Progress,
-    abort_event: Optional[Event] = None,
-):
-    abort_event = abort_event or Event()
-
-    multi_part_upload_progress = progress_bar.add_task(
-        f"[green]Uploading {local_file}:", start=True
-    )
-
-    def upload(part_number: int, seek: int):
-        logger.debug(
-            "Uploading part %d/%d of %s",
-            part_number,
-            multipart_info.num_parts,
-            local_file,
-        )
-        _file_part_upload(
-            url=multipart_upload.part_signed_urls[part_number].signed_url,
-            file_path=local_file,
-            seek=seek,
-            length=multipart_info.part_size,
-            file_size=multipart_info.file_size,
-            abort_event=abort_event,
-        )
-        progress_bar.update(
-            multi_part_upload_progress,
-            advance=multipart_info.part_size,
-            total=multipart_info.file_size,
-        )
-        logger.debug(
-            "Uploaded part %d/%d of %s",
-            part_number,
-            multipart_info.num_parts,
-            local_file,
-        )
-
-    futures: List[Future] = []
-    for part_number, seek in enumerate(
-        range(0, multipart_info.file_size, multipart_info.part_size)
-    ):
-        future = executor.submit(upload, part_number=part_number, seek=seek)
-        futures.append(future)
-
-    done, not_done = wait(futures, return_when=FIRST_EXCEPTION)
-    if len(not_done) > 0:
-        abort_event.set()
-    for future in not_done:
-        future.cancel()
-    for future in done:
-        if future.exception() is not None:
-            raise future.exception()
-
-    logger.debug("Finalizing multipart upload of %s", local_file)
-    if multipart_upload.azure_blob_block_ids:
-        response = requests.put(
-            multipart_upload.finalize_signed_url.signed_url,
-            data=_get_azure_blob_completion_body(
-                block_ids=multipart_upload.azure_blob_block_ids
-            ),
-            timeout=2 * 60,
-        )
-        response.raise_for_status()
-    logger.debug("Multipart upload of %s completed", local_file)
-
-
 def _any_future_has_failed(futures) -> bool:
     return any(
         future.done() and not future.cancelled() and future.exception() is not None
@@ -638,25 +372,27 @@ class MlFoundryArtifactsRepository:
             multipart_upload.storage_provider
             is MultiPartUploadStorageProvider.S3_COMPATIBLE
         ):
-            _s3_compatible_multipart_upload(
-                multipart_upload=multipart_upload,
+            s3_compatible_multipart_upload(
+                multipart_upload=MultiPartUpload.parse_obj(multipart_upload.to_dict()),
                 local_file=local_file,
                 executor=executor,
                 multipart_info=multipart_info,
                 abort_event=abort_event,
                 progress_bar=progress_bar,
+                exception_class=MlFoundryException,  # type: ignore
             )
         elif (
             multipart_upload.storage_provider
             is MultiPartUploadStorageProvider.AZURE_BLOB
         ):
-            _azure_multi_part_upload(
-                multipart_upload=multipart_upload,
+            azure_multi_part_upload(
+                multipart_upload=MultiPartUpload.parse_obj(multipart_upload.to_dict()),
                 local_file=local_file,
                 executor=executor,
                 multipart_info=multipart_info,
                 abort_event=abort_event,
                 progress_bar=progress_bar,
+                exception_class=MlFoundryException,  # type: ignore
             )
         else:
             raise NotImplementedError()
@@ -793,7 +529,7 @@ class MlFoundryArtifactsRepository:
             )
         upload_path = artifact_path
         upload_path = upload_path.lstrip(posixpath.sep)
-        multipart_info = _decide_file_parts(local_file)
+        multipart_info = decide_file_parts(local_file)
         if multipart_info.num_parts == 1:
             files_for_normal_upload.append((upload_path, local_file, multipart_info))
         else:

truefoundry/ml/clients/servicefoundry_client.py

@@ -1,36 +1,57 @@
-from typing import Optional
+import functools
 
 from truefoundry.common.constants import (
     SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
     VERSION_PREFIX,
 )
+from truefoundry.common.exceptions import HttpRequestException
+from truefoundry.common.request_utils import (
+    http_request,
+    request_handling,
+    requests_retry_session,
+)
 from truefoundry.common.servicefoundry_client import (
     ServiceFoundryServiceClient as BaseServiceFoundryServiceClient,
 )
-from truefoundry.ml.clients.entities import (
-    HostCreds,
-)
-from truefoundry.ml.clients.utils import http_request_safe
 from truefoundry.ml.exceptions import MlFoundryException
 
 
 class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
-    # TODO (chiragjn): Rename tracking_uri to tfy_host
-    def __init__(self, tracking_uri: str, token: Optional[str] = None):
-        super().__init__(base_url=tracking_uri)
-        self.host_creds = HostCreds(host=self._api_server_url, token=token)
+    def __init__(self, tfy_host: str, token: str):
+        super().__init__(tfy_host=tfy_host)
+        self._token = token
+
+    @functools.cached_property
+    def _min_cli_version_required(self) -> str:
+        # TODO (chiragjn): read the mlfoundry min cli version from the config?
+        return self.python_sdk_config.truefoundry_cli_min_version
 
     def get_integration_from_id(self, integration_id: str):
         integration_id = integration_id or ""
-        response = http_request_safe(
-            host_creds=self.host_creds,
-            endpoint=f"{VERSION_PREFIX}/provider-accounts/provider-integrations",
-            params={"id": integration_id, "type": "blob-storage"},
+        session = requests_retry_session(retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES)
+        response = http_request(
             method="get",
+            url=f"{self._api_server_url}/{VERSION_PREFIX}/provider-accounts/provider-integrations",
+            token=self._token,
             timeout=3,
-            max_retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
+            params={"id": integration_id, "type": "blob-storage"},
+            session=session,
         )
-        data = response.json()
+
+        try:
+            data = request_handling(response)
+            assert isinstance(data, dict)
+        except HttpRequestException as he:
+            raise MlFoundryException(
+                f"Failed to get storage integration from id: {integration_id}. Error: {he.message}",
+                status_code=he.status_code,
+            ) from None
+        except Exception as e:
+            raise MlFoundryException(
+                f"Failed to get storage integration from id: {integration_id}. Error: {str(e)}"
+            ) from None
+
+        # TODO (chiragjn): Parse this using Pydantic
         if (
             data.get("providerIntegrations")
             and len(data["providerIntegrations"]) > 0

truefoundry/ml/exceptions.py

@@ -1,8 +1,9 @@
 from typing import Optional
 
 
+# TODO (chiragjn): We need to establish uniform exception handling across codebase
 class MlFoundryException(Exception):
-    def __init__(self, message, status_code: Optional[int] = None):
+    def __init__(self, message: str, status_code: Optional[int] = None):
         self.message = str(message)
         self.status_code = status_code
         super().__init__(message)

truefoundry/ml/log_types/artifacts/artifact.py

@@ -9,6 +9,7 @@ import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Union
 
+from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml.artifact.truefoundry_artifact_repo import (
     ArtifactIdentifier,
     MlFoundryArtifactsRepository,
@@ -217,7 +218,7 @@ class ArtifactVersion:
         if not self._artifact_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. tags do not exist, returning empty list",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
         return self._tags
@@ -230,7 +231,7 @@ class ArtifactVersion:
         if not self._artifact_version.manifest:
             warnings.warn(
                 message="This model version was created using an older serialization format. Tags will not be updated",
-                category=DeprecationWarning,
+                category=TrueFoundryDeprecationWarning,
                 stacklevel=2,
             )
             return