diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 8c78cef833..70d8cbade7 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -32,6 +32,62 @@ logger = logging.getLogger("agentplatform_genai.modelgarden") +def _CreateBatchPredictionJobConfig_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + + if getv(from_object, ["display_name"]) is not None: + setv(parent_object, ["displayName"], getv(from_object, ["display_name"])) + + if getv(from_object, ["input_config"]) is not None: + setv( + parent_object, + ["inputConfig"], + _ModelGardenBatchJobSource_to_vertex( + getv(from_object, ["input_config"]), to_object + ), + ) + + if getv(from_object, ["output_config"]) is not None: + setv( + parent_object, + ["outputConfig"], + _ModelGardenBatchJobDestination_to_vertex( + getv(from_object, ["output_config"]), to_object + ), + ) + + if getv(from_object, ["dedicated_resources"]) is not None: + setv( + parent_object, + ["dedicatedResources"], + getv(from_object, ["dedicated_resources"]), + ) + + return to_object + + +def _CreateBatchPredictionJobParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["parent"]) is not None: + setv(to_object, ["_url", "parent"], getv(from_object, ["parent"])) + + if getv(from_object, ["model"]) is not None: + setv(to_object, ["model"], getv(from_object, ["model"])) + + if getv(from_object, ["config"]) is not None: + _CreateBatchPredictionJobConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) + + return to_object + + def _DeployRequestParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -98,6 +154,17 @@ def _ExportPublisherModelRequestParameters_to_vertex( return to_object +def _GetBatchPredictionJobParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + return to_object + + def _GetDeployOperationParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -213,6 +280,154 @@ def _ListPublisherModelsRequestParameters_to_vertex( return to_object +def _ModelGardenBatchJobDestination_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["predictionsFormat"]) is not None: + setv(to_object, ["format"], getv(from_object, ["predictionsFormat"])) + + if getv(from_object, ["gcsDestination", "outputUriPrefix"]) is not None: + setv( + to_object, + ["gcs_uri"], + getv(from_object, ["gcsDestination", "outputUriPrefix"]), + ) + + if getv(from_object, ["bigqueryDestination", "outputUri"]) is not None: + setv( + to_object, + ["bigquery_uri"], + getv(from_object, ["bigqueryDestination", "outputUri"]), + ) + + return to_object + + +def _ModelGardenBatchJobDestination_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["format"]) is not None: + setv(to_object, ["predictionsFormat"], getv(from_object, ["format"])) + + if getv(from_object, ["gcs_uri"]) is not None: + setv( + to_object, + ["gcsDestination", "outputUriPrefix"], + getv(from_object, ["gcs_uri"]), + ) + + if getv(from_object, ["bigquery_uri"]) is not None: + setv( + to_object, + ["bigqueryDestination", "outputUri"], + getv(from_object, ["bigquery_uri"]), + ) + + return to_object + + +def _ModelGardenBatchJobSource_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["instancesFormat"]) is not None: + setv(to_object, ["format"], getv(from_object, ["instancesFormat"])) + + if getv(from_object, ["gcsSource", "uris"]) is not None: + setv(to_object, ["gcs_uri"], getv(from_object, ["gcsSource", "uris"])) + + if getv(from_object, ["bigquerySource", "inputUri"]) is not None: + setv( + to_object, + ["bigquery_uri"], + getv(from_object, ["bigquerySource", "inputUri"]), + ) + + return to_object + + +def _ModelGardenBatchJobSource_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["format"]) is not None: + setv(to_object, ["instancesFormat"], getv(from_object, ["format"])) + + if getv(from_object, ["gcs_uri"]) is not None: + setv(to_object, ["gcsSource", "uris"], getv(from_object, ["gcs_uri"])) + + if getv(from_object, ["bigquery_uri"]) is not None: + setv( + to_object, + ["bigquerySource", "inputUri"], + getv(from_object, ["bigquery_uri"]), + ) + + return to_object + + +def _ModelGardenBatchPredictionJob_from_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["name"]) is not None: + setv(to_object, ["name"], getv(from_object, ["name"])) + + if getv(from_object, ["displayName"]) is not None: + setv(to_object, ["display_name"], getv(from_object, ["displayName"])) + + if getv(from_object, ["model"]) is not None: + setv(to_object, ["model"], getv(from_object, ["model"])) + + if getv(from_object, ["inputConfig"]) is not None: + setv( + to_object, + ["input_config"], + _ModelGardenBatchJobSource_from_vertex( + getv(from_object, ["inputConfig"]), to_object + ), + ) + + if getv(from_object, ["outputConfig"]) is not None: + setv( + to_object, + ["output_config"], + _ModelGardenBatchJobDestination_from_vertex( + getv(from_object, ["outputConfig"]), to_object + ), + ) + + if getv(from_object, ["outputInfo"]) is not None: + setv(to_object, ["output_info"], getv(from_object, ["outputInfo"])) + + if getv(from_object, ["state"]) is not None: + setv(to_object, ["state"], getv(from_object, ["state"])) + + if getv(from_object, ["error"]) is not None: + setv(to_object, ["error"], getv(from_object, ["error"])) + + if getv(from_object, ["createTime"]) is not None: + setv(to_object, ["create_time"], getv(from_object, ["createTime"])) + + if getv(from_object, ["startTime"]) is not None: + setv(to_object, ["start_time"], getv(from_object, ["startTime"])) + + if getv(from_object, ["endTime"]) is not None: + setv(to_object, ["end_time"], getv(from_object, ["endTime"])) + + if getv(from_object, ["updateTime"]) is not None: + setv(to_object, ["update_time"], getv(from_object, ["updateTime"])) + + return to_object + + def _RecommendSpecConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -777,6 +992,160 @@ def get_deploy_publisher_model_operation( self._api_client._verify_response(return_value) return return_value + def _create_batch_prediction_job( + self, + *, + parent: str, + model: str, + config: Optional[types.CreateBatchPredictionJobConfigOrDict] = None, + ) -> types.ModelGardenBatchPredictionJob: + """ + Creates a batch prediction job (internal). + """ + + parameter_model = types._CreateBatchPredictionJobParameters( + parent=parent, + model=model, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateBatchPredictionJobParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/batchPredictionJobs".format_map(request_url_dict) + else: + path = "{parent}/batchPredictionJobs" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ModelGardenBatchPredictionJob_from_vertex(response_dict) + + return_value = types.ModelGardenBatchPredictionJob._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + def get_batch_prediction_job( + self, + *, + name: str, + config: Optional[types.GetBatchPredictionJobConfigOrDict] = None, + ) -> types.ModelGardenBatchPredictionJob: + """ + Fetches the current state of a ``batch_predict`` job. + """ + + parameter_model = types._GetBatchPredictionJobParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetBatchPredictionJobParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("get", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ModelGardenBatchPredictionJob_from_vertex(response_dict) + + return_value = types.ModelGardenBatchPredictionJob._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + # Fallbacks for ``ExportOpenModelConfig`` when the caller does not # override them. 2h matches the legacy SDK's blocking ``.export()`` and is # generous enough for large open weights (e.g. Gemma 3 27B). @@ -1894,6 +2263,210 @@ def export_open_model( ) return operation.response.destination_uri + @staticmethod + def _complete_bq_uri(uri: Optional[str], project: str) -> str: + """Completes a BigQuery URI prefix to a full table URI. + + Mirrors ``vertexai.batch_prediction.BatchPredictionJob``: a missing + project, dataset or table is filled with the client project, the + ``gen_ai_batch_prediction`` dataset and a timestamped ``predictions_*`` + table respectively. + + Args: + uri: A ``'bq://{project}[.{dataset}[.{table}]]'`` URI, or None. + project: The project to use when ``uri`` is None. + + Returns: + A ``'bq://{project}.{dataset}.{table}'`` URI. + + Raises: + ValueError: If ``uri`` has more than three parts or does not start with + ``'bq://'``. + """ + import datetime + + parts = uri.split(".") if uri else [] + if len(parts) > 3: + raise ValueError( + f"Invalid BigQuery URI: {uri}. Supported format: " + "'bq://projectId.bqDatasetId.bqTableId'." + ) + project_part = parts[0] if parts else f"bq://{project}" + if not project_part.startswith("bq://"): + raise ValueError(f"BigQuery URI must start with 'bq://': {uri}.") + dataset = parts[1] if len(parts) >= 2 else "gen_ai_batch_prediction" + table = ( + parts[2] + if len(parts) >= 3 + else "predictions_" + + datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S_%f") + ) + return f"{project_part}.{dataset}.{table}" + + @staticmethod + def _prepare_batch_predict_request( + model: str, + src: Union[str, list[str]], + config: types.BatchPredictConfig, + project: str, + ) -> tuple[str, types.CreateBatchPredictionJobConfig]: + """Translates ``batch_predict`` arguments to a ``CreateBatchPredictionJob`` request. + + Shared by the sync and async surfaces. Infers the input and output + formats from the URI schemes, as ``OpenModel.batch_predict`` does. + + Returns: + The ``(publisher_model_name, request_config)`` pair. + + Raises: + ValueError: If ``model`` is a Hugging Face model ID or not a valid + publisher model name, if ``src`` or ``dest`` use an unsupported + scheme, if ``dest`` is missing for a Cloud Storage ``src``, or if + accelerator or replica settings are given without ``machine_type``. + """ + import datetime + + if ModelGarden._is_hugging_face_model(model): + raise ValueError( + f"batch_predict does not support Hugging Face model IDs (got " + f"{model!r}); pass a Model Garden publisher model such as " + "'google/gemma3@gemma-3-12b-it'." + ) + if not src: + raise ValueError("src must be specified.") + publisher_model_name = ModelGarden._reconcile_model_name(model) + + uris = [src] if isinstance(src, str) else list(src) + if all(uri.startswith("gs://") for uri in uris): + input_config = types.ModelGardenBatchJobSource(format="jsonl", gcs_uri=uris) + elif uris[0].startswith("bq://"): + if len(uris) > 1: + raise ValueError("Multiple BigQuery input tables are not supported.") + input_config = types.ModelGardenBatchJobSource( + format="bigquery", bigquery_uri=uris[0] + ) + else: + raise ValueError( + f"Unsupported src: {src}. Supported formats: " + "'gs://path/to/input/data.jsonl' and " + "'bq://projectId.bqDatasetId.bqTableId'." + ) + + dest = config.dest + if dest and dest.startswith("gs://"): + output_config = types.ModelGardenBatchJobDestination( + format="jsonl", gcs_uri=dest + ) + elif dest and dest.startswith("bq://"): + output_config = types.ModelGardenBatchJobDestination( + format="bigquery", + bigquery_uri=ModelGarden._complete_bq_uri(dest, project), + ) + elif dest: + raise ValueError( + f"Unsupported dest: {dest}. Supported formats: " + "'gs://path/to/output/data' and 'bq://projectId.bqDatasetId'." + ) + elif input_config.format == "bigquery": + output_config = types.ModelGardenBatchJobDestination( + format="bigquery", + bigquery_uri=ModelGarden._complete_bq_uri(None, project), + ) + else: + raise ValueError("dest must be specified when src is a Cloud Storage URI.") + + # Like the legacy SDK, the machine shape is only sent with a machine + # type; reject the settings that would otherwise be dropped silently. + dedicated_resources = None + if config.machine_type: + dedicated_resources = types.BatchDedicatedResources( + machine_spec=types.MachineSpec( + machine_type=config.machine_type, + accelerator_type=config.accelerator_type, + accelerator_count=config.accelerator_count, + ), + starting_replica_count=config.starting_replica_count, + max_replica_count=config.max_replica_count, + ) + else: + dropped = [ + name + for name in ( + "accelerator_type", + "accelerator_count", + "starting_replica_count", + "max_replica_count", + ) + if getattr(config, name) + ] + if dropped: + raise ValueError( + f"{', '.join(dropped)} require machine_type to also be set." + ) + + display_name = config.display_name or ( + "model-garden-batch-predict-" + + datetime.datetime.now().strftime("%Y%m%d%H%M%S") + ) + return publisher_model_name, types.CreateBatchPredictionJobConfig( + display_name=display_name, + input_config=input_config, + output_config=output_config, + dedicated_resources=dedicated_resources, + ) + + def batch_predict( + self, + *, + model: str, + src: Union[str, list[str]], + config: Optional[types.BatchPredictConfigOrDict] = None, + ) -> types.ModelGardenBatchPredictionJob: + """Runs batch prediction on a Model Garden publisher model. + + Submits a batch prediction job and returns without waiting for it to + finish. Use ``get_batch_prediction_job`` with the returned job's ``name`` + to track its ``state``. + + Args: + model: The publisher model to run. Accepts the full resource name + ``'publishers/{publisher}/models/{model}@{version}'`` or the + simplified ``'{publisher}/{model}@{version}'`` (or without the + ``@{version}``). Hugging Face model IDs are not supported. + src: Cloud Storage URI(s) of JSONL input files (e.g. + ``'gs://my-bucket/input/prompts.jsonl'``), or a single BigQuery + table URI (e.g. ``'bq://my-project.my_dataset.my_table'``). + config: Optional configuration (output destination, display name, + machine shape). See ``BatchPredictConfig``. + + Returns: + The created ``ModelGardenBatchPredictionJob``. + + Raises: + ValueError: If ``model`` is not a valid publisher model name or is a + Hugging Face model ID, if ``src`` or ``config.dest`` is malformed, if + ``config.dest`` is missing for a Cloud Storage ``src``, or if + accelerator or replica settings are given without + ``config.machine_type``. + """ + if config is None: + config = types.BatchPredictConfig() + elif isinstance(config, dict): + config = types.BatchPredictConfig.model_validate(config) + + publisher_model_name, api_config = ModelGarden._prepare_batch_predict_request( + model, src, config, self._api_client.project + ) + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + return self._create_batch_prediction_job( + parent=parent, + model=publisher_model_name, + config=api_config, + ) + class AsyncModelGarden(_api_module.BaseModule): """Model Garden module.""" @@ -2437,6 +3010,164 @@ async def get_deploy_publisher_model_operation( self._api_client._verify_response(return_value) return return_value + async def _create_batch_prediction_job( + self, + *, + parent: str, + model: str, + config: Optional[types.CreateBatchPredictionJobConfigOrDict] = None, + ) -> types.ModelGardenBatchPredictionJob: + """ + Creates a batch prediction job (internal). + """ + + parameter_model = types._CreateBatchPredictionJobParameters( + parent=parent, + model=model, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _CreateBatchPredictionJobParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/batchPredictionJobs".format_map(request_url_dict) + else: + path = "{parent}/batchPredictionJobs" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ModelGardenBatchPredictionJob_from_vertex(response_dict) + + return_value = types.ModelGardenBatchPredictionJob._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def get_batch_prediction_job( + self, + *, + name: str, + config: Optional[types.GetBatchPredictionJobConfigOrDict] = None, + ) -> types.ModelGardenBatchPredictionJob: + """ + Fetches the current state of a ``batch_predict`` job. + """ + + parameter_model = types._GetBatchPredictionJobParameters( + name=name, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError( + "This method is only supported in Gemini Enterprise Agent Platform mode, not in Gemini Developer API mode." + ) + else: + request_dict = _GetBatchPredictionJobParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{name}".format_map(request_url_dict) + else: + path = "{name}" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "get", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + if self._api_client.vertexai: + response_dict = _ModelGardenBatchPredictionJob_from_vertex(response_dict) + + return_value = types.ModelGardenBatchPredictionJob._from_response( + response=response_dict, + kwargs=( + { + "config": { + "response_schema": getattr( + parameter_model.config, "response_schema", None + ), + "response_json_schema": getattr( + parameter_model.config, "response_json_schema", None + ), + "include_all_fields": getattr( + parameter_model.config, "include_all_fields", None + ), + } + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + async def _list_all_publisher_models( self, api_config: types.ListPublisherModelsConfig, @@ -2845,3 +3576,29 @@ async def export_open_model( f"{operation!r}" ) return operation.response.destination_uri + + async def batch_predict( + self, + *, + model: str, + src: Union[str, list[str]], + config: Optional[types.BatchPredictConfigOrDict] = None, + ) -> types.ModelGardenBatchPredictionJob: + """Async variant of ``ModelGarden.batch_predict``.""" + if config is None: + config = types.BatchPredictConfig() + elif isinstance(config, dict): + config = types.BatchPredictConfig.model_validate(config) + + publisher_model_name, api_config = ModelGarden._prepare_batch_predict_request( + model, src, config, self._api_client.project + ) + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + return await self._create_batch_prediction_job( + parent=parent, + model=publisher_model_name, + config=api_config, + ) diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index 394f649938..223ce0b189 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -28,6 +28,7 @@ from .common import _AssessDatasetParameters from .common import _CancelQueryJobRuntimeRequestParameters from .common import _CheckQueryJobRuntimeRequestParameters +from .common import _CreateBatchPredictionJobParameters from .common import _CreateDatasetParameters from .common import _CreateDatasetVersionParameters from .common import _CreateEvaluationExperimentParameters @@ -80,6 +81,7 @@ from .common import _GenerateLossClustersParameters from .common import _GenerateMemoriesRequestParameters from .common import _GenerateUserScenariosParameters +from .common import _GetBatchPredictionJobParameters from .common import _GetCorpusOperationParameters from .common import _GetCustomJobParameters from .common import _GetCustomJobParameters @@ -247,6 +249,12 @@ from .common import AutoscalingMetricSpec from .common import AutoscalingMetricSpecDict from .common import AutoscalingMetricSpecOrDict +from .common import BatchDedicatedResources +from .common import BatchDedicatedResourcesDict +from .common import BatchDedicatedResourcesOrDict +from .common import BatchPredictConfig +from .common import BatchPredictConfigDict +from .common import BatchPredictConfigOrDict from .common import BatchPredictionResourceUsageAssessmentConfig from .common import BatchPredictionResourceUsageAssessmentConfigDict from .common import BatchPredictionResourceUsageAssessmentConfigOrDict @@ -329,6 +337,9 @@ from .common import CorpusStatus from .common import CorpusStatusDict from .common import CorpusStatusOrDict +from .common import CreateBatchPredictionJobConfig +from .common import CreateBatchPredictionJobConfigDict +from .common import CreateBatchPredictionJobConfigOrDict from .common import CreateDatasetConfig from .common import CreateDatasetConfigDict from .common import CreateDatasetConfigOrDict @@ -870,6 +881,9 @@ from .common import GenerateUserScenariosResponse from .common import GenerateUserScenariosResponseDict from .common import GenerateUserScenariosResponseOrDict +from .common import GetBatchPredictionJobConfig +from .common import GetBatchPredictionJobConfigDict +from .common import GetBatchPredictionJobConfigOrDict from .common import GetCorpusOperationConfig from .common import GetCorpusOperationConfigDict from .common import GetCorpusOperationConfigOrDict @@ -1370,6 +1384,18 @@ from .common import ModelContainerSpec from .common import ModelContainerSpecDict from .common import ModelContainerSpecOrDict +from .common import ModelGardenBatchJobDestination +from .common import ModelGardenBatchJobDestinationDict +from .common import ModelGardenBatchJobDestinationOrDict +from .common import ModelGardenBatchJobSource +from .common import ModelGardenBatchJobSourceDict +from .common import ModelGardenBatchJobSourceOrDict +from .common import ModelGardenBatchPredictionJob +from .common import ModelGardenBatchPredictionJobDict +from .common import ModelGardenBatchPredictionJobOrDict +from .common import ModelGardenBatchPredictionOutputInfo +from .common import ModelGardenBatchPredictionOutputInfoDict +from .common import ModelGardenBatchPredictionOutputInfoOrDict from .common import ModelProvider from .common import MultimodalDataset from .common import MultimodalDatasetDict @@ -3982,6 +4008,27 @@ "GetDeployOperationConfig", "GetDeployOperationConfigDict", "GetDeployOperationConfigOrDict", + "ModelGardenBatchJobSource", + "ModelGardenBatchJobSourceDict", + "ModelGardenBatchJobSourceOrDict", + "ModelGardenBatchJobDestination", + "ModelGardenBatchJobDestinationDict", + "ModelGardenBatchJobDestinationOrDict", + "BatchDedicatedResources", + "BatchDedicatedResourcesDict", + "BatchDedicatedResourcesOrDict", + "CreateBatchPredictionJobConfig", + "CreateBatchPredictionJobConfigDict", + "CreateBatchPredictionJobConfigOrDict", + "ModelGardenBatchPredictionOutputInfo", + "ModelGardenBatchPredictionOutputInfoDict", + "ModelGardenBatchPredictionOutputInfoOrDict", + "ModelGardenBatchPredictionJob", + "ModelGardenBatchPredictionJobDict", + "ModelGardenBatchPredictionJobOrDict", + "GetBatchPredictionJobConfig", + "GetBatchPredictionJobConfigDict", + "GetBatchPredictionJobConfigOrDict", "CreateRuntimeFeedbackEntryConfig", "CreateRuntimeFeedbackEntryConfigDict", "CreateRuntimeFeedbackEntryConfigOrDict", @@ -4366,6 +4413,9 @@ "DeployCustomModelConfig", "DeployCustomModelConfigDict", "DeployCustomModelConfigOrDict", + "BatchPredictConfig", + "BatchPredictConfigDict", + "BatchPredictConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", @@ -4607,6 +4657,8 @@ "_GetExportPublisherModelOperationParameters", "_DeployRequestParameters", "_GetDeployOperationParameters", + "_CreateBatchPredictionJobParameters", + "_GetBatchPredictionJobParameters", "_CreateRuntimeFeedbackEntryRequestParameters", "_DeleteRuntimeFeedbackEntryRequestParameters", "_GetRuntimeFeedbackRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 98219d5ab8..0f90875cc7 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -25824,6 +25824,382 @@ class _GetDeployOperationParametersDict(TypedDict, total=False): ] +class ModelGardenBatchJobSource(_common.BaseModel): + """Input configuration of a Model Garden batch prediction job.""" + + format: Optional[str] = Field( + default=None, + description="""Storage format of the input. One of ``'jsonl'`` or + ``'bigquery'``.""", + ) + gcs_uri: Optional[list[str]] = Field( + default=None, description="""Cloud Storage URIs of the input files.""" + ) + bigquery_uri: Optional[str] = Field( + default=None, description="""BigQuery URI of the input table.""" + ) + + +class ModelGardenBatchJobSourceDict(TypedDict, total=False): + """Input configuration of a Model Garden batch prediction job.""" + + format: Optional[str] + """Storage format of the input. One of ``'jsonl'`` or + ``'bigquery'``.""" + + gcs_uri: Optional[list[str]] + """Cloud Storage URIs of the input files.""" + + bigquery_uri: Optional[str] + """BigQuery URI of the input table.""" + + +ModelGardenBatchJobSourceOrDict = Union[ + ModelGardenBatchJobSource, ModelGardenBatchJobSourceDict +] + + +class ModelGardenBatchJobDestination(_common.BaseModel): + """Output configuration of a Model Garden batch prediction job.""" + + format: Optional[str] = Field( + default=None, + description="""Storage format of the output. One of ``'jsonl'`` or + ``'bigquery'``.""", + ) + gcs_uri: Optional[str] = Field( + default=None, + description="""Cloud Storage URI prefix the predictions are written under.""", + ) + bigquery_uri: Optional[str] = Field( + default=None, + description="""BigQuery table URI the predictions are written to.""", + ) + + +class ModelGardenBatchJobDestinationDict(TypedDict, total=False): + """Output configuration of a Model Garden batch prediction job.""" + + format: Optional[str] + """Storage format of the output. One of ``'jsonl'`` or + ``'bigquery'``.""" + + gcs_uri: Optional[str] + """Cloud Storage URI prefix the predictions are written under.""" + + bigquery_uri: Optional[str] + """BigQuery table URI the predictions are written to.""" + + +ModelGardenBatchJobDestinationOrDict = Union[ + ModelGardenBatchJobDestination, ModelGardenBatchJobDestinationDict +] + + +class BatchDedicatedResources(_common.BaseModel): + """A description of resources that are used for performing batch operations, are dedicated to a Model, and need manual configuration.""" + + flex_start: Optional[FlexStart] = Field( + default=None, + description="""Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""", + ) + machine_spec: Optional[MachineSpec] = Field( + default=None, + description="""Required. Immutable. The specification of a single machine.""", + ) + max_replica_count: Optional[int] = Field( + default=None, + description="""Immutable. The maximum number of machine replicas the batch operation may be scaled to. The default value is 10.""", + ) + spot: Optional[bool] = Field( + default=None, + description="""Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""", + ) + starting_replica_count: Optional[int] = Field( + default=None, + description="""Immutable. The number of machine replicas used at the start of the batch operation. If not set, Agent Platform decides starting number, not greater than max_replica_count""", + ) + + +class BatchDedicatedResourcesDict(TypedDict, total=False): + """A description of resources that are used for performing batch operations, are dedicated to a Model, and need manual configuration.""" + + flex_start: Optional[FlexStartDict] + """Optional. Immutable. If set, use DWS resource to schedule the deployment workload. reference: (https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler)""" + + machine_spec: Optional[MachineSpecDict] + """Required. Immutable. The specification of a single machine.""" + + max_replica_count: Optional[int] + """Immutable. The maximum number of machine replicas the batch operation may be scaled to. The default value is 10.""" + + spot: Optional[bool] + """Optional. If true, schedule the deployment workload on [spot VMs](https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms).""" + + starting_replica_count: Optional[int] + """Immutable. The number of machine replicas used at the start of the batch operation. If not set, Agent Platform decides starting number, not greater than max_replica_count""" + + +BatchDedicatedResourcesOrDict = Union[ + BatchDedicatedResources, BatchDedicatedResourcesDict +] + + +class CreateBatchPredictionJobConfig(_common.BaseModel): + """RPC-level config for ``create_batch_prediction_job``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + display_name: Optional[str] = Field(default=None, description="""""") + input_config: Optional[ModelGardenBatchJobSource] = Field( + default=None, description="""""" + ) + output_config: Optional[ModelGardenBatchJobDestination] = Field( + default=None, description="""""" + ) + dedicated_resources: Optional[BatchDedicatedResources] = Field( + default=None, description="""""" + ) + + +class CreateBatchPredictionJobConfigDict(TypedDict, total=False): + """RPC-level config for ``create_batch_prediction_job``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + display_name: Optional[str] + """""" + + input_config: Optional[ModelGardenBatchJobSourceDict] + """""" + + output_config: Optional[ModelGardenBatchJobDestinationDict] + """""" + + dedicated_resources: Optional[BatchDedicatedResourcesDict] + """""" + + +CreateBatchPredictionJobConfigOrDict = Union[ + CreateBatchPredictionJobConfig, CreateBatchPredictionJobConfigDict +] + + +class _CreateBatchPredictionJobParameters(_common.BaseModel): + """Parameters for ``create_batch_prediction_job``.""" + + parent: Optional[str] = Field(default=None, description="""""") + model: Optional[str] = Field(default=None, description="""""") + config: Optional[CreateBatchPredictionJobConfig] = Field( + default=None, description="""""" + ) + + +class _CreateBatchPredictionJobParametersDict(TypedDict, total=False): + """Parameters for ``create_batch_prediction_job``.""" + + parent: Optional[str] + """""" + + model: Optional[str] + """""" + + config: Optional[CreateBatchPredictionJobConfigDict] + """""" + + +_CreateBatchPredictionJobParametersOrDict = Union[ + _CreateBatchPredictionJobParameters, _CreateBatchPredictionJobParametersDict +] + + +class ModelGardenBatchPredictionOutputInfo(_common.BaseModel): + """Where a Model Garden batch prediction job wrote its output.""" + + gcs_output_directory: Optional[str] = Field( + default=None, + description="""Cloud Storage directory holding the prediction files.""", + ) + bigquery_output_dataset: Optional[str] = Field( + default=None, description="""BigQuery dataset holding the output table.""" + ) + bigquery_output_table: Optional[str] = Field( + default=None, + description="""Name of the BigQuery table holding the predictions.""", + ) + + +class ModelGardenBatchPredictionOutputInfoDict(TypedDict, total=False): + """Where a Model Garden batch prediction job wrote its output.""" + + gcs_output_directory: Optional[str] + """Cloud Storage directory holding the prediction files.""" + + bigquery_output_dataset: Optional[str] + """BigQuery dataset holding the output table.""" + + bigquery_output_table: Optional[str] + """Name of the BigQuery table holding the predictions.""" + + +ModelGardenBatchPredictionOutputInfoOrDict = Union[ + ModelGardenBatchPredictionOutputInfo, ModelGardenBatchPredictionOutputInfoDict +] + + +class ModelGardenBatchPredictionJob(_common.BaseModel): + """A batch prediction job running a Model Garden publisher model. + + Declares only the ``BatchPredictionJob`` fields needed to identify the job + and track it to completion, rather than patching the full discovery class + and its monitoring and explanation subtrees. + """ + + name: Optional[str] = Field( + default=None, + description="""Resource name of the job, in the format + ``'projects/{project}/locations/{location}/batchPredictionJobs/{id}'``.""", + ) + display_name: Optional[str] = Field( + default=None, description="""User-defined name of the job.""" + ) + model: Optional[str] = Field( + default=None, description="""Publisher model resource name the job runs.""" + ) + input_config: Optional[ModelGardenBatchJobSource] = Field( + default=None, description="""Input configuration of the job.""" + ) + output_config: Optional[ModelGardenBatchJobDestination] = Field( + default=None, description="""Output configuration of the job.""" + ) + output_info: Optional[ModelGardenBatchPredictionOutputInfo] = Field( + default=None, + description="""Where the job wrote its output. Populated once the job runs.""", + ) + state: Optional[genai_types.JobState] = Field( + default=None, description="""Detailed state of the job.""" + ) + error: Optional[genai_types.GoogleRpcStatus] = Field( + default=None, + description="""Only populated when the job state is ``JOB_STATE_FAILED`` or + ``JOB_STATE_CANCELLED``.""", + ) + create_time: Optional[datetime.datetime] = Field( + default=None, description="""Time the job was created.""" + ) + start_time: Optional[datetime.datetime] = Field( + default=None, + description="""Time the job first entered ``JOB_STATE_RUNNING``.""", + ) + end_time: Optional[datetime.datetime] = Field( + default=None, description="""Time the job entered a terminal state.""" + ) + update_time: Optional[datetime.datetime] = Field( + default=None, description="""Time the job was most recently updated.""" + ) + + +class ModelGardenBatchPredictionJobDict(TypedDict, total=False): + """A batch prediction job running a Model Garden publisher model. + + Declares only the ``BatchPredictionJob`` fields needed to identify the job + and track it to completion, rather than patching the full discovery class + and its monitoring and explanation subtrees. + """ + + name: Optional[str] + """Resource name of the job, in the format + ``'projects/{project}/locations/{location}/batchPredictionJobs/{id}'``.""" + + display_name: Optional[str] + """User-defined name of the job.""" + + model: Optional[str] + """Publisher model resource name the job runs.""" + + input_config: Optional[ModelGardenBatchJobSourceDict] + """Input configuration of the job.""" + + output_config: Optional[ModelGardenBatchJobDestinationDict] + """Output configuration of the job.""" + + output_info: Optional[ModelGardenBatchPredictionOutputInfoDict] + """Where the job wrote its output. Populated once the job runs.""" + + state: Optional[genai_types.JobState] + """Detailed state of the job.""" + + error: Optional[genai_types.GoogleRpcStatus] + """Only populated when the job state is ``JOB_STATE_FAILED`` or + ``JOB_STATE_CANCELLED``.""" + + create_time: Optional[datetime.datetime] + """Time the job was created.""" + + start_time: Optional[datetime.datetime] + """Time the job first entered ``JOB_STATE_RUNNING``.""" + + end_time: Optional[datetime.datetime] + """Time the job entered a terminal state.""" + + update_time: Optional[datetime.datetime] + """Time the job was most recently updated.""" + + +ModelGardenBatchPredictionJobOrDict = Union[ + ModelGardenBatchPredictionJob, ModelGardenBatchPredictionJobDict +] + + +class GetBatchPredictionJobConfig(_common.BaseModel): + """Config for ``get_batch_prediction_job``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetBatchPredictionJobConfigDict(TypedDict, total=False): + """Config for ``get_batch_prediction_job``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetBatchPredictionJobConfigOrDict = Union[ + GetBatchPredictionJobConfig, GetBatchPredictionJobConfigDict +] + + +class _GetBatchPredictionJobParameters(_common.BaseModel): + """Parameters for ``get_batch_prediction_job``.""" + + name: Optional[str] = Field( + default=None, + description="""Resource name of the job, as returned by ``batch_predict``.""", + ) + config: Optional[GetBatchPredictionJobConfig] = Field( + default=None, description="""""" + ) + + +class _GetBatchPredictionJobParametersDict(TypedDict, total=False): + """Parameters for ``get_batch_prediction_job``.""" + + name: Optional[str] + """Resource name of the job, as returned by ``batch_predict``.""" + + config: Optional[GetBatchPredictionJobConfigDict] + """""" + + +_GetBatchPredictionJobParametersOrDict = Union[ + _GetBatchPredictionJobParameters, _GetBatchPredictionJobParametersDict +] + + class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): """Config for creating a Feedback Entry.""" @@ -32024,6 +32400,77 @@ class DeployCustomModelConfigDict(TypedDict, total=False): ] +class BatchPredictConfig(_common.BaseModel): + """Config for ``batch_predict``.""" + + dest: Optional[str] = Field( + default=None, + description="""Cloud Storage or BigQuery URI prefix for the output predictions, + e.g. ``'gs://my-bucket/output/'`` or ``'bq://my-project.my_dataset'``. + A BigQuery dataset without a table gets a timestamped + ``predictions_*`` table. Required for a Cloud Storage ``src``; for a + BigQuery ``src`` it defaults to + ``'bq://{project}.gen_ai_batch_prediction.predictions_{timestamp}'``.""", + ) + display_name: Optional[str] = Field( + default=None, + description="""User-defined name of the batch prediction job, up to 128 + characters. Defaults to a timestamped name.""", + ) + machine_type: Optional[str] = Field( + default=None, + description="""Machine type for the job (e.g. ``'g2-standard-12'``). The + accelerator and replica settings are only sent when this is set.""", + ) + accelerator_type: Optional[str] = Field( + default=None, description="""Accelerator type (e.g. ``'NVIDIA_L4'``).""" + ) + accelerator_count: Optional[int] = Field( + default=None, description="""Number of accelerators per replica.""" + ) + starting_replica_count: Optional[int] = Field( + default=None, description="""Number of replicas the job starts with.""" + ) + max_replica_count: Optional[int] = Field( + default=None, description="""Maximum number of replicas the job may scale to.""" + ) + + +class BatchPredictConfigDict(TypedDict, total=False): + """Config for ``batch_predict``.""" + + dest: Optional[str] + """Cloud Storage or BigQuery URI prefix for the output predictions, + e.g. ``'gs://my-bucket/output/'`` or ``'bq://my-project.my_dataset'``. + A BigQuery dataset without a table gets a timestamped + ``predictions_*`` table. Required for a Cloud Storage ``src``; for a + BigQuery ``src`` it defaults to + ``'bq://{project}.gen_ai_batch_prediction.predictions_{timestamp}'``.""" + + display_name: Optional[str] + """User-defined name of the batch prediction job, up to 128 + characters. Defaults to a timestamped name.""" + + machine_type: Optional[str] + """Machine type for the job (e.g. ``'g2-standard-12'``). The + accelerator and replica settings are only sent when this is set.""" + + accelerator_type: Optional[str] + """Accelerator type (e.g. ``'NVIDIA_L4'``).""" + + accelerator_count: Optional[int] + """Number of accelerators per replica.""" + + starting_replica_count: Optional[int] + """Number of replicas the job starts with.""" + + max_replica_count: Optional[int] + """Maximum number of replicas the job may scale to.""" + + +BatchPredictConfigOrDict = Union[BatchPredictConfig, BatchPredictConfigDict] + + class DeployOption(_common.BaseModel): """A verified deploy option for a model.""" diff --git a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py index b122588248..c4887c4975 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -524,6 +524,90 @@ def test_deploy_custom_model_no_wait_dict_config(client): assert operation.name +_BATCH_MODEL = "google/gemma3@gemma-3-1b-it" +_BATCH_SRC = "gs://sdk-mg-custom-model-replay/batch-predict/input.jsonl" +_BATCH_DEST = "gs://sdk-mg-custom-model-replay/batch-predict/output/" + + +def _batch_predict_config(display_name): + # A fixed display name keeps the recorded request body deterministic. + return types.BatchPredictConfig( + dest=_BATCH_DEST, + display_name=display_name, + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + starting_replica_count=1, + max_replica_count=1, + ) + + +def test_batch_predict_returns_job(client): + """batch_predict returns the created job without waiting for it. + + Also pins the wire mapping: the JSONL source, GCS destination and machine + shape must serialize to ``inputConfig``, ``outputConfig`` and + ``dedicatedResources`` exactly as recorded. + """ + job = client.model_garden.batch_predict( + model=_BATCH_MODEL, + src=_BATCH_SRC, + config=_batch_predict_config("sdk-replay-batch-predict"), + ) + assert isinstance(job, types.ModelGardenBatchPredictionJob) + assert "/batchPredictionJobs/" in job.name + assert job.display_name == "sdk-replay-batch-predict" + assert job.model.startswith("publishers/google/models/gemma3") + assert job.input_config.gcs_uri == [_BATCH_SRC] + assert job.output_config.gcs_uri == _BATCH_DEST + + +def test_get_batch_prediction_job(client): + """get_batch_prediction_job fetches a job created by batch_predict.""" + job = client.model_garden.batch_predict( + model=_BATCH_MODEL, + src=_BATCH_SRC, + config=_batch_predict_config("sdk-replay-batch-predict-get"), + ) + fetched = client.model_garden.get_batch_prediction_job(name=job.name) + assert isinstance(fetched, types.ModelGardenBatchPredictionJob) + assert fetched.name == job.name + assert fetched.state is not None + + +def test_batch_predict_dict_config(client): + """The config may be passed as a plain dict.""" + job = client.model_garden.batch_predict( + model=_BATCH_MODEL, + src=_BATCH_SRC, + config={ + "dest": _BATCH_DEST, + "display_name": "sdk-replay-batch-predict-dict", + "machine_type": "g2-standard-12", + "accelerator_type": "NVIDIA_L4", + "accelerator_count": 1, + }, + ) + assert isinstance(job, types.ModelGardenBatchPredictionJob) + assert job.display_name == "sdk-replay-batch-predict-dict" + + +def test_batch_predict_missing_dest_raises(client): + """A Cloud Storage src needs an explicit dest; rejected client-side.""" + with pytest.raises(ValueError, match="dest must be specified"): + client.model_garden.batch_predict(model=_BATCH_MODEL, src=_BATCH_SRC) + + +def test_batch_predict_hugging_face_model_raises(client): + """Hugging Face model IDs are rejected client-side.""" + with pytest.raises(ValueError, match="does not support Hugging Face"): + client.model_garden.batch_predict( + model="meta-llama/Llama-3.3-70B-Instruct", + src=_BATCH_SRC, + config=types.BatchPredictConfig(dest=_BATCH_DEST), + ) + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), @@ -690,3 +774,29 @@ async def test_deploy_custom_model_async_no_wait_returns_operation(client): assert isinstance(operation, types.DeployModelOperation) assert "/operations/" in operation.name assert operation.done is not True + + +@pytest.mark.asyncio +async def test_batch_predict_async_returns_job(client): + """Async batch_predict returns the created job without waiting for it.""" + job = await client.aio.model_garden.batch_predict( + model=_BATCH_MODEL, + src=_BATCH_SRC, + config=_batch_predict_config("sdk-replay-batch-predict-async"), + ) + assert isinstance(job, types.ModelGardenBatchPredictionJob) + assert "/batchPredictionJobs/" in job.name + + +@pytest.mark.asyncio +async def test_get_batch_prediction_job_async(client): + """Async get_batch_prediction_job fetches a job created by batch_predict.""" + job = await client.aio.model_garden.batch_predict( + model=_BATCH_MODEL, + src=_BATCH_SRC, + config=_batch_predict_config("sdk-replay-batch-predict-get-async"), + ) + fetched = await client.aio.model_garden.get_batch_prediction_job( + name=job.name + ) + assert fetched.name == job.name diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index f5ebbb5218..82a92e1d06 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -3141,6 +3141,310 @@ def test_deploy_custom_model_async_raises_on_operation_error(mock_async_client): ) +# --- batch_predict ------------------------------------------------------- + +_TEST_BATCH_JOB_NAME = ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + "/batchPredictionJobs/789" +) +_TEST_BATCH_SRC = "gs://my-bucket/input/prompts.jsonl" +_TEST_BATCH_DEST = "gs://my-bucket/output/" + + +def _make_batch_job(): + return types.ModelGardenBatchPredictionJob( + name=_TEST_BATCH_JOB_NAME, + state="JOB_STATE_PENDING", + ) + + +def _patch_batch_rpc(mg_client): + return mock.patch.object( + mg_client, + "_create_batch_prediction_job", + return_value=_make_batch_job(), + ) + + +def _patch_async_batch_rpc(mg_client): + return mock.patch.object( + mg_client, + "_create_batch_prediction_job", + new=mock.AsyncMock(return_value=_make_batch_job()), + ) + + +def test_batch_predict_returns_created_job(mock_client): + """The job is returned as soon as it is created; nothing is polled.""" + with _patch_batch_rpc(mock_client): + job = mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + assert isinstance(job, types.ModelGardenBatchPredictionJob) + assert job.name == _TEST_BATCH_JOB_NAME + + +def test_batch_predict_gcs_request_shape(mock_client): + """A GCS src/dest pair is sent as JSONL in and JSONL out.""" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig( + dest=_TEST_BATCH_DEST, display_name="my-job" + ), + ) + kwargs = m_rpc.call_args.kwargs + assert kwargs["parent"] == ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + ) + assert kwargs["model"] == "publishers/google/models/gemma3@gemma-3-12b-it" + api_config = kwargs["config"] + assert api_config.display_name == "my-job" + assert api_config.input_config.format == "jsonl" + assert api_config.input_config.gcs_uri == [_TEST_BATCH_SRC] + assert api_config.input_config.bigquery_uri is None + assert api_config.output_config.format == "jsonl" + assert api_config.output_config.gcs_uri == _TEST_BATCH_DEST + assert api_config.dedicated_resources is None + + +def test_batch_predict_full_resource_name_passed_through(mock_client): + """A full publisher model resource name is sent unchanged.""" + model = "publishers/google/models/gemma3@gemma-3-12b-it" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=model, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + assert m_rpc.call_args.kwargs["model"] == model + + +def test_batch_predict_multiple_gcs_sources(mock_client): + """A list of GCS URIs is sent as-is.""" + srcs = ["gs://my-bucket/a.jsonl", "gs://my-bucket/b.jsonl"] + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src=srcs, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + assert m_rpc.call_args.kwargs["config"].input_config.gcs_uri == srcs + + +def test_batch_predict_bigquery_full_table_dest_kept(mock_client): + """A fully qualified BigQuery table dest is sent unchanged.""" + dest = "bq://my-project.my_dataset.my_table" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src="bq://my-project.my_dataset.input", + config=types.BatchPredictConfig(dest=dest), + ) + api_config = m_rpc.call_args.kwargs["config"] + assert api_config.input_config.format == "bigquery" + assert api_config.input_config.bigquery_uri == ( + "bq://my-project.my_dataset.input" + ) + assert api_config.input_config.gcs_uri is None + assert api_config.output_config.format == "bigquery" + assert api_config.output_config.bigquery_uri == dest + + +def test_batch_predict_bigquery_dataset_dest_gets_table(mock_client): + """A BigQuery dataset dest gets a timestamped predictions table.""" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest="bq://my-project.my_dataset"), + ) + output_uri = m_rpc.call_args.kwargs["config"].output_config.bigquery_uri + assert output_uri.startswith("bq://my-project.my_dataset.predictions_") + + +def test_batch_predict_bigquery_src_defaults_dest(mock_client): + """A BigQuery src without dest writes to the default dataset.""" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, src="bq://my-project.my_dataset.input" + ) + output_config = m_rpc.call_args.kwargs["config"].output_config + assert output_config.format == "bigquery" + assert output_config.bigquery_uri.startswith( + f"bq://{_TEST_PROJECT}.gen_ai_batch_prediction.predictions_" + ) + + +def test_batch_predict_machine_shape_forwarded(mock_client): + """Machine shape and replica counts ride on dedicated_resources.""" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig( + dest=_TEST_BATCH_DEST, + machine_type="g2-standard-48", + accelerator_type="NVIDIA_L4", + accelerator_count=4, + starting_replica_count=1, + max_replica_count=2, + ), + ) + resources = m_rpc.call_args.kwargs["config"].dedicated_resources + assert resources.machine_spec.machine_type == "g2-standard-48" + assert resources.machine_spec.accelerator_type == "NVIDIA_L4" + assert resources.machine_spec.accelerator_count == 4 + assert resources.starting_replica_count == 1 + assert resources.max_replica_count == 2 + + +def test_batch_predict_default_display_name(mock_client): + """Without a display name, a timestamped one is generated.""" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + assert m_rpc.call_args.kwargs["config"].display_name.startswith( + "model-garden-batch-predict-" + ) + + +def test_batch_predict_dict_config_validated(mock_client): + """The config may be passed as a plain dict.""" + with _patch_batch_rpc(mock_client) as m_rpc: + mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config={"dest": _TEST_BATCH_DEST, "display_name": "dict-job"}, + ) + api_config = m_rpc.call_args.kwargs["config"] + assert api_config.display_name == "dict-job" + assert api_config.output_config.gcs_uri == _TEST_BATCH_DEST + + +@pytest.mark.parametrize( + "kwargs,match", + [ + ( + {"src": _TEST_BATCH_SRC}, + "dest must be specified", + ), + ( + {"src": ""}, + "src must be specified", + ), + ( + {"src": "s3://bucket/input.jsonl", "dest": _TEST_BATCH_DEST}, + "Unsupported src", + ), + ( + {"src": ["bq://p.d.a", "bq://p.d.b"]}, + "Multiple BigQuery input tables", + ), + ( + {"src": [_TEST_BATCH_SRC, "bq://p.d.t"], "dest": _TEST_BATCH_DEST}, + "Unsupported src", + ), + ( + {"src": _TEST_BATCH_SRC, "dest": "s3://bucket/output/"}, + "Unsupported dest", + ), + ( + {"src": _TEST_BATCH_SRC, "dest": "bq://p.d.t.extra"}, + "Invalid BigQuery URI", + ), + ], +) +def test_batch_predict_invalid_io_raises_before_rpc(mock_client, kwargs, match): + """Malformed src/dest combinations are rejected locally.""" + config = {"dest": kwargs["dest"]} if "dest" in kwargs else None + with _patch_batch_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match=match): + mock_client.batch_predict( + model=_OPEN_MODEL, src=kwargs["src"], config=config + ) + m_rpc.assert_not_called() + + +def test_batch_predict_hugging_face_model_raises(mock_client): + """Hugging Face model IDs are rejected locally.""" + with _patch_batch_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="does not support Hugging Face"): + mock_client.batch_predict( + model="meta-llama/Llama-3.3-70B-Instruct", + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + m_rpc.assert_not_called() + + +def test_batch_predict_invalid_model_name_raises(mock_client): + """Model Registry names are not publisher models.""" + with _patch_batch_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="not a valid publisher model name"): + mock_client.batch_predict( + model="projects/p/locations/l/models/123", + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + m_rpc.assert_not_called() + + +def test_batch_predict_accelerator_without_machine_type_raises(mock_client): + """Machine settings that would be dropped without machine_type raise.""" + with _patch_batch_rpc(mock_client) as m_rpc: + with pytest.raises( + ValueError, + match="accelerator_type, max_replica_count require machine_type", + ): + mock_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig( + dest=_TEST_BATCH_DEST, + accelerator_type="NVIDIA_L4", + max_replica_count=2, + ), + ) + m_rpc.assert_not_called() + + +def test_batch_predict_async_returns_created_job(mock_async_client): + """The async surface shapes the same request and returns the job.""" + with _patch_async_batch_rpc(mock_async_client) as m_rpc: + job = asyncio.run( + mock_async_client.batch_predict( + model=_OPEN_MODEL, + src=_TEST_BATCH_SRC, + config=types.BatchPredictConfig(dest=_TEST_BATCH_DEST), + ) + ) + assert isinstance(job, types.ModelGardenBatchPredictionJob) + assert job.name == _TEST_BATCH_JOB_NAME + kwargs = m_rpc.call_args.kwargs + assert kwargs["model"] == "publishers/google/models/gemma3@gemma-3-12b-it" + assert kwargs["config"].input_config.gcs_uri == [_TEST_BATCH_SRC] + assert kwargs["config"].output_config.gcs_uri == _TEST_BATCH_DEST + + +def test_batch_predict_async_missing_dest_raises(mock_async_client): + """Async validation matches the sync surface.""" + with _patch_async_batch_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="dest must be specified"): + asyncio.run( + mock_async_client.batch_predict( + model=_OPEN_MODEL, src=_TEST_BATCH_SRC + ) + ) + m_rpc.assert_not_called() + + # --- Private (pre-GA) surface ------------------------------------------ # # No direct unit tests: `agentplatform/private/**` is excluded from the