add price estimation

Author: newokaerinasai

src/together/cli/api/finetune.py

@@ -17,6 +17,8 @@
     DownloadCheckpointType,
     FinetuneEventType,
     FinetuneTrainingLimits,
+    FullTrainingType,
+    LoRATrainingType,
 )
 from together.utils import (
     finetune_price_to_dollars,
@@ -36,6 +38,15 @@
     "Do you want to proceed?"
 )
 
+_PRICE_ESTIMATION_CONFIRMATION_MESSAGE = (
+    "The estimated price of the fine-tuning job is {} which is significantly "
+    "greater than your current credit limit and balance. "
+    "It will likely fail due to insufficient funds. "
+    "Please consider increasing your credit limit at https://api.together.xyz/settings/profile\n"
+    "You can pass `-y` or `--confirm` to your command to skip this message.\n\n"
+    "Do you want to proceed?"
+)
+
 
 class DownloadCheckpointTypeChoice(click.Choice):
     def __init__(self) -> None:
@@ -358,20 +369,49 @@ def create(
         )
 
     if confirm or click.confirm(_CONFIRMATION_MESSAGE, default=True, show_default=True):
-        response = client.fine_tuning.create(
-            **training_args,
-            verbose=True,
+        price_estimation_response = client.fine_tuning.estimate_price(
+            training_file=training_file,
+            validation_file=validation_file,
+            model=model,
+            n_epochs=n_epochs,
+            n_evals=n_evals,
+            training_type="lora" if lora else "full",
+            training_method=training_method,
         )
-
-        report_string = f"Successfully submitted a fine-tuning job {response.id}"
-        if response.created_at is not None:
-            created_time = datetime.strptime(
-                response.created_at, "%Y-%m-%dT%H:%M:%S.%f%z"
+        proceed = (
+            confirm
+            or price_estimation_response.allowed_to_proceed
+            or (
+                not price_estimation_response.allowed_to_proceed
+                and click.confirm(
+                    click.style(
+                        _PRICE_ESTIMATION_CONFIRMATION_MESSAGE.format(
+                            price_estimation_response.estimated_total_price
+                        ),
+                        fg="red",
+                        bold=True,
+                    ),
+                    default=True,
+                    show_default=True,
+                )
+            )
+        )
+        if proceed:
+            response = client.fine_tuning.create(
+                **training_args,
+                verbose=True,
             )
-            # created_at reports UTC time, we use .astimezone() to convert to local time
-            formatted_time = created_time.astimezone().strftime("%m/%d/%Y, %H:%M:%S")
-            report_string += f" at {formatted_time}"
-        rprint(report_string)
+            report_string = f"Successfully submitted a fine-tuning job {response.id}"
+            if response.created_at is not None:
+                created_time = datetime.strptime(
+                    response.created_at, "%Y-%m-%dT%H:%M:%S.%f%z"
+                )
+                # created_at reports UTC time, we use .astimezone() to convert to local time
+                formatted_time = created_time.astimezone().strftime(
+                    "%m/%d/%Y, %H:%M:%S"
+                )
+                report_string += f" at {formatted_time}"
+            rprint(report_string)
     else:
         click.echo("No confirmation received, stopping job launch")
 

src/together/resources/finetune.py

@@ -20,6 +20,8 @@
     FinetuneLRScheduler,
     FinetuneRequest,
     FinetuneResponse,
+    FinetunePriceEstimationRequest,
+    FinetunePriceEstimationResponse,
     FinetuneTrainingLimits,
     FullTrainingType,
     LinearLRScheduler,
@@ -31,7 +33,7 @@
     TrainingMethodSFT,
     TrainingType,
 )
-from together.types.finetune import DownloadCheckpointType
+from together.types.finetune import DownloadCheckpointType, TrainingMethod
 from together.utils import log_warn_once, normalize_key
 
 
@@ -42,6 +44,12 @@
     TrainingMethodSFT().method,
     TrainingMethodDPO().method,
 }
+_CONFIRMATION_MESSAGE_INSUFFICIENT_FUNDS = (
+    "The estimated price of the fine-tuning job is {} which is significantly "
+    "greater than your current credit limit and balance. "
+    "It will likely fail due to insufficient funds. "
+    "Please proceed at your own risk."
+)
 
 
 def create_finetune_request(
@@ -474,11 +482,29 @@ def create(
             hf_output_repo_name=hf_output_repo_name,
         )
 
+        price_estimation_result = self.estimate_price(
+            training_file=training_file,
+            validation_file=validation_file,
+            model=model_name,
+            n_epochs=n_epochs,
+            n_evals=n_evals,
+            training_type="lora" if lora else "full",
+            training_method=training_method,
+        )
+
         if verbose:
             rprint(
                 "Submitting a fine-tuning job with the following parameters:",
                 finetune_request,
             )
+            if not price_estimation_result.allowed_to_proceed:
+                rprint(
+                    "[red]"
+                    + _CONFIRMATION_MESSAGE_INSUFFICIENT_FUNDS.format(
+                        price_estimation_result.estimated_total_price
+                    )
+                    + "[/red]",
+                )
         parameter_payload = finetune_request.model_dump(exclude_none=True)
 
         response, _, _ = requestor.request(
@@ -493,6 +519,73 @@ def create(
 
         return FinetuneResponse(**response.data)
 
+    def estimate_price(
+        self,
+        *,
+        training_file: str,
+        model: str | None,
+        validation_file: str | None = None,
+        n_epochs: int | None = None,
+        n_evals: int | None = None,
+        training_type: str = "lora",
+        training_method: str = "sft",
+    ) -> FinetunePriceEstimationResponse:
+        """
+        Estimates the price of a fine-tuning job
+
+        Args:
+            request (FinetunePriceEstimationRequest): Request object containing the parameters for the price estimation.
+
+        Returns:
+            FinetunePriceEstimationResponse: Object containing the estimated price.
+        """
+        training_type_cls: TrainingType | None = None
+        training_method_cls: TrainingMethod | None = None
+
+        if training_method == "sft":
+            training_method_cls = TrainingMethodSFT(method="sft")
+        elif training_method == "dpo":
+            training_method_cls = TrainingMethodDPO(method="dpo")
+        else:
+            raise ValueError(f"Unknown training method: {training_method}")
+
+        if training_type.lower() == "lora":
+            training_type_cls = LoRATrainingType(
+                type="Lora",
+                lora_r=16,
+                lora_alpha=16,
+                lora_dropout=0.0,
+                lora_trainable_modules="all-linear",
+            )
+        elif training_type.lower() == "full":
+            training_type_cls = FullTrainingType(type="Full")
+        else:
+            raise ValueError(f"Unknown training type: {training_type}")
+
+        request = FinetunePriceEstimationRequest(
+            training_file=training_file,
+            validation_file=validation_file,
+            model=model,
+            n_epochs=n_epochs,
+            n_evals=n_evals,
+            training_type=training_type_cls,
+            training_method=training_method_cls,
+        )
+        parameter_payload = request.model_dump(exclude_none=True)
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=TogetherRequest(
+                method="POST", url="fine-tunes/estimate-price", params=parameter_payload
+            ),
+            stream=False,
+        )
+        assert isinstance(response, TogetherResponse)
+
+        return FinetunePriceEstimationResponse(**response.data)
+
     def list(self) -> FinetuneList:
         """
         Lists fine-tune job history
@@ -941,11 +1034,29 @@ async def create(
             hf_output_repo_name=hf_output_repo_name,
         )
 
+        price_estimation_result = await self.estimate_price(
+            training_file=training_file,
+            validation_file=validation_file,
+            model=model_name,
+            n_epochs=n_epochs,
+            n_evals=n_evals,
+            training_type=finetune_request.training_type,
+            training_method=finetune_request.training_method,
+        )
+
         if verbose:
             rprint(
                 "Submitting a fine-tuning job with the following parameters:",
                 finetune_request,
             )
+            if not price_estimation_result.allowed_to_proceed:
+                rprint(
+                    "[red]"
+                    + _CONFIRMATION_MESSAGE_INSUFFICIENT_FUNDS.format(
+                        price_estimation_result.estimated_total_price
+                    )
+                    + "[/red]",
+                )
         parameter_payload = finetune_request.model_dump(exclude_none=True)
 
         response, _, _ = await requestor.arequest(
@@ -961,6 +1072,50 @@ async def create(
 
         return FinetuneResponse(**response.data)
 
+    async def estimate_price(
+        self,
+        *,
+        training_file: str,
+        model: str,
+        validation_file: str | None = None,
+        n_epochs: int | None = None,
+        n_evals: int | None = None,
+        training_type: TrainingType | None = None,
+        training_method: TrainingMethodSFT | TrainingMethodDPO | None = None,
+    ) -> FinetunePriceEstimationResponse:
+        """
+        Async method to estimate the price of a fine-tuning job
+
+        Args:
+            request (FinetunePriceEstimationRequest): Request object containing the parameters for the price estimation.
+
+        Returns:
+            FinetunePriceEstimationResponse: Object containing the estimated price.
+        """
+        request = FinetunePriceEstimationRequest(
+            training_file=training_file,
+            validation_file=validation_file,
+            model=model,
+            n_epochs=n_epochs,
+            n_evals=n_evals,
+            training_type=training_type,
+            training_method=training_method,
+        )
+        parameter_payload = request.model_dump(exclude_none=True)
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=TogetherRequest(
+                method="POST", url="fine-tunes/estimate-price", params=parameter_payload
+            ),
+            stream=False,
+        )
+        assert isinstance(response, TogetherResponse)
+
+        return FinetunePriceEstimationResponse(**response.data)
+
     async def list(self) -> FinetuneList:
         """
         Async method to list fine-tune job history

src/together/types/__init__.py

@@ -54,6 +54,8 @@
     FinetuneListEvents,
     FinetuneRequest,
     FinetuneResponse,
+    FinetunePriceEstimationRequest,
+    FinetunePriceEstimationResponse,
     FinetuneDeleteResponse,
     FinetuneTrainingLimits,
     FullTrainingType,
@@ -103,6 +105,8 @@
     "FinetuneDeleteResponse",
     "FinetuneDownloadResult",
     "FinetuneLRScheduler",
+    "FinetunePriceEstimationRequest",
+    "FinetunePriceEstimationResponse",
     "LinearLRScheduler",
     "LinearLRSchedulerArgs",
     "CosineLRScheduler",

src/together/types/finetune.py

@@ -308,6 +308,32 @@ def validate_training_type(cls, v: TrainingType) -> TrainingType:
             raise ValueError("Unknown training type")
 
 
+class FinetunePriceEstimationRequest(BaseModel):
+    """
+    Fine-tune price estimation request type
+    """
+
+    training_file: str
+    validation_file: str | None = None
+    model: str
+    n_epochs: int | None = None
+    n_evals: int | None = None
+    training_type: TrainingType | None = None
+    training_method: TrainingMethodSFT | TrainingMethodDPO
+
+
+class FinetunePriceEstimationResponse(BaseModel):
+    """
+    Fine-tune price estimation response type
+    """
+
+    estimated_total_price: float
+    user_limit: float
+    estimated_train_token_count: int
+    estimated_eval_token_count: int
+    allowed_to_proceed: bool
+
+
 class FinetuneList(BaseModel):
     # object type
     object: Literal["list"] | None = None