diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py index f0ebae0477..75a5001030 100644 --- a/vertexai/language_models/_language_models.py +++ b/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, }