address None and comments from review

Author: newokaerinasai

src/together/cli/api/finetune.py

@@ -42,8 +42,8 @@
 
 _WARNING_MESSAGE_INSUFFICIENT_FUNDS = (
     "The estimated price of this job is significantly greater than your current credit limit and balance combined. "
-    "It will likely fail due to insufficient funds. "
-    "Please consider increasing your credit limit at https://api.together.xyz/settings/profile\n"
+    "It will likely get cancelled due to insufficient funds. "
+    "Consider increasing your credit limit at https://api.together.xyz/settings/profile\n"
 )
 
 

src/together/resources/finetune.py

@@ -47,8 +47,8 @@
 _WARNING_MESSAGE_INSUFFICIENT_FUNDS = (
     "The estimated price of the fine-tuning job is {} which is significantly "
     "greater than your current credit limit and balance combined. "
-    "It will likely fail due to insufficient funds. "
-    "Please proceed at your own risk."
+    "It will likely get cancelled due to insufficient funds. "
+    "Proceed at your own risk."
 )
 
 
@@ -481,16 +481,25 @@ def create(
             hf_api_token=hf_api_token,
             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 from_checkpoint is None:
+            price_estimation_result = self.estimate_price(
+                training_file=training_file,
+                validation_file=validation_file,
+                model=model_name,
+                n_epochs=finetune_request.n_epochs,
+                n_evals=finetune_request.n_evals,
+                training_type="lora" if lora else "full",
+                training_method=training_method,
+            )
+        else:
+            # unsupported case
+            price_estimation_result = FinetunePriceEstimationResponse(
+                estimated_total_price=0.0,
+                allowed_to_proceed=True,
+                estimated_train_token_count=0,
+                estimated_eval_token_count=0,
+                user_limit=0.0,
+            )
 
         if verbose:
             rprint(
@@ -523,10 +532,10 @@ def estimate_price(
         self,
         *,
         training_file: str,
-        model: str | None,
+        model: str,
         validation_file: str | None = None,
-        n_epochs: int | None = None,
-        n_evals: int | None = None,
+        n_epochs: int | None = 1,
+        n_evals: int | None = 0,
         training_type: str = "lora",
         training_method: str = "sft",
     ) -> FinetunePriceEstimationResponse:
@@ -539,8 +548,8 @@ def estimate_price(
         Returns:
             FinetunePriceEstimationResponse: Object containing the estimated price.
         """
-        training_type_cls: TrainingType | None = None
-        training_method_cls: TrainingMethod | None = None
+        training_type_cls: TrainingType
+        training_method_cls: TrainingMethod
 
         if training_method == "sft":
             training_method_cls = TrainingMethodSFT(method="sft")
@@ -1036,15 +1045,25 @@ 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 from_checkpoint is not None:
+            price_estimation_result = await self.estimate_price(
+                training_file=training_file,
+                validation_file=validation_file,
+                model=model_name,
+                n_epochs=finetune_request.n_epochs,
+                n_evals=finetune_request.n_evals,
+                training_type="lora" if lora else "full",
+                training_method=training_method,
+            )
+        else:
+            # unsupported case
+            price_estimation_result = FinetunePriceEstimationResponse(
+                estimated_total_price=0.0,
+                allowed_to_proceed=True,
+                estimated_train_token_count=0,
+                estimated_eval_token_count=0,
+                user_limit=0.0,
+            )
 
         if verbose:
             rprint(
@@ -1080,28 +1099,53 @@ async def estimate_price(
         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,
+        n_epochs: int | None = 1,
+        n_evals: int | None = 0,
+        training_type: str = "lora",
+        training_method: str = "sft",
     ) -> FinetunePriceEstimationResponse:
         """
-        Async method to estimate the price of a fine-tuning job
+        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
+        training_method_cls: TrainingMethod
+
+        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":
+            # parameters of lora are unused in price estimation
+            # but we need to set them to valid values
+            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,
-            training_method=training_method,
+            training_type=training_type_cls,
+            training_method=training_method_cls,
         )
         parameter_payload = request.model_dump(exclude_none=True)
         requestor = api_requestor.APIRequestor(

src/together/types/finetune.py

@@ -316,10 +316,10 @@ class FinetunePriceEstimationRequest(BaseModel):
     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
+    n_epochs: int
+    n_evals: int
+    training_type: TrainingType
+    training_method: TrainingMethod
 
 
 class FinetunePriceEstimationResponse(BaseModel):

tests/unit/test_finetune_resources.py

@@ -1,6 +1,10 @@
 import pytest
+from unittest.mock import MagicMock, Mock, patch
 
+from together.client import Together
 from together.resources.finetune import create_finetune_request
+from together.together_response import TogetherResponse
+from together.types import TogetherRequest
 from together.types.finetune import (
     FinetuneFullTrainingLimits,
     FinetuneLoraTrainingLimits,
@@ -31,6 +35,41 @@
 )
 
 
+def mock_request(options: TogetherRequest, *args, **kwargs):
+    if options.url == "fine-tunes/estimate-price":
+        return (
+            TogetherResponse(
+                data={
+                    "estimated_total_price": 100,
+                    "allowed_to_proceed": True,
+                    "estimated_train_token_count": 1000,
+                    "estimated_eval_token_count": 100,
+                    "user_limit": 1000,
+                },
+                headers={},
+            ),
+            None,
+            None,
+        )
+    elif options.url == "fine-tunes":
+        return (
+            TogetherResponse(
+                data={
+                    "id": "ft-12345678-1234-1234-1234-1234567890ab",
+                },
+                headers={},
+            ),
+            None,
+            None,
+        )
+    else:
+        return (
+            TogetherResponse(data=_MODEL_LIMITS.model_dump(), headers={}),
+            None,
+            None,
+        )
+
+
 def test_simple_request():
     request = create_finetune_request(
         model_limits=_MODEL_LIMITS,
@@ -335,3 +374,87 @@ def test_train_on_inputs_not_supported_for_dpo():
             training_method="dpo",
             train_on_inputs=True,
         )
+
+
+@patch("together.abstract.api_requestor.APIRequestor.request")
+def test_price_estimation_request(mocker):
+    test_data = [
+        {
+            "training_type": "lora",
+            "training_method": "sft",
+        },
+        {
+            "training_type": "lora",
+            "training_method": "dpo",
+        },
+        {
+            "training_type": "full",
+            "training_method": "sft",
+        },
+    ]
+    mocker.return_value = (
+        TogetherResponse(
+            data={
+                "estimated_total_price": 100,
+                "allowed_to_proceed": True,
+                "estimated_train_token_count": 1000,
+                "estimated_eval_token_count": 100,
+                "user_limit": 1000,
+            },
+            headers={},
+        ),
+        None,
+        None,
+    )
+    client = Together()
+    for test_case in test_data:
+        response = client.fine_tuning.estimate_price(
+            training_file=_TRAINING_FILE,
+            model=_MODEL_NAME,
+            validation_file=_VALIDATION_FILE,
+            n_epochs=1,
+            n_evals=0,
+            training_type=test_case["training_type"],
+            training_method=test_case["training_method"],
+        )
+        assert response.estimated_total_price > 0
+        assert response.allowed_to_proceed
+        assert response.estimated_train_token_count > 0
+        assert response.estimated_eval_token_count > 0
+
+
+def test_create_ft_job(mocker):
+    mock_requestor = Mock()
+    mock_requestor.request = MagicMock()
+    mock_requestor.request.side_effect = mock_request
+    mocker.patch(
+        "together.abstract.api_requestor.APIRequestor", return_value=mock_requestor
+    )
+
+    client = Together()
+    response = client.fine_tuning.create(
+        training_file=_TRAINING_FILE,
+        model=_MODEL_NAME,
+        validation_file=_VALIDATION_FILE,
+        n_epochs=1,
+        n_evals=0,
+        lora=True,
+        training_method="sft",
+    )
+
+    assert mock_requestor.request.call_count == 3
+    assert response.id == "ft-12345678-1234-1234-1234-1234567890ab"
+
+    response = client.fine_tuning.create(
+        training_file=_TRAINING_FILE,
+        model=None,
+        validation_file=_VALIDATION_FILE,
+        n_epochs=1,
+        n_evals=0,
+        lora=True,
+        training_method="sft",
+        from_checkpoint=_FROM_CHECKPOINT,
+    )
+
+    assert mock_requestor.request.call_count == 5
+    assert response.id == "ft-12345678-1234-1234-1234-1234567890ab"