kiln-ai 0.13.2__py3-none-any.whl → 0.15.0__py3-none-any.whl

kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py

@@ -315,7 +315,7 @@ async def test_generate_and_upload_jsonl_success(
             "thinking_instructions": thinking_instructions,
         }
 
-        assert result == mock_dataset_id
+        assert result == "kiln-" + mock_dataset_id
         assert mock_client.post.call_count == 2
         assert mock_client.get.call_count == 1
 
@@ -448,7 +448,7 @@ def test_available_parameters(fireworks_finetune):
     assert payload_parameters == {"loraRank": 16, "epochs": 3}
 
 
-async def test_deploy_success(fireworks_finetune, mock_api_key):
+async def test_deploy_serverless_success(fireworks_finetune, mock_api_key):
     # Mock response for successful deployment
     success_response = MagicMock(spec=httpx.Response)
     success_response.status_code = 200
@@ -467,12 +467,12 @@ async def test_deploy_success(fireworks_finetune, mock_api_key):
         mock_client.post.return_value = success_response
         mock_client_class.return_value.__aenter__.return_value = mock_client
 
-        result = await fireworks_finetune._deploy()
+        result = await fireworks_finetune._deploy_serverless()
         assert result is True
         assert fireworks_finetune.datamodel.fine_tune_model_id == "ftm-123"
 
 
-async def test_deploy_already_deployed(fireworks_finetune, mock_api_key):
+async def test_deploy_serverless_already_deployed(fireworks_finetune, mock_api_key):
     # Mock response for already deployed model
     already_deployed_response = MagicMock(spec=httpx.Response)
     already_deployed_response.status_code = 400
@@ -494,12 +494,12 @@ async def test_deploy_already_deployed(fireworks_finetune, mock_api_key):
         mock_client.post.return_value = already_deployed_response
         mock_client_class.return_value.__aenter__.return_value = mock_client
 
-        result = await fireworks_finetune._deploy()
+        result = await fireworks_finetune._deploy_serverless()
         assert result is True
         assert fireworks_finetune.datamodel.fine_tune_model_id == "ftm-123"
 
 
-async def test_deploy_failure(fireworks_finetune, mock_api_key):
+async def test_deploy_serverless_failure(fireworks_finetune, mock_api_key):
     # Mock response for failed deployment
     failure_response = MagicMock(spec=httpx.Response)
     failure_response.status_code = 500
@@ -510,18 +510,28 @@ async def test_deploy_failure(fireworks_finetune, mock_api_key):
         mock_client.post.return_value = failure_response
         mock_client_class.return_value.__aenter__.return_value = mock_client
 
-        result = await fireworks_finetune._deploy()
+        result = await fireworks_finetune._deploy_serverless()
         assert result is False
 
 
-async def test_deploy_missing_credentials(fireworks_finetune):
+async def test_deploy_serverless_missing_credentials(fireworks_finetune):
     # Test missing API key or account ID
     with patch.object(Config, "shared") as mock_config:
         mock_config.return_value.fireworks_api_key = None
         mock_config.return_value.fireworks_account_id = None
 
         with pytest.raises(ValueError, match="Fireworks API key or account ID not set"):
-            await fireworks_finetune._deploy()
+            await fireworks_finetune._deploy_serverless()
+
+
+async def test_deploy_server_missing_credentials(fireworks_finetune):
+    # Test missing API key or account ID
+    with patch.object(Config, "shared") as mock_config:
+        mock_config.return_value.fireworks_api_key = None
+        mock_config.return_value.fireworks_account_id = None
+
+        response = await fireworks_finetune._check_or_deploy_server()
+        assert response is False
 
 
 async def test_deploy_missing_model_id(fireworks_finetune, mock_api_key):
@@ -564,3 +574,479 @@ async def test_status_with_deploy(fireworks_finetune, mock_api_key):
         # Verify message was updated due to failed deployment
         assert status.status == FineTuneStatusType.completed
         assert status.message == "Fine-tuning job completed but failed to deploy model."
+
+
+@pytest.mark.paid
+async def test_fetch_all_deployments(fireworks_finetune):
+    deployments = await fireworks_finetune._fetch_all_deployments()
+    assert isinstance(deployments, list)
+
+
+async def test_api_key_and_account_id(fireworks_finetune, mock_api_key):
+    # Test successful retrieval of API key and account ID
+    api_key, account_id = fireworks_finetune.api_key_and_account_id()
+    assert api_key == "test-api-key"
+    assert account_id == "test-account-id"
+
+
+async def test_api_key_and_account_id_missing_credentials(fireworks_finetune):
+    # Test missing API key or account ID
+    with patch.object(Config, "shared") as mock_config:
+        mock_config.return_value.fireworks_api_key = None
+        mock_config.return_value.fireworks_account_id = None
+
+        with pytest.raises(ValueError, match="Fireworks API key or account ID not set"):
+            fireworks_finetune.api_key_and_account_id()
+
+
+def test_deployment_display_name(fireworks_finetune):
+    # Test with default ID and name
+    display_name = fireworks_finetune.deployment_display_name()
+    expected = f"Kiln AI fine-tuned model [ID:{fireworks_finetune.datamodel.id}][name:test-finetune]"[
+        :60
+    ]
+    assert display_name == expected
+
+    # Test with a very long name to ensure 60 character limit
+    fireworks_finetune.datamodel.name = "x" * 100
+    display_name = fireworks_finetune.deployment_display_name()
+    assert len(display_name) == 60
+    assert display_name.startswith("Kiln AI fine-tuned model [ID:")
+
+
+async def test_model_id_checking_status_completed(fireworks_finetune):
+    # Test with completed status and valid model ID
+    status_response = (
+        FineTuneStatus(status=FineTuneStatusType.completed, message=""),
+        "model-123",
+    )
+
+    with patch.object(fireworks_finetune, "_status", return_value=status_response):
+        model_id = await fireworks_finetune.model_id_checking_status()
+        assert model_id == "model-123"
+
+
+async def test_model_id_checking_status_not_completed(fireworks_finetune):
+    # Test with non-completed status
+    status_response = (
+        FineTuneStatus(status=FineTuneStatusType.running, message=""),
+        "model-123",
+    )
+
+    with patch.object(fireworks_finetune, "_status", return_value=status_response):
+        model_id = await fireworks_finetune.model_id_checking_status()
+        assert model_id is None
+
+
+async def test_model_id_checking_status_invalid_model_id(fireworks_finetune):
+    # Test with completed status but invalid model ID
+    status_response = (
+        FineTuneStatus(status=FineTuneStatusType.completed, message=""),
+        None,
+    )
+
+    with patch.object(fireworks_finetune, "_status", return_value=status_response):
+        model_id = await fireworks_finetune.model_id_checking_status()
+        assert model_id is None
+
+    # Test with non-string model ID
+    status_response = (
+        FineTuneStatus(status=FineTuneStatusType.completed, message=""),
+        {"id": "model-123"},  # Not a string
+    )
+
+    with patch.object(fireworks_finetune, "_status", return_value=status_response):
+        model_id = await fireworks_finetune.model_id_checking_status()
+        assert model_id is None
+
+
+@pytest.mark.parametrize(
+    "base_model_id,expected_method",
+    [
+        ("accounts/fireworks/models/llama-v3p1-8b-instruct", "_deploy_serverless"),
+        ("accounts/fireworks/models/llama-v3p1-70b-instruct", "_deploy_serverless"),
+        ("some-other-model", "_check_or_deploy_server"),
+    ],
+)
+async def test_deploy_model_selection(
+    fireworks_finetune, base_model_id, expected_method, mock_api_key
+):
+    # Set the base model ID
+    fireworks_finetune.datamodel.base_model_id = base_model_id
+
+    # Mock the deployment methods
+    with (
+        patch.object(
+            fireworks_finetune, "_deploy_serverless", return_value=True
+        ) as mock_serverless,
+        patch.object(
+            fireworks_finetune, "_check_or_deploy_server", return_value=True
+        ) as mock_server,
+    ):
+        result = await fireworks_finetune._deploy()
+
+        # Verify the correct method was called based on the model
+        if expected_method == "_deploy_serverless":
+            mock_serverless.assert_called_once()
+            mock_server.assert_not_called()
+        else:
+            mock_serverless.assert_not_called()
+            mock_server.assert_called_once()
+
+        assert result is True
+
+
+async def test_fetch_all_deployments_request_error(fireworks_finetune, mock_api_key):
+    # Test with error response
+    error_response = MagicMock(spec=httpx.Response)
+    error_response.status_code = 500
+    error_response.text = "Internal Server Error"
+
+    with patch("httpx.AsyncClient") as mock_client_class:
+        mock_client = AsyncMock()
+        mock_client.get.side_effect = Exception("API request failed")
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        with pytest.raises(Exception, match="API request failed"):
+            await fireworks_finetune._fetch_all_deployments()
+
+        # Verify API was called with correct parameters
+        mock_client.get.assert_called_once()
+        call_args = mock_client.get.call_args[1]
+        assert "params" in call_args
+        assert call_args["params"]["pageSize"] == 200
+
+
+async def test_fetch_all_deployments_standard_case(fireworks_finetune, mock_api_key):
+    # Test with single page of results
+    mock_deployments = [
+        {"id": "deploy-1", "baseModel": "model-1", "state": "READY"},
+        {"id": "deploy-2", "baseModel": "model-2", "state": "READY"},
+    ]
+
+    success_response = MagicMock(spec=httpx.Response)
+    success_response.status_code = 200
+    success_response.json.return_value = {
+        "deployments": mock_deployments,
+        "nextPageToken": None,
+    }
+
+    with patch("httpx.AsyncClient") as mock_client_class:
+        mock_client = AsyncMock()
+        mock_client.get.return_value = success_response
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        deployments = await fireworks_finetune._fetch_all_deployments()
+
+        # Verify API was called correctly
+        mock_client.get.assert_called_once()
+
+        # Verify correct deployments were returned
+        assert deployments == mock_deployments
+        assert len(deployments) == 2
+        assert deployments[0]["id"] == "deploy-1"
+        assert deployments[1]["id"] == "deploy-2"
+
+
+async def test_fetch_all_deployments_paged_case(fireworks_finetune, mock_api_key):
+    # Test with multiple pages of results
+    mock_deployments_page1 = [
+        {"id": "deploy-1", "baseModel": "model-1", "state": "READY"},
+        {"id": "deploy-2", "baseModel": "model-2", "state": "READY"},
+    ]
+
+    mock_deployments_page2 = [
+        {"id": "deploy-3", "baseModel": "model-3", "state": "READY"},
+        {"id": "deploy-4", "baseModel": "model-4", "state": "READY"},
+    ]
+
+    page1_response = MagicMock(spec=httpx.Response)
+    page1_response.status_code = 200
+    page1_response.json.return_value = {
+        "deployments": mock_deployments_page1,
+        "nextPageToken": "page2token",
+    }
+
+    page2_response = MagicMock(spec=httpx.Response)
+    page2_response.status_code = 200
+    page2_response.json.return_value = {
+        "deployments": mock_deployments_page2,
+        "nextPageToken": None,
+    }
+
+    with patch("httpx.AsyncClient") as mock_client_class:
+        mock_client = AsyncMock()
+        mock_client.get.side_effect = [page1_response, page2_response]
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        deployments = await fireworks_finetune._fetch_all_deployments()
+
+        # Verify API was called twice (once for each page)
+        assert mock_client.get.call_count == 2
+
+        # Verify first call had no page token
+        first_call_args = mock_client.get.call_args_list[0][1]
+        assert "pageToken" not in first_call_args["params"]
+
+        # Verify second call included the page token
+        second_call_args = mock_client.get.call_args_list[1][1]
+        assert second_call_args["params"]["pageToken"] == "page2token"
+
+        # Verify all deployments from both pages were returned
+        assert len(deployments) == 4
+        assert deployments == mock_deployments_page1 + mock_deployments_page2
+        for deployment in deployments:
+            assert deployment["id"] in [
+                "deploy-1",
+                "deploy-2",
+                "deploy-3",
+                "deploy-4",
+            ]
+
+
+async def test_deploy_server_success(fireworks_finetune, mock_api_key):
+    # Mock response for successful deployment
+    success_response = MagicMock(spec=httpx.Response)
+    success_response.status_code = 200
+    success_response.json.return_value = {"baseModel": "model-123"}
+
+    status_response = (
+        FineTuneStatus(status=FineTuneStatusType.completed, message=""),
+        "model-123",
+    )
+
+    with (
+        patch("httpx.AsyncClient") as mock_client_class,
+        patch.object(
+            fireworks_finetune, "model_id_checking_status", return_value="model-123"
+        ),
+    ):
+        mock_client = AsyncMock()
+        mock_client.post.return_value = success_response
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        result = await fireworks_finetune._deploy_server()
+
+        # Verify result
+        assert result is True
+
+        # Verify fine_tune_model_id was updated
+        assert fireworks_finetune.datamodel.fine_tune_model_id == "model-123"
+
+        # Verify API was called with correct parameters
+        mock_client.post.assert_called_once()
+        call_args = mock_client.post.call_args[1]
+        assert "json" in call_args
+        assert call_args["json"]["baseModel"] == "model-123"
+        assert call_args["json"]["minReplicaCount"] == 0
+        assert "autoscalingPolicy" in call_args["json"]
+        assert call_args["json"]["autoscalingPolicy"]["scaleToZeroWindow"] == "300s"
+
+        # load the datamodel from the file and confirm the fine_tune_model_id was updated
+        loaded_datamodel = FinetuneModel.load_from_file(
+            fireworks_finetune.datamodel.path
+        )
+        assert loaded_datamodel.fine_tune_model_id == "model-123"
+
+
+async def test_deploy_server_failure(fireworks_finetune, mock_api_key):
+    # Mock response for failed deployment
+    failure_response = MagicMock(spec=httpx.Response)
+    failure_response.status_code = 500
+    failure_response.text = "Internal Server Error"
+
+    with (
+        patch("httpx.AsyncClient") as mock_client_class,
+        patch.object(
+            fireworks_finetune, "model_id_checking_status", return_value="model-123"
+        ),
+    ):
+        mock_client = AsyncMock()
+        mock_client.post.return_value = failure_response
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        result = await fireworks_finetune._deploy_server()
+
+        # Verify result
+        assert result is False
+
+        # Verify API was called
+        mock_client.post.assert_called_once()
+
+
+async def test_deploy_server_non_200_but_valid_response(
+    fireworks_finetune, mock_api_key
+):
+    # Mock response with non-200 status but valid JSON response
+    mixed_response = MagicMock(spec=httpx.Response)
+    mixed_response.status_code = 200
+    mixed_response.json.return_value = {"not_baseModel": "something-else"}
+
+    with (
+        patch("httpx.AsyncClient") as mock_client_class,
+        patch.object(
+            fireworks_finetune, "model_id_checking_status", return_value="model-123"
+        ),
+    ):
+        mock_client = AsyncMock()
+        mock_client.post.return_value = mixed_response
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        result = await fireworks_finetune._deploy_server()
+
+        # Verify result - should fail because baseModel is missing
+        assert result is False
+
+
+async def test_deploy_server_missing_model_id(fireworks_finetune, mock_api_key):
+    # Test when model_id_checking_status returns None
+    with patch.object(
+        fireworks_finetune, "model_id_checking_status", return_value=None
+    ):
+        result = await fireworks_finetune._deploy_server()
+
+        # Verify result - should fail because model ID is missing
+        assert result is False
+
+
+@pytest.mark.parametrize(
+    "state,expected_already_deployed",
+    [
+        ("READY", True),
+        ("CREATING", True),
+        ("FAILED", False),
+    ],
+)
+async def test_check_or_deploy_server_already_deployed(
+    fireworks_finetune, mock_api_key, state, expected_already_deployed
+):
+    # Test when model is already deployed (should return True without calling _deploy_server)
+
+    # Set a fine_tune_model_id so we search for deployments
+    fireworks_finetune.datamodel.fine_tune_model_id = "model-123"
+
+    # Mock deployments including one matching our model ID
+    mock_deployments = [
+        {"id": "deploy-1", "baseModel": "different-model", "state": "READY"},
+        {"id": "deploy-2", "baseModel": "model-123", "state": state},
+    ]
+
+    with (
+        patch.object(
+            fireworks_finetune, "_fetch_all_deployments", return_value=mock_deployments
+        ) as mock_fetch,
+        patch.object(fireworks_finetune, "_deploy_server") as mock_deploy,
+    ):
+        mock_deploy.return_value = True
+        result = await fireworks_finetune._check_or_deploy_server()
+        # Even true if the model is in a non-ready state, as we'll call deploy (checked below)
+        assert result is True
+
+        if expected_already_deployed:
+            assert mock_deploy.call_count == 0
+        else:
+            assert mock_deploy.call_count == 1
+
+        # Verify _fetch_all_deployments was called
+        mock_fetch.assert_called_once()
+
+
+async def test_check_or_deploy_server_not_deployed(fireworks_finetune, mock_api_key):
+    # Test when model exists but isn't deployed (should call _deploy_server)
+
+    # Set a fine_tune_model_id so we search for deployments
+    fireworks_finetune.datamodel.fine_tune_model_id = "model-123"
+
+    # Mock deployments without our model ID
+    mock_deployments = [
+        {"id": "deploy-1", "baseModel": "different-model-1", "state": "READY"},
+        {"id": "deploy-2", "baseModel": "different-model-2", "state": "READY"},
+    ]
+
+    with (
+        patch.object(
+            fireworks_finetune, "_fetch_all_deployments", return_value=mock_deployments
+        ) as mock_fetch,
+        patch.object(
+            fireworks_finetune, "_deploy_server", return_value=True
+        ) as mock_deploy,
+    ):
+        result = await fireworks_finetune._check_or_deploy_server()
+
+        # Verify method returned True (from _deploy_server)
+        assert result is True
+
+        # Verify _fetch_all_deployments was called
+        mock_fetch.assert_called_once()
+
+        # Verify _deploy_server was called since model is not deployed
+        mock_deploy.assert_called_once()
+
+
+async def test_check_or_deploy_server_no_model_id(fireworks_finetune, mock_api_key):
+    # Test when no fine_tune_model_id exists (should skip fetch and call _deploy_server directly)
+
+    # Ensure no fine_tune_model_id is set
+    fireworks_finetune.datamodel.fine_tune_model_id = None
+
+    with (
+        patch.object(fireworks_finetune, "_fetch_all_deployments") as mock_fetch,
+        patch.object(
+            fireworks_finetune, "_deploy_server", return_value=True
+        ) as mock_deploy,
+    ):
+        result = await fireworks_finetune._check_or_deploy_server()
+
+        # Verify method returned True (from _deploy_server)
+        assert result is True
+
+        # Verify _fetch_all_deployments was NOT called
+        mock_fetch.assert_not_called()
+
+        # Verify _deploy_server was called directly
+        mock_deploy.assert_called_once()
+
+
+async def test_check_or_deploy_server_deploy_fails(fireworks_finetune, mock_api_key):
+    # Test when deployment fails
+
+    # Ensure no fine_tune_model_id is set
+    fireworks_finetune.datamodel.fine_tune_model_id = None
+
+    with (
+        patch.object(
+            fireworks_finetune, "_deploy_server", return_value=False
+        ) as mock_deploy,
+    ):
+        result = await fireworks_finetune._check_or_deploy_server()
+
+        # Verify method returned False (from _deploy_server)
+        assert result is False
+
+        # Verify _deploy_server was called
+        mock_deploy.assert_called_once()
+
+
+async def test_fetch_all_deployments_invalid_json(fireworks_finetune, mock_api_key):
+    # Test with invalid JSON response (missing 'deployments' key)
+    invalid_response = MagicMock(spec=httpx.Response)
+    invalid_response.status_code = 200
+    invalid_response.json.return_value = {
+        "some_other_key": "value",
+        # No 'deployments' key
+    }
+    invalid_response.text = '{"some_other_key": "value"}'
+
+    with patch("httpx.AsyncClient") as mock_client_class:
+        mock_client = AsyncMock()
+        mock_client.get.return_value = invalid_response
+        mock_client_class.return_value.__aenter__.return_value = mock_client
+
+        with pytest.raises(
+            ValueError,
+            match="Invalid response from Fireworks. Expected list of deployments in 'deployments' key",
+        ):
+            await fireworks_finetune._fetch_all_deployments()
+
+        # Verify API was called
+        mock_client.get.assert_called_once()