Repository: bentoml/BentoML
Pull request #4988 (2024-09-23 09:49:19)

refactor: unify same schema
Signed-off-by: Frost Ming <me@frostming.com>

## What does this PR address?

<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a few steps ahead.

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
  - style: (formatting, missing semicolons, etc; no production code change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
  - test: (adding missing tests, refactoring tests; no production code change)
  - chore: (updating grunt tasks etc; no production code change)
  - build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community member will help review your PR (see "Who can help review?" section for potential reviewers.). If no one has reviewed your PR after a week have passed, don't hesitate to post a new comment and ping @-the same person. Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->

Fixes #(issue)

## Before submitting:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->

- [ ] Does the Pull Request follow [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary) naming? Here are [GitHub's
      guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request) on how to create a pull request.
- [ ] Does the code follow BentoML's code style, `pre-commit run -a` script has passed ([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [ ] Did you read through [contribution guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute) and follow [development guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [ ] Did your changes require updates to the documentation? Have you updated
      those accordingly? Here are [documentation guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips on writting docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?



Patch:
diff --git a/docs/source/guides/model-composition.rst b/docs/source/guides/model-composition.rst
index 3c30af85..d6ead848 100644
--- a/docs/source/guides/model-composition.rst
+++ b/docs/source/guides/model-composition.rst
@@ -270,7 +270,7 @@ This example runs two independent models concurrently to generate different type
 Inference graph
 """""""""""""""
 
-The following is a complex example of model composition within BentoML, including both parallel and sequential processing steps to for an advanced inference workflow.
+The following is a complex example of model composition within BentoML, including both parallel and sequential processing steps for an advanced inference workflow.
 
 .. code-block:: python
 
diff --git a/src/_bentoml_sdk/models/base.py b/src/_bentoml_sdk/models/base.py
index 351eeeac..f8c8fcd9 100644
--- a/src/_bentoml_sdk/models/base.py
+++ b/src/_bentoml_sdk/models/base.py
@@ -9,7 +9,6 @@ from simple_di import Provide
 from simple_di import inject
 
 from bentoml._internal.bento.bento import BentoModelInfo
-from bentoml._internal.cloud import BentoCloudClient
 from bentoml._internal.cloud.client import RestApiClient
 from bentoml._internal.cloud.schemas.modelschemas import LabelItemSchema
 from bentoml._internal.cloud.schemas.modelschemas import ModelManifestSchema
@@ -120,7 +119,6 @@ class BentoModel(Model[StoredModel]):
         self,
         base_path: t.Union[PathType, FS, None] = None,
         global_model_store: ModelStore = Provide[BentoMLContainer.model_store],
-        cloud_client: BentoCloudClient = Provide[BentoMLContainer.bentocloud_client],
     ) -> StoredModel:
         stored = self.stored
         if base_path is not None:
@@ -135,6 +133,7 @@ class BentoModel(Model[StoredModel]):
                     target_model_store=model_store,
                 )
             return stored
+        cloud_client = BentoMLContainer.bentocloud_client.get()
         model = cloud_client.model.pull(self.tag, model_store=model_store)
         assert model is not None, "non-bentoml model"
         return model
diff --git a/src/bentoml/_internal/cloud/bento.py b/src/bentoml/_internal/cloud/bento.py
index 08125c92..82df2006 100644
--- a/src/bentoml/_internal/cloud/bento.py
+++ b/src/bentoml/_internal/cloud/bento.py
@@ -27,14 +27,14 @@ from .model import ModelAPI
 from .schemas.modelschemas import BentoApiSchema
 from .schemas.modelschemas import BentoRunnerResourceSchema
 from .schemas.modelschemas import BentoRunnerSchema
-from .schemas.modelschemas import BentoUploadStatus
+from .schemas.modelschemas import UploadStatus
 from .schemas.schemasv1 import BentoManifestSchema
 from .schemas.schemasv1 import BentoSchema
 from .schemas.schemasv1 import CompleteMultipartUploadSchema
 from .schemas.schemasv1 import CompletePartSchema
 from .schemas.schemasv1 import CreateBentoRepositorySchema
 from .schemas.schemasv1 import CreateBentoSchema
-from .schemas.schemasv1 import FinishUploadBentoSchema
+from .schemas.schemasv1 import FinishUploadSchema
 from .schemas.schemasv1 import LabelItemSchema
 from .schemas.schemasv1 import PreSignMultipartUploadUrlSchema
 from .schemas.schemasv1 import TransmissionStrategy
@@ -146,7 +146,7 @@ class BentoAPI:
         if (
             not force
             and remote_bento
-            and remote_bento.upload_status == BentoUploadStatus.SUCCESS.value
+            and remote_bento.upload_status == UploadStatus.SUCCESS.value
         ):
             self.spinner.log(
                 f'[bold blue]Push skipped: Bento "{bento.tag}" already exists in remote Bento store'
@@ -218,8 +218,8 @@ class BentoAPI:
                 rest_client.v1.finish_upload_bento(
                     bento_repository_name=bento_repository.name,
                     version=version,
-                    req=FinishUploadBentoSchema(
-                        status=BentoUploadStatus.SUCCESS.value,
+                    req=FinishUploadSchema(
+                        status=UploadStatus.SUCCESS.value,
                         reason="bento for development",
                     ),
                 )
@@ -286,8 +286,8 @@ class BentoAPI:
                     raise e
                 self.spinner.log(f'[bold green]Successfully pushed bento "{bento.tag}"')
                 return
-            finish_req = FinishUploadBentoSchema(
-                status=BentoUploadStatus.SUCCESS.value, reason=""
+            finish_req = FinishUploadSchema(
+                status=UploadStatus.SUCCESS.value, reason=""
             )
             try:
                 if presigned_upload_url is not None:
@@ -295,8 +295,8 @@ class BentoAPI:
                         presigned_upload_url, content=io_with_cb, timeout=36000
                     )
                     if resp.status_code != 200:
-                        finish_req = FinishUploadBentoSchema(
-                            status=BentoUploadStatus.FAILED.value,
+                        finish_req = FinishUploadSchema(
+                            status=UploadStatus.FAILED.value,
                             reason=resp.text,
                         )
                 else:
@@ -319,7 +319,7 @@ class BentoAPI:
 
                     def chunk_upload(
                         upload_id: str, chunk_number: int
-                    ) -> FinishUploadBentoSchema | tuple[str, int]:
+                    ) -> FinishUploadSchema | tuple[str, int]:
                         with self.spinner.spin(
                             text=f'({chunk_number}/{chunks_count}) Presign multipart upload url of Bento "{bento.tag}"...'
                         ):
@@ -355,8 +355,8 @@ class BentoAPI:
                                     if resp.status_code == 200:
                                         break
                                     if i == UPLOAD_RETRY_COUNT - 1:
-                                        return FinishUploadBentoSchema(
-                                            status=BentoUploadStatus.FAILED.value,
+                                        return FinishUploadSchema(
+                                            status=UploadStatus.FAILED.value,
                                             reason=resp.text,
                                         )
                                     else:  # retry and reset and update progress
@@ -366,9 +366,7 @@ class BentoAPI:
                                         )
                                 return resp.headers["ETag"], chunk_number
 
-                    futures_: list[
-                        Future[FinishUploadBentoSchema | tuple[str, int]]
-                    ] = []
+                    futures_: list[Future[FinishUploadSchema | tuple[str, int]]] = []
 
                     with ThreadPoolExecutor(
                         max_workers=min(max(chunks_count, 1), threads)
@@ -385,7 +383,7 @@ class BentoAPI:
 
                     for future in futures_:
                         result = future.result()
-                        if isinstance(result, FinishUploadBentoSchema):
+                        if isinstance(result, FinishUploadSchema):
                             finish_req = result
                             break
                         else:
@@ -410,11 +408,11 @@ class BentoAPI:
                         )
 
             except Exception as e:  # pylint: disable=broad-except
-                finish_req = FinishUploadBentoSchema(
-                    status=BentoUploadStatus.FAILED.value,
+                finish_req = FinishUploadSchema(
+                    status=UploadStatus.FAILED.value,
                     reason=str(e),
                 )
-            if finish_req.status == BentoUploadStatus.FAILED.value:
+            if finish_req.status == UploadStatus.FAILED.value:
                 self.spinner.log(f'[bold red]Failed to upload Bento "{bento.tag}"')
             with self.spinner.spin(
                 text="Submitting upload status to remote Bento store"
@@ -424,7 +422,7 @@ class BentoAPI:
                     version=version,
                     req=finish_req,
                 )
-            if finish_req.status != BentoUploadStatus.SUCCESS.value:
+            if finish_req.status != UploadStatus.SUCCESS.value:
                 self.spinner.log(
                     f'[bold red]Failed pushing Bento "{bento.tag}": {finish_req.reason}'
                 )
diff --git a/src/bentoml/_internal/cloud/client.py b/src/bentoml/_internal/cloud/client.py
index f284fa27..e920be19 100644
--- a/src/bentoml/_internal/cloud/client.py
+++ b/src/bentoml/_internal/cloud/client.py
@@ -31,8 +31,7 @@ from .schemas.schemasv1 import CreateModelSchema
 from .schemas.schemasv1 import CreateSecretSchema
 from .schemas.schemasv1 import DeploymentFullSchema
 from .schemas.schemasv1 import DeploymentListSchema
-from .schemas.schemasv1 import FinishUploadBentoSchema
-from .schemas.schemasv1 import FinishUploadModelSchema
+from .schemas.schemasv1 import FinishUploadSchema
 from .schemas.schemasv1 import ModelRepositorySchema
 from .schemas.schemasv1 import ModelSchema
 from .schemas.schemasv1 import ModelWithRepositoryListSchema
@@ -222,7 +221,7 @@ class RestApiClientV1(BaseRestApiClient):
         return schema_from_object(resp.json(), BentoSchema)
 
     def finish_upload_bento(
-        self, bento_repository_name: str, version: str, req: FinishUploadBentoSchema
+        self, bento_repository_name: str, version: str, req: FinishUploadSchema
     ) -> BentoSchema:
         url = f"/api/v1/bento_repositories/{bento_repository_name}/bentos/{version}/finish_upload"
         resp = self.session.patch(url, json=schema_to_object(req))
@@ -341,7 +340,7 @@ class RestApiClientV1(BaseRestApiClient):
         return schema_from_object(resp.json(), ModelSchema)
 
     def finish_upload_model(
-        self, model_repository_name: str, version: str, req: FinishUploadModelSchema
+        self, model_repository_name: str, version: str, req: FinishUploadSchema
     ) -> ModelSchema:
         url = f"/api/v1/model_repositories/{model_repository_name}/models/{version}/finish_upload"
         resp = self.session.patch(url, json=schema_to_object(req))
diff --git a/src/bentoml/_internal/cloud/model.py b/src/bentoml/_internal/cloud/model.py
index dd9b81b3..e9c5dcbe 100644
--- a/src/bentoml/_internal/cloud/model.py
+++ b/src/bentoml/_internal/cloud/model.py
@@ -24,11 +24,11 @@ from .base import FILE_CHUNK_SIZE
 from .base import UPLOAD_RETRY_COUNT
 from .base import CallbackIOWrapper
 from .base import Spinner
-from .schemas.modelschemas import ModelUploadStatus
+from .schemas.modelschemas import UploadStatus
 from .schemas.schemasv1 import CompleteMultipartUploadSchema
 from .schemas.schemasv1 import CompletePartSchema
 from .schemas.schemasv1 import CreateModelRepositorySchema
-from .schemas.schemasv1 import FinishUploadModelSchema
+from .schemas.schemasv1 import FinishUploadSchema
 from .schemas.schemasv1 import PreSignMultipartUploadUrlSchema
 from .schemas.schemasv1 import TransmissionStrategy
 
@@ -108,7 +108,7 @@ class ModelAPI:
         if (
             not force
             and remote_model
-            and remote_model.upload_status == ModelUploadStatus.SUCCESS.value
+            and remote_model.upload_status == UploadStatus.SUCCESS.value
         ):
             self.spinner.log(
                 f'[bold blue]Model "{model}" already exists in remote model store, skipping'
@@ -179,8 +179,8 @@ class ModelAPI:
                     raise e
                 self.spinner.log(f'[bold green]Successfully pushed model "{model.tag}"')
                 return
-            finish_req = FinishUploadModelSchema(
-                status=ModelUploadStatus.SUCCESS.value, reason=""
+            finish_req = FinishUploadSchema(
+                status=UploadStatus.SUCCESS.value, reason=""
             )
             try:
                 if presigned_upload_url is not None:
@@ -188,8 +188,8 @@ class ModelAPI:
                         presigned_upload_url, content=io_with_cb, timeout=36000
                     )
                     if resp.status_code != 200:
-                        finish_req = FinishUploadModelSchema(
-                            status=ModelUploadStatus.FAILED.value,
+                        finish_req = FinishUploadSchema(
+                            status=UploadStatus.FAILED.value,
                             reason=resp.text,
                         )
                 else:
@@ -212,7 +212,7 @@ class ModelAPI:
 
                     def chunk_upload(
                         upload_id: str, chunk_number: int
-                    ) -> FinishUploadModelSchema | tuple[str, int]:
+                    ) -> FinishUploadSchema | tuple[str, int]:
                         with self.spinner.spin(
                             text=f'({chunk_number}/{chunks_count}) Presign multipart upload url of model "{model.tag}"...'
                         ):
@@ -249,8 +249,8 @@ class ModelAPI:
                                     if resp.status_code == 200:
                                         break
                                     if i == UPLOAD_RETRY_COUNT - 1:
-                                        return FinishUploadModelSchema(
-                                            status=ModelUploadStatus.FAILED.value,
+                                        return FinishUploadSchema(
+                                            status=UploadStatus.FAILED.value,
                                             reason=resp.text,
                                         )
                                     else:  # retry and reset and update progress
@@ -260,9 +260,7 @@ class ModelAPI:
                                         )
                                 return resp.headers["ETag"], chunk_number
 
-                    futures_: list[
-                        Future[FinishUploadModelSchema | tuple[str, int]]
-                    ] = []
+                    futures_: list[Future[FinishUploadSchema | tuple[str, int]]] = []
 
                     with ThreadPoolExecutor(
                         max_workers=min(max(chunks_count, 1), threads)
@@ -279,7 +277,7 @@ class ModelAPI:
 
                     for future in futures_:
                         result = future.result()
-                        if isinstance(result, FinishUploadModelSchema):
+                        if isinstance(result, FinishUploadSchema):
                             finish_req = result
                             break
                         else:
@@ -304,11 +302,11 @@ class ModelAPI:
                         )
 
             except Exception as e:  # pylint: disable=broad-except
-                finish_req = FinishUploadModelSchema(
-                    status=ModelUploadStatus.FAILED.value,
+                finish_req = FinishUploadSchema(
+                    status=UploadStatus.FAILED.value,
                     reason=str(e),
                 )
-            if finish_req.status == ModelUploadStatus.FAILED.value:
+            if finish_req.status == UploadStatus.FAILED.value:
                 self.spinner.log(f'[bold red]Failed to upload model "{model.tag}"')
             with self.spinner.spin(
                 text="Submitting upload status to remote model store"
@@ -319,7 +317,7 @@ class ModelAPI:
                     req=finish_req,
                 )
 
-            if finish_req.status != ModelUploadStatus.SUCCESS.value:
+            if finish_req.status != UploadStatus.SUCCESS.value:
                 self.spinner.log(
                     f'[bold red]Failed pushing model "{model.tag}" : {finish_req.reason}'
                 )
diff --git a/src/bentoml/_internal/cloud/schemas/modelschemas.py b/src/bentoml/_internal/cloud/schemas/modelschemas.py
index 21fa3b2d..441ecd3b 100644
--- a/src/bentoml/_internal/cloud/schemas/modelschemas.py
+++ b/src/bentoml/_internal/cloud/schemas/modelschemas.py
@@ -41,7 +41,7 @@ class BentoImageBuildStatus(Enum):
     FAILED = "failed"
 
 
-class BentoUploadStatus(Enum):
+class UploadStatus(Enum):
     PENDING = "pending"
     BUILDING = "uploading"
     SUCCESS = "success"
@@ -97,13 +97,6 @@ else:
     TransmissionStrategy = str
 
 
-class ModelUploadStatus(Enum):
-    PENDING = "pending"
-    BUILDING = "uploading"
-    SUCCESS = "success"
-    FAILED = "failed"
-
-
 @attr.define
 class ModelManifestSchema:
     module: str
diff --git a/src/bentoml/_internal/cloud/schemas/schemasv1.py b/src/bentoml/_internal/cloud/schemas/schemasv1.py
index 2b481987..6294b84a 100644
--- a/src/bentoml/_internal/cloud/schemas/schemasv1.py
+++ b/src/bentoml/_internal/cloud/schemas/schemasv1.py
@@ -156,7 +156,7 @@ class CompleteMultipartUploadSchema:
 
 
 @attr.define
-class FinishUploadBentoSchema:
+class FinishUploadSchema:
     status: t.Optional[str]
     reason: t.Optional[str]
 
@@ -211,12 +211,6 @@ class CreateModelSchema:
     labels: t.List[LabelItemSchema] = attr.field(factory=list)
 
 
-@attr.define
-class FinishUploadModelSchema:
-    status: t.Optional[str]
-    reason: t.Optional[str]
-
-
 @attr.define
 class BentoRepositoryListSchema(BaseListSchema):
     items: t.List[BentoRepositorySchema]
diff --git a/src/bentoml/_internal/cloud/yatai.py b/src/bentoml/_internal/cloud/yatai.py
index 919b0e5d..695d0517 100644
--- a/src/bentoml/_internal/cloud/yatai.py
+++ b/src/bentoml/_internal/cloud/yatai.py
@@ -31,8 +31,7 @@ from .config import get_rest_api_client
 from .schemas.modelschemas import BentoApiSchema
 from .schemas.modelschemas import BentoRunnerResourceSchema
 from .schemas.modelschemas import BentoRunnerSchema
-from .schemas.modelschemas import BentoUploadStatus
-from .schemas.modelschemas import ModelUploadStatus
+from .schemas.modelschemas import UploadStatus
 from .schemas.schemasv1 import BentoManifestSchema
 from .schemas.schemasv1 import CompleteMultipartUploadSchema
 from .schemas.schemasv1 import CompletePartSchema
@@ -40,8 +39,7 @@ from .schemas.schemasv1 import CreateBentoRepositorySchema
 from .schemas.schemasv1 import CreateBentoSchema
 from .schemas.schemasv1 import CreateModelRepositorySchema
 from .schemas.schemasv1 import CreateModelSchema
-from .schemas.schemasv1 import FinishUploadBentoSchema
-from .schemas.schemasv1 import FinishUploadModelSchema
+from .schemas.schemasv1 import FinishUploadSchema
 from .schemas.schemasv1 import LabelItemSchema
 from .schemas.schemasv1 import ModelManifestSchema
 from .schemas.schemasv1 import PreSignMultipartUploadUrlSchema
@@ -129,7 +127,7 @@ class YataiClient:
         if (
             not force
             and remote_bento
-            and remote_bento.upload_status == BentoUploadStatus.SUCCESS
+            and remote_bento.upload_status == UploadStatus.SUCCESS
         ):
             self.spinner.log(
                 f'[bold blue]Push failed: Bento "{bento.tag}" already exists in Yatai'
@@ -258,16 +256,16 @@ class YataiClient:
                     raise e
                 self.spinner.log(f'[bold green]Successfully pushed bento "{bento.tag}"')
                 return
-            finish_req = FinishUploadBentoSchema(
-                status=BentoUploadStatus.SUCCESS,
+            finish_req = FinishUploadSchema(
+                status=UploadStatus.SUCCESS,
                 reason="",
             )
             try:
                 if presigned_upload_url is not None:
                     resp = httpx.put(presigned_upload_url, data=tar_io)
                     if resp.status_code != 200:
-                        finish_req = FinishUploadBentoSchema(
-                            status=BentoUploadStatus.FAILED,
+                        finish_req = FinishUploadSchema(
+                            status=UploadStatus.FAILED,
                             reason=resp.text,
                         )
                 else:
@@ -289,7 +287,7 @@ class YataiClient:
 
                     def chunk_upload(
                         upload_id: str, chunk_number: int
-                    ) -> FinishUploadBentoSchema | tuple[str, int]:
+                    ) -> FinishUploadSchema | tuple[str, int]:
                         with self.spinner.spin(
                             text=f'({chunk_number}/{chunks_count}) Presign multipart upload url of Bento "{bento.tag}"...'
                         ):
@@ -322,15 +320,13 @@ class YataiClient:
                                     remote_bento.presigned_upload_url, data=chunk_io
                                 )
                                 if resp.status_code != 200:
-                                    return FinishUploadBentoSchema(
-                                        status=BentoUploadStatus.FAILED,
+                                    return FinishUploadSchema(
+                                        status=UploadStatus.FAILED,
                                         reason=resp.text,
                                     )
                                 return resp.headers["ETag"], chunk_number
 
-                    futures_: list[
-                        Future[FinishUploadBentoSchema | tuple[str, int]]
-                    ] = []
+                    futures_: list[Future[FinishUploadSchema | tuple[str, int]]] = []
 
                     with ThreadPoolExecutor(
                         max_workers=min(max(chunks_count, 1), threads)
@@ -347,7 +343,7 @@ class YataiClient:
 
                     for future in futures_:
                         result = future.result()
-                        if isinstance(result, FinishUploadBentoSchema):
+                        if isinstance(result, FinishUploadSchema):
                             finish_req = result
                             break
                         else:
@@ -374,11 +370,11 @@ class YataiClient:
                         )
 
             except Exception as e:  # pylint: disable=broad-except
-                finish_req = FinishUploadBentoSchema(
-                    status=BentoUploadStatus.FAILED,
+                finish_req = FinishUploadSchema(
+                    status=UploadStatus.FAILED,
                     reason=str(e),
                 )
-            if finish_req.status is BentoUploadStatus.FAILED:
+            if finish_req.status is UploadStatus.FAILED:
                 self.spinner.log(f'[bold red]Failed to upload Bento "{bento.tag}"')
             with self.spinner.spin(text="Submitting upload status to Yatai"):
                 yatai_rest_client.finish_upload_bento(
@@ -386,7 +382,7 @@ class YataiClient:
                     version=version,
                     req=finish_req,
                 )
-            if finish_req.status != BentoUploadStatus.SUCCESS:
+            if finish_req.status != UploadStatus.SUCCESS:
                 self.spinner.log(
                     f'[bold red]Failed pushing Bento "{bento.tag}": {finish_req.reason}'
                 )
@@ -606,7 +602,7 @@ class YataiClient:
         if (
             not force
             and remote_model
-            and remote_model.upload_status == ModelUploadStatus.SUCCESS
+            and remote_model.upload_status == UploadStatus.SUCCESS
         ):
             self.spinner.log(
                 f'[bold blue]Model "{model.tag}" already exists in Yatai, skipping'
@@ -693,16 +689,16 @@ class YataiClient:
                     raise e
                 self.spinner.log(f'[bold green]Successfully pushed model "{model.tag}"')
                 return
-            finish_req = FinishUploadModelSchema(
-                status=ModelUploadStatus.SUCCESS,
+            finish_req = FinishUploadSchema(
+                status=UploadStatus.SUCCESS,
                 reason="",
             )
             try:
                 if presigned_upload_url is not None:
                     resp = httpx.put(presigned_upload_url, data=tar_io)
                     if resp.status_code != 200:
-                        finish_req = FinishUploadModelSchema(
-                            status=ModelUploadStatus.FAILED,
+                        finish_req = FinishUploadSchema(
+                            status=UploadStatus.FAILED,
                             reason=resp.text,
                         )
                 else:
@@ -724,7 +720,7 @@ class YataiClient:
 
                     def chunk_upload(
                         upload_id: str, chunk_number: int
-                    ) -> FinishUploadModelSchema | tuple[str, int]:
+                    ) -> FinishUploadSchema | tuple[str, int]:
                         with self.spinner.spin(
                             text=f'({chunk_number}/{chunks_count}) Presign multipart upload url of model "{model.tag}"...'
                         ):
@@ -758,15 +754,13 @@ class YataiClient:
                                     remote_model.presigned_upload_url, content=chunk_io
                                 )
                                 if resp.status_code != 200:
-                                    return FinishUploadModelSchema(
-                                        status=ModelUploadStatus.FAILED,
+                                    return FinishUploadSchema(
+                                        status=UploadStatus.FAILED,
                                         reason=resp.text,
                                     )
                                 return resp.headers["ETag"], chunk_number
 
-                    futures_: list[
-                        Future[FinishUploadModelSchema | tuple[str, int]]
-                    ] = []
+                    futures_: list[Future[FinishUploadSchema | tuple[str, int]]] = []
 
                     with ThreadPoolExecutor(
                         max_workers=min(max(chunks_count, 1), threads)
@@ -783,7 +777,7 @@ class YataiClient:
 
                     for future in futures_:
                         result = future.result()
-                        if isinstance(result, FinishUploadModelSchema):
+                        if isinstance(result, FinishUploadSchema):
                             finish_req = result
                             break
                         else:
@@ -810,11 +804,11 @@ class YataiClient:
                         )
 
             except Exception as e:  # pylint: disable=broad-except
-                finish_req = FinishUploadModelSchema(
-                    status=ModelUploadStatus.FAILED,
+                finish_req = FinishUploadSchema(
+                    status=UploadStatus.FAILED,
                     reason=str(e),
                 )
-            if finish_req.status is ModelUploadStatus.FAILED:
+            if finish_req.status is UploadStatus.FAILED:
                 self.spinner.log(f'[bold red]Failed to upload model "{model.tag}"')
             with self.spinner.spin(text="Submitting upload status to Yatai"):
                 yatai_rest_client.finish_upload_model(
@@ -822,7 +816,7 @@ class YataiClient:
                     version=version,
                     req=finish_req,
                 )
-            if finish_req.status != ModelUploadStatus.SUCCESS:
+            if finish_req.status != UploadStatus.SUCCESS:
                 self.spinner.log(
                     f'[bold red]Failed pushing model "{model.tag}" : {finish_req.reason}'
                 )


Test patch:
diff --git a/tests/unit/_internal/cloud/test_deployment.py b/tests/unit/_internal/cloud/test_deployment.py
index 4af2a1aa..1ef1d183 100644
--- a/tests/unit/_internal/cloud/test_deployment.py
+++ b/tests/unit/_internal/cloud/test_deployment.py
@@ -10,13 +10,13 @@ from bentoml._internal.cloud.client import RestApiClient
 from bentoml._internal.cloud.deployment import DeploymentAPI
 from bentoml._internal.cloud.deployment import DeploymentConfigParameters
 from bentoml._internal.cloud.schemas.modelschemas import BentoImageBuildStatus
-from bentoml._internal.cloud.schemas.modelschemas import BentoUploadStatus
 from bentoml._internal.cloud.schemas.modelschemas import DeploymentRevisionStatus
 from bentoml._internal.cloud.schemas.modelschemas import DeploymentServiceConfig
 from bentoml._internal.cloud.schemas.modelschemas import DeploymentStatus
 from bentoml._internal.cloud.schemas.modelschemas import DeploymentTargetHPAConf
 from bentoml._internal.cloud.schemas.modelschemas import EnvItemSchema
 from bentoml._internal.cloud.schemas.modelschemas import ResourceType
+from bentoml._internal.cloud.schemas.modelschemas import UploadStatus
 from bentoml._internal.cloud.schemas.schemasv1 import BentoFullSchema
 from bentoml._internal.cloud.schemas.schemasv1 import BentoManifestSchema
 from bentoml._internal.cloud.schemas.schemasv1 import BentoRepositorySchema
@@ -76,7 +76,7 @@ def dummy_generate_deployment_schema(
         description="",
         version=bento_version,
         image_build_status=BentoImageBuildStatus.PENDING,
-        upload_status=BentoUploadStatus.SUCCESS,
+        upload_status=UploadStatus.SUCCESS,
         upload_finished_reason="",
         presigned_upload_url="",
         presigned_download_url="",