feat: LLM - Added new regions for tuning and tuned model inference

PiperOrigin-RevId: 576294181

Author: Ark-kun

vertexai/language_models/_language_models.py

@@ -190,8 +190,7 @@ def tune_model(
                 Learning rate to use in tuning.
             learning_rate_multiplier: Learning rate multiplier to use in tuning.
             tuning_job_location: GCP location where the tuning job should be run.
-                Only "europe-west4" and "us-central1" locations are supported for now.
-            tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
+            tuned_model_location: GCP location where the tuned model should be deployed.
             model_display_name: Custom display name for the tuned model.
             tuning_evaluation_spec: Specification for the model evaluation during tuning.
             default_context: The context to use for all training samples by default.
@@ -294,8 +293,7 @@ def _tune_model(
                 See https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-models#dataset_format
             tuning_parameters: Tuning pipeline parameter values.
             tuning_job_location: GCP location where the tuning job should be run.
-                Only "europe-west4" and "us-central1" locations are supported for now.
-            tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
+            tuned_model_location: GCP location where the tuned model should be deployed.
             model_display_name: Custom display name for the tuned model.
 
         Returns:
@@ -312,9 +310,10 @@ def _tune_model(
                 "Please specify the tuning job location (`tuning_job_location`)."
                 f"Tuning is supported in the following locations: {_TUNING_LOCATIONS}"
             )
-        if tuned_model_location != _TUNED_MODEL_LOCATION:
+        if tuned_model_location not in _TUNED_MODEL_LOCATIONS:
             raise ValueError(
-                f'Model deployment is only supported in the following locations: tuned_model_location="{_TUNED_MODEL_LOCATION}"'
+                "Tuned model deployment is only supported in the following locations: "
+                f"{_TUNED_MODEL_LOCATIONS}"
             )
         model_info = _model_garden_models._get_model_info(
             model_id=self._model_id,
@@ -329,6 +328,7 @@ def _tune_model(
             tuning_parameters=tuning_parameters,
             model_display_name=model_display_name,
             tuning_job_location=tuning_job_location,
+            tuned_model_location=tuned_model_location,
         )
 
         job = _LanguageModelTuningJob(
@@ -369,8 +369,7 @@ def tune_model(
             train_steps: Number of training batches to tune on (batch size is 8 samples).
             learning_rate_multiplier: Learning rate multiplier to use in tuning.
             tuning_job_location: GCP location where the tuning job should be run.
-                Only "europe-west4" and "us-central1" locations are supported for now.
-            tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
+            tuned_model_location: GCP location where the tuned model should be deployed.
             model_display_name: Custom display name for the tuned model.
             tuning_evaluation_spec: Specification for the model evaluation during tuning.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
@@ -436,8 +435,7 @@ def tune_model(
                 Learning rate to use in tuning.
             learning_rate_multiplier: Learning rate multiplier to use in tuning.
             tuning_job_location: GCP location where the tuning job should be run.
-                Only "europe-west4" and "us-central1" locations are supported for now.
-            tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
+            tuned_model_location: GCP location where the tuned model should be deployed.
             model_display_name: Custom display name for the tuned model.
             tuning_evaluation_spec: Specification for the model evaluation during tuning.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
@@ -503,8 +501,7 @@ def tune_model(
                 Learning rate to use in tuning.
             learning_rate_multiplier: Learning rate multiplier to use in tuning.
             tuning_job_location: GCP location where the tuning job should be run.
-                Only "europe-west4" and "us-central1" locations are supported for now.
-            tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
+            tuned_model_location: GCP location where the tuned model should be deployed.
             model_display_name: Custom display name for the tuned model.
             default_context: The context to use for all training samples by default.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
@@ -570,8 +567,7 @@ def tune_model(
                 Learning rate to use in tuning.
             learning_rate_multiplier: Learning rate multiplier to use in tuning.
             tuning_job_location: GCP location where the tuning job should be run.
-                Only "europe-west4" and "us-central1" locations are supported for now.
-            tuned_model_location: GCP location where the tuned model should be deployed. Only "us-central1" is supported for now.
+            tuned_model_location: GCP location where the tuned model should be deployed.
             model_display_name: Custom display name for the tuned model.
             default_context: The context to use for all training samples by default.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
@@ -2697,9 +2693,29 @@ class _PreviewCodeGenerationModel(
 
 ###### Model tuning
 # Currently, tuning can only work in this location
-_TUNING_LOCATIONS = ("europe-west4", "us-central1")
-# Currently, deployment can only work in this location
-_TUNED_MODEL_LOCATION = "us-central1"
+
+_SUPPORTED_LOCATIONS = [
+    # 1
+    "us-central1",
+    "europe-west4",
+    "asia-southeast1",
+    # 2
+    "us-west1",
+    "europe-west3",
+    "europe-west2",
+    "asia-northeast1",
+    # 3
+    "us-east4",
+    "us-west4",
+    "northamerica-northeast1",
+    "europe-west9",
+    "europe-west1",
+    "asia-northeast3",
+]
+
+_TUNING_LOCATIONS = _SUPPORTED_LOCATIONS
+# Currently, deployment can only work in these locations
+_TUNED_MODEL_LOCATIONS = _SUPPORTED_LOCATIONS
 
 
 class _LanguageModelTuningJob:
@@ -2771,7 +2787,7 @@ def _list_tuned_model_names(model_id: str) -> List[str]:
     tuned_models = aiplatform.Model.list(
         filter=f'labels.{_TUNING_BASE_MODEL_ID_LABEL_KEY}="{model_id.replace("@", "-")}"',
         # TODO(b/275444096): Remove the explicit location once models are deployed to the user's selected location
-        location=_TUNED_MODEL_LOCATION,
+        location=aiplatform_initializer.global_config.location,
     )
     model_names = [model.resource_name for model in tuned_models]
     return model_names
@@ -2789,8 +2805,9 @@ def _launch_tuning_job(
     model_id: str,
     tuning_pipeline_uri: str,
     tuning_parameters: Dict[str, Any],
+    tuning_job_location: str,
+    tuned_model_location: str,
     model_display_name: Optional[str] = None,
-    tuning_job_location: str = _TUNING_LOCATIONS[0],
 ) -> aiplatform.PipelineJob:
     output_dir_uri = _generate_tuned_model_dir_uri(model_id=model_id)
     if isinstance(training_data, str):
@@ -2836,7 +2853,7 @@ def _launch_tuning_job(
         "project": aiplatform_initializer.global_config.project,
         # TODO(b/275444096): Remove the explicit location once tuning can happen in all regions
         # "location": aiplatform_initializer.global_config.location,
-        "location": _TUNED_MODEL_LOCATION,
+        "location": tuned_model_location,
         "large_model_reference": model_id,
         "model_display_name": model_display_name,
     }