From 693ab3759caa4da2ff33276d3eb30dbcb4af3661 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 2 Aug 2024 17:25:03 -0700
Subject: [PATCH] add vertex ft support

---
 litellm/fine_tuning/main.py | 40 +++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py
index 1400d0b1ad..c8c23c3ffc 100644
--- a/litellm/fine_tuning/main.py
+++ b/litellm/fine_tuning/main.py
@@ -25,6 +25,7 @@ from litellm.llms.fine_tuning_apis.openai import (
     FineTuningJobCreate,
     OpenAIFineTuningAPI,
 )
+from litellm.llms.fine_tuning_apis.vertex_ai import VertexFineTuningAPI
 from litellm.types.llms.openai import Hyperparameters
 from litellm.types.router import *
 from litellm.utils import supports_httpx_timeout
@@ -32,6 +33,7 @@ from litellm.utils import supports_httpx_timeout
 ####### ENVIRONMENT VARIABLES ###################
 openai_fine_tuning_apis_instance = OpenAIFineTuningAPI()
 azure_fine_tuning_apis_instance = AzureOpenAIFineTuningAPI()
+vertex_fine_tuning_apis_instance = VertexFineTuningAPI()
 #################################################
 
 
@@ -43,7 +45,7 @@ async def acreate_fine_tuning_job(
     validation_file: Optional[str] = None,
     integrations: Optional[List[str]] = None,
     seed: Optional[int] = None,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
@@ -96,7 +98,7 @@ def create_fine_tuning_job(
     validation_file: Optional[str] = None,
     integrations: Optional[List[str]] = None,
     seed: Optional[int] = None,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
@@ -221,6 +223,39 @@ def create_fine_tuning_job(
                 max_retries=optional_params.max_retries,
                 _is_async=_is_async,
             )
+        elif custom_llm_provider == "vertex_ai":
+            api_base = optional_params.api_base
+            vertex_ai_project = (
+                optional_params.vertex_project
+                or litellm.vertex_project
+                or get_secret("VERTEXAI_PROJECT")
+            )
+            vertex_ai_location = (
+                optional_params.vertex_location
+                or litellm.vertex_location
+                or get_secret("VERTEXAI_LOCATION")
+            )
+            vertex_credentials = optional_params.vertex_credentials or get_secret(
+                "VERTEXAI_CREDENTIALS"
+            )
+            create_fine_tuning_job_data = FineTuningJobCreate(
+                model=model,
+                training_file=training_file,
+                hyperparameters=hyperparameters,
+                suffix=suffix,
+                validation_file=validation_file,
+                integrations=integrations,
+                seed=seed,
+            )
+            response = vertex_fine_tuning_apis_instance.create_fine_tuning_job(
+                _is_async=_is_async,
+                create_fine_tuning_job_data=create_fine_tuning_job_data,
+                vertex_credentials=vertex_credentials,
+                vertex_project=vertex_ai_project,
+                vertex_location=vertex_ai_location,
+                timeout=timeout,
+                api_base=api_base,
+            )
         else:
             raise litellm.exceptions.BadRequestError(
                 message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
@@ -236,6 +271,7 @@ def create_fine_tuning_job(
             )
         return response
     except Exception as e:
+        verbose_logger.error("got exception in create_fine_tuning_job=%s", str(e))
         raise e