diff --git a/agentplatform/_genai/model_garden.py b/agentplatform/_genai/model_garden.py index 511f958367..871553491f 100644 --- a/agentplatform/_genai/model_garden.py +++ b/agentplatform/_genai/model_garden.py @@ -22,14 +22,62 @@ from google.genai import _api_module from google.genai import _common +from google.genai import types as genai_types from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv +from . import _operations_utils from . import types logger = logging.getLogger("agentplatform_genai.modelgarden") +def _ExportPublisherModelConfig_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, ["destination"]) is not None: + setv(parent_object, ["destination"], getv(from_object, ["destination"])) + + return to_object + + +def _ExportPublisherModelRequestParameters_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, ["name"]) is not None: + setv(to_object, ["_url", "name"], getv(from_object, ["name"])) + + if getv(from_object, ["config"]) is not None: + _ExportPublisherModelConfig_to_vertex( + getv(from_object, ["config"]), to_object + ) + + return to_object + + +def _GetExportPublisherModelOperationParameters_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, ["operation_name"]) is not None: + setv( + to_object, + ["_url", "operationName"], + getv(from_object, ["operation_name"]), + ) + + return to_object + + def _GetPublisherModelConfig_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -156,59 +204,59 @@ def _RecommendSpecRequestParameters_to_vertex( class ModelGarden(_api_module.BaseModule): - """Model Garden module.""" + """Model Garden module.""" - def _list_publisher_models( + def _list_publisher_models( self, *, parent: Optional[str] = None, config: Optional[types.ListPublisherModelsConfigOrDict] = None, ) -> types.ListPublisherModelsResponse: - """ + """ Lists publisher models (internal). """ - parameter_model = types._ListPublisherModelsRequestParameters( + parameter_model = types._ListPublisherModelsRequestParameters( parent=parent, config=config, ) - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( + 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 = _ListPublisherModelsRequestParameters_to_vertex( + else: + request_dict = _ListPublisherModelsRequestParameters_to_vertex( parameter_model ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{parent}/models".format_map(request_url_dict) - else: - path = "{parent}/models" - - 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 ( + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/models".format_map(request_url_dict) + else: + path = "{parent}/models" + + 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 + http_options = parameter_model.config.http_options - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) + 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 = self._api_client.request("get", path, request_dict, http_options) - response_dict = {} if not response.body else json.loads(response.body) + response_dict = {} if not response.body else json.loads(response.body) - return_value = types.ListPublisherModelsResponse._from_response( + return_value = types.ListPublisherModelsResponse._from_response( response=response_dict, kwargs=( { @@ -229,57 +277,57 @@ def _list_publisher_models( ), ) - self._api_client._verify_response(return_value) - return return_value + self._api_client._verify_response(return_value) + return return_value - def _get_publisher_model( + def _get_publisher_model( self, *, name: str, config: Optional[types.GetPublisherModelConfigOrDict] = None ) -> types.PublisherModel: - """ + """ Gets a publisher model (internal). """ - parameter_model = types._GetPublisherModelRequestParameters( + parameter_model = types._GetPublisherModelRequestParameters( name=name, config=config, ) - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( + 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 = _GetPublisherModelRequestParameters_to_vertex( + else: + request_dict = _GetPublisherModelRequestParameters_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 ( + 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 + http_options = parameter_model.config.http_options - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) + 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 = self._api_client.request("get", path, request_dict, http_options) - response_dict = {} if not response.body else json.loads(response.body) + response_dict = {} if not response.body else json.loads(response.body) - return_value = types.PublisherModel._from_response( + return_value = types.PublisherModel._from_response( response=response_dict, kwargs=( { @@ -300,60 +348,60 @@ def _get_publisher_model( ), ) - self._api_client._verify_response(return_value) - return return_value + self._api_client._verify_response(return_value) + return return_value - def _recommend_spec( + def _recommend_spec( self, *, parent: str, gcs_uri: str, config: Optional[types.RecommendSpecConfigOrDict] = None, ) -> types.RecommendSpecResponse: - """ + """ Recommends spec for a custom model (internal). """ - parameter_model = types._RecommendSpecRequestParameters( + parameter_model = types._RecommendSpecRequestParameters( parent=parent, gcs_uri=gcs_uri, config=config, ) - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( + 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 = _RecommendSpecRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{parent}:recommendSpec".format_map(request_url_dict) - else: - path = "{parent}:recommendSpec" - - 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 ( + else: + request_dict = _RecommendSpecRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}:recommendSpec".format_map(request_url_dict) + else: + path = "{parent}:recommendSpec" + + 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 + http_options = parameter_model.config.http_options - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) + 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 = self._api_client.request("post", path, request_dict, http_options) - response_dict = {} if not response.body else json.loads(response.body) + response_dict = {} if not response.body else json.loads(response.body) - return_value = types.RecommendSpecResponse._from_response( + return_value = types.RecommendSpecResponse._from_response( response=response_dict, kwargs=( { @@ -374,865 +422,1120 @@ def _recommend_spec( ), ) - self._api_client._verify_response(return_value) - return return_value - - @staticmethod - def _build_filter_str( - model_filter: Optional[str], - include_hugging_face_models: bool, - deployable_only: bool, - ) -> str: - """Builds the filter string for the ListPublisherModels API. - - Args: - model_filter: Optional substring to match against model IDs and display - names (case-insensitive). - include_hugging_face_models: Whether to include HuggingFace models. If - True, uses ``is_hf_wildcard(true)``; otherwise ``is_hf_wildcard(false)``. - deployable_only: Whether to restrict to models with verified deployment - configurations via the ``VERIFIED_DEPLOYMENT_SUCCEED`` label. - - Returns: - A filter string suitable for the ``filter`` parameter of the - ListPublisherModels API. - """ - import re - - if include_hugging_face_models: - filter_str = "is_hf_wildcard(true)" - if deployable_only: - filter_str += ( - " AND labels.VERIFIED_DEPLOYMENT_CONFIG=VERIFIED_DEPLOYMENT_SUCCEED" - ) - else: - filter_str = "is_hf_wildcard(false)" - - if model_filter: - escaped = re.escape(model_filter) - filter_str = ( - f'{filter_str} AND (model_user_id=~"(?i).*{escaped}.*"' - f' OR display_name=~"(?i).*{escaped}.*")' - ) - - return filter_str - - @staticmethod - def _format_model_name( - model: types.PublisherModel, - include_hugging_face_models: bool, - ) -> str: - """Formats a PublisherModel into a human-readable model name string. - - Args: - model: The PublisherModel to format. - include_hugging_face_models: Whether HuggingFace models are included in - the listing. Controls whether the ``@version`` suffix is appended. - - Returns: - A formatted model name string in one of the following formats: - - - ``'{publisher}/{model}@{version}'`` when - ``include_hugging_face_models`` is False. - - ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True. - """ - import re - - name = model.name or "" - formatted = re.sub(r"publishers/(hf-|)|models/", "", name) - if include_hugging_face_models: - return formatted - return formatted + "@" + (model.version_id or "") - - @staticmethod - def _has_deploy_config(model: types.PublisherModel) -> bool: - """Checks whether a model has verified deployment configurations. - - Args: - model: The PublisherModel to check. - - Returns: - True if the model has at least one entry in - ``supported_actions.multi_deploy_vertex``. - """ - return bool( - model.supported_actions - and model.supported_actions.multi_deploy_vertex - and model.supported_actions.multi_deploy_vertex.multi_deploy_vertex - ) - - @staticmethod - def _reconcile_model_name(model_name: str) -> str: - """Normalizes a model name into a publisher model resource name. - - Args: - model_name: A Model Garden model resource name in the format - ``'publishers/{publisher}/models/{model}@{version}'``, a simplified name - in the format ``'{publisher}/{model}@{version}'`` (or without - ``@{version}``), or a Hugging Face model ID ``'{organization}/{model}'``. - - Returns: - The resource name in the format - ``'publishers/{publisher}/models/{model}@{version}'``. - - Raises: - ValueError: If ``model_name`` is not a valid publisher model name. - """ - import re - - model_name = model_name.lower() # Hugging Face IDs are lower-case. - # A full resource name must carry an @version, matching the legacy SDK; a - # versionless full name is not accepted. - full_match = re.match( - r"^publishers/(?P[^/]+)/models/(?P[^@]+)@(?P[^@]+)$", - model_name, - ) - if full_match: - return ( - f"publishers/{full_match.group('publisher')}/models/" - f"{full_match.group('model')}@{full_match.group('version')}" - ) - # Reject Model Registry names; they would otherwise match the simplified - # branch and be silently mangled. - if re.match(r"^projects/.+/locations/.+/models/.+$", model_name): - raise ValueError(f"`{model_name}` is not a valid publisher model name") - simplified_match = re.match( - r"^(?P[^/]+)/(?P[^@]+)(?:@(?P.+))?$", - model_name, - ) - if simplified_match: - model = simplified_match.group("model") - if simplified_match.group("version"): - model = f"{model}@{simplified_match.group('version')}" - return f"publishers/{simplified_match.group('publisher')}/models/{model}" - raise ValueError(f"`{model_name}` is not a valid publisher model name") - - @staticmethod - def _is_hugging_face_model(model_name: str) -> bool: - """Returns whether a model name looks like a Hugging Face model ID. - - Matches the bare ``'{organization}/{model}'`` shape (a single slash and no - ``@version``), e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``. - - Args: - model_name: The model name to inspect. - - Returns: - True if ``model_name`` matches the Hugging Face ID shape. - """ - import re - - return bool(re.match(r"^(?P[^/]+)/(?P[^/@]+)$", model_name)) - - @staticmethod - def _matches_filter( - value: Optional[str], - model_filter: Optional[Union[str, list[str]]], - ) -> bool: - """Returns whether ``value`` matches the (optional) keyword filter. - - Mirrors the legacy SDK: the filter may be a single keyword or a list of - keywords, and matching is a case-insensitive substring test where the value - matches if it contains *any* of the keywords. - - Args: - value: The field value to test (e.g. a machine type), or None. - model_filter: A keyword, a list of keywords, or None (no filtering). - - Returns: - True if there is no filter, or if ``value`` contains any of the keywords. - """ - if not model_filter: - return True - if value is None: - return False - keywords = [model_filter] if isinstance(model_filter, str) else model_filter - value_lower = value.lower() - return any(keyword.lower() in value_lower for keyword in keywords) - - @staticmethod - def _extract_and_filter_deploy_options( - publisher_model: types.PublisherModel, - machine_type_filter: Optional[Union[str, list[str]]] = None, - accelerator_type_filter: Optional[Union[str, list[str]]] = None, - serving_container_image_uri_filter: Optional[Union[str, list[str]]] = None, - ) -> list[types.DeployOption]: - """Extracts and filters deploy options from a publisher model. - - Args: - publisher_model: The publisher model to extract deploy options from. - machine_type_filter: Optional case-insensitive keyword (or list of - keywords) matched against the machine type; an option is kept if its - machine type contains any of them (e.g. ``'g2'`` or ``['n1', 'g2']``). - accelerator_type_filter: Optional case-insensitive keyword (or list of - keywords) matched against the accelerator type (e.g. ``'L4'`` or - ``['T4', 'L4']``). - serving_container_image_uri_filter: Optional case-insensitive keyword (or - list of keywords) matched against the serving container image URI - (e.g. ``'vllm'`` or ``['vllm', 'tgi']``). - - Returns: - A list of ``DeployOption`` objects matching the provided filters. - - Raises: - ValueError: If the model does not support deployment, or if no deploy - options remain after applying the filters. - """ - if not ( - publisher_model.supported_actions - and publisher_model.supported_actions.multi_deploy_vertex - and publisher_model.supported_actions.multi_deploy_vertex.multi_deploy_vertex - ): - raise ValueError( - "Model does not support deployment. " - "Use `list_deployable_models()` to find supported models." - ) - - options = ( - publisher_model.supported_actions.multi_deploy_vertex.multi_deploy_vertex - ) - result = [] - for opt in options: - container = opt.container_spec.image_uri if opt.container_spec else None - machine = ( - opt.dedicated_resources.machine_spec - if opt.dedicated_resources - else None - ) - machine_type = machine.machine_type if machine else None - - # Restore the proto3 defaults the JSON transport drops, so structured - # output matches the gRPC SDK on CPU/TPU machines. - accelerator_enum = machine.accelerator_type if machine else None - accelerator_value = accelerator_enum.value if accelerator_enum else None - has_accelerator = ( - accelerator_value is not None - and accelerator_value != "ACCELERATOR_TYPE_UNSPECIFIED" - ) - if machine: - accelerator_type = ( - accelerator_value - if accelerator_value is not None - else "ACCELERATOR_TYPE_UNSPECIFIED" - ) - accelerator_count = ( - machine.accelerator_count - if machine.accelerator_count is not None - else 0 - ) - else: - accelerator_type = None - accelerator_count = None - - if not ModelGarden._matches_filter(machine_type, machine_type_filter): - continue - # ACCELERATOR_TYPE_UNSPECIFIED means "no accelerator" and never matches. - if accelerator_type_filter and not has_accelerator: - continue - if not ModelGarden._matches_filter( - accelerator_type, accelerator_type_filter - ): - continue - if not ModelGarden._matches_filter( - container, serving_container_image_uri_filter - ): - continue - - result.append( - types.DeployOption( - option_name=opt.deploy_task_name, - serving_container_image_uri=container, - machine_type=machine_type, - accelerator_type=accelerator_type, - accelerator_count=accelerator_count, - ) - ) - - if not result: - raise ValueError("No deploy options found.") - - return result - - @staticmethod - def _format_concise_deploy_options( - options: list[types.DeployOption], - ) -> str: - """Formats deploy options into a human-readable string. - - Mirrors the legacy ``vertexai.model_garden`` SDK output: each option is - rendered as a ``[Option N: ]`` block followed by its non-null - fields (container image, machine type, accelerator type/count). - - Args: - options: The deploy options to format. - - Returns: - A human-readable, multi-line string describing the deploy options. - """ - fields = [ - "serving_container_image_uri", - "machine_type", - "accelerator_type", - "accelerator_count", - ] - blocks = [] - for i, option in enumerate(options): - if option.option_name: - header = f"[Option {i + 1}: {option.option_name}]\n" - else: - header = f"[Option {i + 1}]\n" - lines = [] - for field in fields: - value = getattr(option, field) - if value is None: - continue - if field == "accelerator_count": - lines.append(f" {field}={value},") - else: - lines.append(f' {field}="{value}",') - blocks.append(header + "\n".join(lines)) - return "\n\n".join(blocks) - - def _list_all_publisher_models( - self, - api_config: types.ListPublisherModelsConfig, - ) -> list[types.PublisherModel]: - """Fetches all pages of publisher models from the API. - - Args: - api_config: The configuration for the ListPublisherModels API call, - including filter and version settings. - - Returns: - A list of all ``PublisherModel`` objects across all pages. - """ - all_models = [] - page_token = None - while True: - if page_token: - api_config = types.ListPublisherModelsConfig( - filter=api_config.filter, - list_all_versions=api_config.list_all_versions, - page_token=page_token, - ) - response = self._list_publisher_models( - parent="publishers/*", - config=api_config, - ) - all_models.extend(response.publisher_models or []) - page_token = response.next_page_token - if not page_token: - break - return all_models - - def _list( - self, - model_filter: Optional[str], - include_hugging_face_models: Optional[bool], - deployable_only: bool, - ) -> list[str]: - """Shared implementation for listing models. - - Args: - model_filter: Optional substring to filter models by. - include_hugging_face_models: Whether to include HuggingFace models. - deployable_only: If True, only return models with deployment configs. - - Returns: - A list of formatted model name strings. - """ - include_hf = include_hugging_face_models is True - - filter_str = self._build_filter_str( - model_filter, include_hf, deployable_only=deployable_only - ) - - api_config = types.ListPublisherModelsConfig( - filter=filter_str, - list_all_versions=True, - ) - - models = self._list_all_publisher_models(api_config) - - if deployable_only: - # The VERIFIED_DEPLOYMENT_SUCCEED server filter only applies to HF - # models; filter native models client-side via multi_deploy_vertex. - models = [m for m in models if self._has_deploy_config(m)] - - return [self._format_model_name(m, include_hf) for m in models] - - def list_deployable_models( - self, - config: Optional[types.ListDeployableModelsConfigOrDict] = None, - ) -> list[str]: - """Lists models in Model Garden that support deployment. - - Returns models that have at least one verified deployment configuration. - When ``include_hugging_face_models`` is False (the default), - HuggingFace models are excluded from the results. - - Args: - config: Optional configuration for filtering results. Accepts a - ``ListDeployableModelsConfig`` instance or an equivalent dict. - - Returns: - A list of model name strings in the format - ``'{publisher}/{model}@{version}'`` (e.g. ``'google/gemma2@gemma-2-2b-it'``) - or ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True - (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). - """ - if config is None: - config = types.ListDeployableModelsConfig() - if isinstance(config, dict): - config = types.ListDeployableModelsConfig.model_validate(config) - - return self._list( - config.model_filter, - config.include_hugging_face_models, - deployable_only=True, - ) - - def list_models( - self, - config: Optional[types.ListModelGardenModelsConfigOrDict] = None, - ) -> list[str]: - """Lists all models available in Model Garden. - - Returns all models regardless of deployment support. When - ``include_hugging_face_models`` is False (the default), HuggingFace - models are excluded from the results. - - Args: - config: Optional configuration for filtering results. Accepts a - ``ListModelGardenModelsConfig`` instance or an equivalent dict. - - Returns: - A list of model name strings in the format - ``'{publisher}/{model}@{version}'`` (e.g. ``'google/gemma2@gemma-2-2b-it'``) - or ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True - (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). - """ - if config is None: - config = types.ListModelGardenModelsConfig() - if isinstance(config, dict): - config = types.ListModelGardenModelsConfig.model_validate(config) - - return self._list( - config.model_filter, - config.include_hugging_face_models, - deployable_only=False, - ) - - def list_publisher_model_deploy_options( - self, - model: str, - config: Optional[types.ListPublisherModelDeployOptionsConfigOrDict] = None, - ) -> Union[str, list[types.DeployOption]]: - """Lists the verified deploy options for a Model Garden publisher model. - - Supports Google open models (e.g. ``'google/gemma3@gemma-3-12b-it'``), - partner publisher models (e.g. - ``'mistralai/mistral-7b@mistral-7b-instruct-v0.2'``), and Hugging Face - model IDs (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). - - Args: - model: The publisher model to list deploy options for. Accepts the full - resource name ``'publishers/{publisher}/models/{model}@{version}'``, a - simplified ``'{publisher}/{model}@{version}'`` (or without the - ``@{version}``), or a Hugging Face model ID ``'{organization}/{model}'``. - config: Optional configuration for filtering the deploy options. Accepts a - ``ListPublisherModelDeployOptionsConfig`` instance or an equivalent - dict. - - Returns: - A list of ``DeployOption`` objects, one per verified deployment - configuration (container image URI, machine type, accelerator type and - count). If ``config.concise`` is True, returns a human-readable string - describing those deploy options instead. - - Raises: - ValueError: If ``model`` is not a valid publisher model name, if the - model does not support deployment, or if no deploy options match the - provided filters. - """ - if config is None: - config = types.ListPublisherModelDeployOptionsConfig() - if isinstance(config, dict): - config = types.ListPublisherModelDeployOptionsConfig.model_validate(config) - - get_publisher_model_config = types.GetPublisherModelConfig( - is_hugging_face_model=self._is_hugging_face_model(model), - include_equivalent_model_garden_model_deployment_configs=True, + self._api_client._verify_response(return_value) + return return_value + + def _export_publisher_model( + self, + *, + parent: str, + name: str, + config: Optional[types.ExportPublisherModelConfigOrDict] = None, + ) -> types.ExportModelOperation: + """Exports a publisher model (internal).""" + + parameter_model = types._ExportPublisherModelRequestParameters( + parent=parent, + 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 = _ExportPublisherModelRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/{name}:export".format_map(request_url_dict) + else: + path = "{parent}/{name}:export" + + 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) + + return_value = types.ExportModelOperation._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_export_publisher_model_operation( + self, + *, + operation_name: str, + config: Optional[ + types.GetExportPublisherModelOperationConfigOrDict + ] = None, + ) -> types.ExportModelOperation: + """Fetches the status of an in-flight ``export_open_model`` LRO.""" + + parameter_model = types._GetExportPublisherModelOperationParameters( + operation_name=operation_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 = _GetExportPublisherModelOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + 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) + + return_value = types.ExportModelOperation._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). + _DEFAULT_EXPORT_TIMEOUT_SECONDS = 2 * 60 * 60 + _DEFAULT_EXPORT_POLL_INTERVAL_SECONDS = 30 + + @staticmethod + def _build_filter_str( + model_filter: Optional[str], + include_hugging_face_models: bool, + deployable_only: bool, + ) -> str: + """Builds the filter string for the ListPublisherModels API. + + Args: + model_filter: Optional substring to match against model IDs and display + names (case-insensitive). + include_hugging_face_models: Whether to include HuggingFace models. If + True, uses ``is_hf_wildcard(true)``; otherwise + ``is_hf_wildcard(false)``. + deployable_only: Whether to restrict to models with verified deployment + configurations via the ``VERIFIED_DEPLOYMENT_SUCCEED`` label. + + Returns: + A filter string suitable for the ``filter`` parameter of the + ListPublisherModels API. + """ + import re + + if include_hugging_face_models: + filter_str = "is_hf_wildcard(true)" + if deployable_only: + filter_str += ( + " AND labels.VERIFIED_DEPLOYMENT_CONFIG=VERIFIED_DEPLOYMENT_SUCCEED" ) - publisher_model = self._get_publisher_model( - name=self._reconcile_model_name(model), - config=get_publisher_model_config, + else: + filter_str = "is_hf_wildcard(false)" + + if model_filter: + escaped = re.escape(model_filter) + filter_str = ( + f'{filter_str} AND (model_user_id=~"(?i).*{escaped}.*"' + f' OR display_name=~"(?i).*{escaped}.*")' + ) + + return filter_str + + @staticmethod + def _format_model_name( + model: types.PublisherModel, + include_hugging_face_models: bool, + ) -> str: + """Formats a PublisherModel into a human-readable model name string. + + Args: + model: The PublisherModel to format. + include_hugging_face_models: Whether HuggingFace models are included in + the listing. Controls whether the ``@version`` suffix is appended. + + Returns: + A formatted model name string in one of the following formats: + + - ``'{publisher}/{model}@{version}'`` when + ``include_hugging_face_models`` is False. + - ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True. + """ + import re + + name = model.name or "" + formatted = re.sub(r"publishers/(hf-|)|models/", "", name) + if include_hugging_face_models: + return formatted + return formatted + "@" + (model.version_id or "") + + @staticmethod + def _has_deploy_config(model: types.PublisherModel) -> bool: + """Checks whether a model has verified deployment configurations. + + Args: + model: The PublisherModel to check. + + Returns: + True if the model has at least one entry in + ``supported_actions.multi_deploy_vertex``. + """ + return bool( + model.supported_actions + and model.supported_actions.multi_deploy_vertex + and model.supported_actions.multi_deploy_vertex.multi_deploy_vertex + ) + + @staticmethod + def _reconcile_model_name(model_name: str) -> str: + """Normalizes a model name into a publisher model resource name. + + Args: + model_name: A Model Garden model resource name in the format + ``'publishers/{publisher}/models/{model}@{version}'``, a simplified name + in the format ``'{publisher}/{model}@{version}'`` (or without + ``@{version}``), or a Hugging Face model ID + ``'{organization}/{model}'``. + + Returns: + The resource name in the format + ``'publishers/{publisher}/models/{model}@{version}'``. + + Raises: + ValueError: If ``model_name`` is not a valid publisher model name. + """ + import re + + model_name = model_name.lower() # Hugging Face IDs are lower-case. + # A full resource name must carry an @version, matching the legacy SDK; a + # versionless full name is not accepted. + full_match = re.match( + r"^publishers/(?P[^/]+)/models/(?P[^@]+)@(?P[^@]+)$", + model_name, + ) + if full_match: + return ( + f"publishers/{full_match.group('publisher')}/models/" + f"{full_match.group('model')}@{full_match.group('version')}" + ) + # Reject Model Registry names; they would otherwise match the simplified + # branch and be silently mangled. + if re.match(r"^projects/.+/locations/.+/models/.+$", model_name): + raise ValueError(f"`{model_name}` is not a valid publisher model name") + simplified_match = re.match( + r"^(?P[^/]+)/(?P[^@]+)(?:@(?P.+))?$", + model_name, + ) + if simplified_match: + model = simplified_match.group("model") + if simplified_match.group("version"): + model = f"{model}@{simplified_match.group('version')}" + return f"publishers/{simplified_match.group('publisher')}/models/{model}" + raise ValueError(f"`{model_name}` is not a valid publisher model name") + + @staticmethod + def _is_hugging_face_model(model_name: str) -> bool: + """Returns whether a model name looks like a Hugging Face model ID. + + Matches the bare ``'{organization}/{model}'`` shape (a single slash and no + ``@version``), e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``. + + Args: + model_name: The model name to inspect. + + Returns: + True if ``model_name`` matches the Hugging Face ID shape. + """ + import re + + return bool( + re.match(r"^(?P[^/]+)/(?P[^/@]+)$", model_name) + ) + + @staticmethod + def _matches_filter( + value: Optional[str], + model_filter: Optional[Union[str, list[str]]], + ) -> bool: + """Returns whether ``value`` matches the (optional) keyword filter. + + Mirrors the legacy SDK: the filter may be a single keyword or a list of + keywords, and matching is a case-insensitive substring test where the value + matches if it contains *any* of the keywords. + + Args: + value: The field value to test (e.g. a machine type), or None. + model_filter: A keyword, a list of keywords, or None (no filtering). + + Returns: + True if there is no filter, or if ``value`` contains any of the keywords. + """ + if not model_filter: + return True + if value is None: + return False + keywords = [model_filter] if isinstance(model_filter, str) else model_filter + value_lower = value.lower() + return any(keyword.lower() in value_lower for keyword in keywords) + + @staticmethod + def _extract_and_filter_deploy_options( + publisher_model: types.PublisherModel, + machine_type_filter: Optional[Union[str, list[str]]] = None, + accelerator_type_filter: Optional[Union[str, list[str]]] = None, + serving_container_image_uri_filter: Optional[ + Union[str, list[str]] + ] = None, + ) -> list[types.DeployOption]: + """Extracts and filters deploy options from a publisher model. + + Args: + publisher_model: The publisher model to extract deploy options from. + machine_type_filter: Optional case-insensitive keyword (or list of + keywords) matched against the machine type; an option is kept if its + machine type contains any of them (e.g. ``'g2'`` or ``['n1', 'g2']``). + accelerator_type_filter: Optional case-insensitive keyword (or list of + keywords) matched against the accelerator type (e.g. ``'L4'`` or + ``['T4', 'L4']``). + serving_container_image_uri_filter: Optional case-insensitive keyword (or + list of keywords) matched against the serving container image URI (e.g. + ``'vllm'`` or ``['vllm', 'tgi']``). + + Returns: + A list of ``DeployOption`` objects matching the provided filters. + + Raises: + ValueError: If the model does not support deployment, or if no deploy + options remain after applying the filters. + """ + if not ( + publisher_model.supported_actions + and publisher_model.supported_actions.multi_deploy_vertex + and publisher_model.supported_actions.multi_deploy_vertex.multi_deploy_vertex + ): + raise ValueError( + "Model does not support deployment. " + "Use `list_deployable_models()` to find supported models." + ) + + options = ( + publisher_model.supported_actions.multi_deploy_vertex.multi_deploy_vertex + ) + result = [] + for opt in options: + container = opt.container_spec.image_uri if opt.container_spec else None + machine = ( + opt.dedicated_resources.machine_spec + if opt.dedicated_resources + else None + ) + machine_type = machine.machine_type if machine else None + + # Restore the proto3 defaults the JSON transport drops, so structured + # output matches the gRPC SDK on CPU/TPU machines. + accelerator_enum = machine.accelerator_type if machine else None + accelerator_value = accelerator_enum.value if accelerator_enum else None + has_accelerator = ( + accelerator_value is not None + and accelerator_value != "ACCELERATOR_TYPE_UNSPECIFIED" + ) + if machine: + accelerator_type = ( + accelerator_value + if accelerator_value is not None + else "ACCELERATOR_TYPE_UNSPECIFIED" ) - - options = self._extract_and_filter_deploy_options( - publisher_model, - machine_type_filter=config.machine_type_filter, - accelerator_type_filter=config.accelerator_type_filter, - serving_container_image_uri_filter=config.serving_container_image_uri_filter, + accelerator_count = ( + machine.accelerator_count + if machine.accelerator_count is not None + else 0 ) - - if config.concise is True: - return self._format_concise_deploy_options(options) - - return options - - @staticmethod - def _extract_recommend_spec(spec) -> dict[str, Any]: - """Extracts machine spec fields from a single recommend-spec entry. - - Args: - spec: A ``RecommendSpecResponseMachineAndModelContainerSpec`` describing a - recommended machine and container configuration. - - Returns: - A dict with the ``machine_type``, ``accelerator_type`` and - ``accelerator_count`` of the spec (values may be ``None``). - """ - machine_spec = spec.machine_spec - machine_type = None + else: accelerator_type = None accelerator_count = None - if machine_spec: - machine_type = getattr(machine_spec, "machine_type", None) - accelerator_enum = getattr(machine_spec, "accelerator_type", None) - if accelerator_enum: - accelerator_type = getattr(accelerator_enum, "name", None) - accelerator_count = getattr(machine_spec, "accelerator_count", None) - - return { - "machine_type": machine_type, - "accelerator_type": accelerator_type, - "accelerator_count": accelerator_count, - } - - @staticmethod - def _extract_recommendation(recommendation) -> dict[str, Any]: - """Extracts the spec, region and user quota state from a recommendation. - - Args: - recommendation: A ``RecommendSpecResponseRecommendation`` returned when - machine availability is requested. - - Returns: - A dict with the machine spec fields plus ``region`` and, when known, the - ``user_quota_state``. - """ - extracted_spec = ModelGarden._extract_recommend_spec(recommendation.spec) - extracted_spec["region"] = getattr(recommendation, "region", None) - if ( - recommendation.user_quota_state - and recommendation.user_quota_state - != types.QuotaState.QUOTA_STATE_UNSPECIFIED - ): - extracted_spec["user_quota_state"] = recommendation.user_quota_state.name - return extracted_spec - - @staticmethod - def _format_custom_deploy_options(options: list[dict[str, Any]]) -> str: - """Formats custom model deploy options into a human-readable string. - - Mirrors the legacy ``vertexai.model_garden`` ``CustomModel`` SDK output: - each option is rendered as an ``[Option N]`` block followed by its non-null - fields; ``accelerator_count`` is rendered unquoted. - - Args: - options: The extracted deploy option dicts to format. - - Returns: - A human-readable, multi-line string describing the deploy options. - """ - return "\n\n".join( - f"[Option {i + 1}]\n" - + ",\n".join( - f' {k}="{v}"' if k != "accelerator_count" else f" {k}={v}" - for k, v in option.items() - if v is not None - ) - for i, option in enumerate(options) - ) - - def list_custom_model_deploy_options( - self, - src: str, - config: Optional[types.ListCustomModelDeployOptionsConfigOrDict] = None, - ) -> str: - """Lists the recommended deploy options for a Model Garden custom model. - - Args: - src: The Google Cloud Storage URI of the custom model, storing the model - weights and config files (e.g. ``'gs://my-bucket/weights/'``). - config: Optional configuration. Accepts a - ``ListCustomModelDeployOptionsConfig`` instance or an equivalent dict. - - Returns: - A human-readable string describing the recommended deploy options - (machine type, accelerator type/count, region and, when available, the - user quota state). - - Raises: - ValueError: If ``src`` is not specified, or if no deploy options are - returned by the API (either because the backend produced none or - because the ``filter_by_user_quota`` filter dropped them all). - """ - if not src: - raise ValueError("src must be specified.") - if config is None: - config = types.ListCustomModelDeployOptionsConfig() - if isinstance(config, dict): - config = types.ListCustomModelDeployOptionsConfig.model_validate(config) - - parent = ( - f"projects/{self._api_client.project}/locations/" - f"{self._api_client.location}" - ) - - api_config = types.RecommendSpecConfig( - check_machine_availability=config.check_machine_availability, - check_user_quota=config.filter_by_user_quota, + if not ModelGarden._matches_filter(machine_type, machine_type_filter): + continue + # ACCELERATOR_TYPE_UNSPECIFIED means "no accelerator" and never matches. + if accelerator_type_filter and not has_accelerator: + continue + if not ModelGarden._matches_filter( + accelerator_type, accelerator_type_filter + ): + continue + if not ModelGarden._matches_filter( + container, serving_container_image_uri_filter + ): + continue + + result.append( + types.DeployOption( + option_name=opt.deploy_task_name, + serving_container_image_uri=container, + machine_type=machine_type, + accelerator_type=accelerator_type, + accelerator_count=accelerator_count, + ) + ) + + if not result: + raise ValueError("No deploy options found.") + + return result + + @staticmethod + def _format_concise_deploy_options( + options: list[types.DeployOption], + ) -> str: + """Formats deploy options into a human-readable string. + + Mirrors the legacy ``vertexai.model_garden`` SDK output: each option is + rendered as a ``[Option N: ]`` block followed by its non-null + fields (container image, machine type, accelerator type/count). + + Args: + options: The deploy options to format. + + Returns: + A human-readable, multi-line string describing the deploy options. + """ + fields = [ + "serving_container_image_uri", + "machine_type", + "accelerator_type", + "accelerator_count", + ] + blocks = [] + for i, option in enumerate(options): + if option.option_name: + header = f"[Option {i + 1}: {option.option_name}]\n" + else: + header = f"[Option {i + 1}]\n" + lines = [] + for field in fields: + value = getattr(option, field) + if value is None: + continue + if field == "accelerator_count": + lines.append(f" {field}={value},") + else: + lines.append(f' {field}="{value}",') + blocks.append(header + "\n".join(lines)) + return "\n\n".join(blocks) + + def _list_all_publisher_models( + self, + api_config: types.ListPublisherModelsConfig, + ) -> list[types.PublisherModel]: + """Fetches all pages of publisher models from the API. + + Args: + api_config: The configuration for the ListPublisherModels API call, + including filter and version settings. + + Returns: + A list of all ``PublisherModel`` objects across all pages. + """ + all_models = [] + page_token = None + while True: + if page_token: + api_config = types.ListPublisherModelsConfig( + filter=api_config.filter, + list_all_versions=api_config.list_all_versions, + page_token=page_token, ) - - response = self._recommend_spec( - parent=parent, - gcs_uri=src, - config=api_config, + response = self._list_publisher_models( + parent="publishers/*", + config=api_config, + ) + all_models.extend(response.publisher_models or []) + page_token = response.next_page_token + if not page_token: + break + return all_models + + def _list( + self, + model_filter: Optional[str], + include_hugging_face_models: Optional[bool], + deployable_only: bool, + ) -> list[str]: + """Shared implementation for listing models. + + Args: + model_filter: Optional substring to filter models by. + include_hugging_face_models: Whether to include HuggingFace models. + deployable_only: If True, only return models with deployment configs. + + Returns: + A list of formatted model name strings. + """ + include_hf = include_hugging_face_models is True + + filter_str = self._build_filter_str( + model_filter, include_hf, deployable_only=deployable_only + ) + + api_config = types.ListPublisherModelsConfig( + filter=filter_str, + list_all_versions=True, + ) + + models = self._list_all_publisher_models(api_config) + + if deployable_only: + # The VERIFIED_DEPLOYMENT_SUCCEED server filter only applies to HF + # models; filter native models client-side via multi_deploy_vertex. + models = [m for m in models if self._has_deploy_config(m)] + + return [self._format_model_name(m, include_hf) for m in models] + + def list_deployable_models( + self, + config: Optional[types.ListDeployableModelsConfigOrDict] = None, + ) -> list[str]: + """Lists models in Model Garden that support deployment. + + Returns models that have at least one verified deployment configuration. + When ``include_hugging_face_models`` is False (the default), + HuggingFace models are excluded from the results. + + Args: + config: Optional configuration for filtering results. Accepts a + ``ListDeployableModelsConfig`` instance or an equivalent dict. + + Returns: + A list of model name strings in the format + ``'{publisher}/{model}@{version}'`` (e.g. + ``'google/gemma2@gemma-2-2b-it'``) + or ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True + (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). + """ + if config is None: + config = types.ListDeployableModelsConfig() + if isinstance(config, dict): + config = types.ListDeployableModelsConfig.model_validate(config) + + return self._list( + config.model_filter, + config.include_hugging_face_models, + deployable_only=True, + ) + + def list_models( + self, + config: Optional[types.ListModelGardenModelsConfigOrDict] = None, + ) -> list[str]: + """Lists all models available in Model Garden. + + Returns all models regardless of deployment support. When + ``include_hugging_face_models`` is False (the default), HuggingFace + models are excluded from the results. + + Args: + config: Optional configuration for filtering results. Accepts a + ``ListModelGardenModelsConfig`` instance or an equivalent dict. + + Returns: + A list of model name strings in the format + ``'{publisher}/{model}@{version}'`` (e.g. + ``'google/gemma2@gemma-2-2b-it'``) + or ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True + (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). + """ + if config is None: + config = types.ListModelGardenModelsConfig() + if isinstance(config, dict): + config = types.ListModelGardenModelsConfig.model_validate(config) + + return self._list( + config.model_filter, + config.include_hugging_face_models, + deployable_only=False, + ) + + def list_publisher_model_deploy_options( + self, + model: str, + config: Optional[ + types.ListPublisherModelDeployOptionsConfigOrDict + ] = None, + ) -> Union[str, list[types.DeployOption]]: + """Lists the verified deploy options for a Model Garden publisher model. + + Supports Google open models (e.g. ``'google/gemma3@gemma-3-12b-it'``), + partner publisher models (e.g. + ``'mistralai/mistral-7b@mistral-7b-instruct-v0.2'``), and Hugging Face + model IDs (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). + + Args: + model: The publisher model to list deploy options for. Accepts the full + resource name ``'publishers/{publisher}/models/{model}@{version}'``, a + simplified ``'{publisher}/{model}@{version}'`` (or without the + ``@{version}``), or a Hugging Face model ID + ``'{organization}/{model}'``. + config: Optional configuration for filtering the deploy options. Accepts a + ``ListPublisherModelDeployOptionsConfig`` instance or an equivalent + dict. + + Returns: + A list of ``DeployOption`` objects, one per verified deployment + configuration (container image URI, machine type, accelerator type and + count). If ``config.concise`` is True, returns a human-readable string + describing those deploy options instead. + + Raises: + ValueError: If ``model`` is not a valid publisher model name, if the + model does not support deployment, or if no deploy options match the + provided filters. + """ + if config is None: + config = types.ListPublisherModelDeployOptionsConfig() + if isinstance(config, dict): + config = types.ListPublisherModelDeployOptionsConfig.model_validate( + config + ) + + get_publisher_model_config = types.GetPublisherModelConfig( + is_hugging_face_model=self._is_hugging_face_model(model), + include_equivalent_model_garden_model_deployment_configs=True, + ) + publisher_model = self._get_publisher_model( + name=self._reconcile_model_name(model), + config=get_publisher_model_config, + ) + + options = self._extract_and_filter_deploy_options( + publisher_model, + machine_type_filter=config.machine_type_filter, + accelerator_type_filter=config.accelerator_type_filter, + serving_container_image_uri_filter=config.serving_container_image_uri_filter, + ) + + if config.concise is True: + return self._format_concise_deploy_options(options) + + return options + + @staticmethod + def _extract_recommend_spec(spec) -> dict[str, Any]: + """Extracts machine spec fields from a single recommend-spec entry. + + Args: + spec: A ``RecommendSpecResponseMachineAndModelContainerSpec`` describing a + recommended machine and container configuration. + + Returns: + A dict with the ``machine_type``, ``accelerator_type`` and + ``accelerator_count`` of the spec (values may be ``None``). + """ + machine_spec = spec.machine_spec + machine_type = None + accelerator_type = None + accelerator_count = None + + if machine_spec: + machine_type = getattr(machine_spec, "machine_type", None) + accelerator_enum = getattr(machine_spec, "accelerator_type", None) + if accelerator_enum: + accelerator_type = getattr(accelerator_enum, "name", None) + accelerator_count = getattr(machine_spec, "accelerator_count", None) + + return { + "machine_type": machine_type, + "accelerator_type": accelerator_type, + "accelerator_count": accelerator_count, + } + + @staticmethod + def _extract_recommendation(recommendation) -> dict[str, Any]: + """Extracts the spec, region and user quota state from a recommendation. + + Args: + recommendation: A ``RecommendSpecResponseRecommendation`` returned when + machine availability is requested. + + Returns: + A dict with the machine spec fields plus ``region`` and, when known, the + ``user_quota_state``. + """ + extracted_spec = ModelGarden._extract_recommend_spec(recommendation.spec) + extracted_spec["region"] = getattr(recommendation, "region", None) + if ( + recommendation.user_quota_state + and recommendation.user_quota_state + != types.QuotaState.QUOTA_STATE_UNSPECIFIED + ): + extracted_spec["user_quota_state"] = recommendation.user_quota_state.name + return extracted_spec + + @staticmethod + def _format_custom_deploy_options(options: list[dict[str, Any]]) -> str: + """Formats custom model deploy options into a human-readable string. + + Mirrors the legacy ``vertexai.model_garden`` ``CustomModel`` SDK output: + each option is rendered as an ``[Option N]`` block followed by its non-null + fields; ``accelerator_count`` is rendered unquoted. + + Args: + options: The extracted deploy option dicts to format. + + Returns: + A human-readable, multi-line string describing the deploy options. + """ + return "\n\n".join( + f"[Option {i + 1}]\n" + + ",\n".join( + f' {k}="{v}"' if k != "accelerator_count" else f" {k}={v}" + for k, v in option.items() + if v is not None ) - - options = [] - if response.recommendations: - options = [ - self._extract_recommendation(recommendation) - for recommendation in response.recommendations - if recommendation.spec - ] - if config.filter_by_user_quota: - options = [ - option - for option in options - if option.get("user_quota_state") - == types.QuotaState.QUOTA_STATE_USER_HAS_QUOTA.name - ] - elif response.specs: - options = [ - self._extract_recommend_spec(spec) for spec in response.specs if spec - ] - - if not options: - raise ValueError("No deploy options found.") - - return self._format_custom_deploy_options(options) + for i, option in enumerate(options) + ) + + def list_custom_model_deploy_options( + self, + src: str, + config: Optional[types.ListCustomModelDeployOptionsConfigOrDict] = None, + ) -> str: + """Lists the recommended deploy options for a Model Garden custom model. + + Args: + src: The Google Cloud Storage URI of the custom model, storing the model + weights and config files (e.g. ``'gs://my-bucket/weights/'``). + config: Optional configuration. Accepts a + ``ListCustomModelDeployOptionsConfig`` instance or an equivalent dict. + + Returns: + A human-readable string describing the recommended deploy options + (machine type, accelerator type/count, region and, when available, the + user quota state). + + Raises: + ValueError: If ``src`` is not specified, or if no deploy options are + returned by the API (either because the backend produced none or + because the ``filter_by_user_quota`` filter dropped them all). + """ + if not src: + raise ValueError("src must be specified.") + if config is None: + config = types.ListCustomModelDeployOptionsConfig() + if isinstance(config, dict): + config = types.ListCustomModelDeployOptionsConfig.model_validate(config) + + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + + api_config = types.RecommendSpecConfig( + check_machine_availability=config.check_machine_availability, + check_user_quota=config.filter_by_user_quota, + ) + + response = self._recommend_spec( + parent=parent, + gcs_uri=src, + config=api_config, + ) + + options = [] + if response.recommendations: + options = [ + self._extract_recommendation(recommendation) + for recommendation in response.recommendations + if recommendation.spec + ] + if config.filter_by_user_quota: + options = [ + option + for option in options + if option.get("user_quota_state") + == types.QuotaState.QUOTA_STATE_USER_HAS_QUOTA.name + ] + elif response.specs: + options = [ + self._extract_recommend_spec(spec) for spec in response.specs if spec + ] + + if not options: + raise ValueError("No deploy options found.") + + return self._format_custom_deploy_options(options) + + def export_open_model( + self, + *, + model: str, + output_gcs_uri: str, + config: Optional[types.ExportOpenModelConfigOrDict] = None, + ) -> Union[str, types.ExportModelOperation]: + """Exports Google open model weights to a Cloud Storage bucket. + + Args: + model: The publisher model to export. Accepts a full resource name + ``'publishers/{publisher}/models/{model}@{version}'`` or the simplified + form ``'{publisher}/{model}@{version}'``. Hugging Face model IDs are not + supported. + output_gcs_uri: Cloud Storage URI prefix to write the model weights to + (e.g. ``'gs://my-bucket/gemma-weights/'``). + config: Optional export configuration (blocking behavior, poll timing). + See ``ExportOpenModelConfig``. + + Returns: + When ``config.wait_for_completion`` is ``True`` (default), the + Cloud Storage URI (``str``) where the weights were written. + When ``config.wait_for_completion`` is ``False``, the + ``ExportModelOperation`` for the caller to poll (via + ``get_export_publisher_model_operation`` or their own strategy). + + Raises: + ValueError: If ``output_gcs_uri`` is empty, if ``model`` is not a + valid Google publisher model resource name, or if ``model`` looks + like a Hugging Face model ID. + TimeoutError: If ``wait_for_completion=True`` and the LRO does not + complete within ``config.timeout_seconds`` (default 2 hours). + RuntimeError: If the LRO completes with an error, or completes + successfully but without a ``destination_uri``. + """ + if not output_gcs_uri: + raise ValueError("output_gcs_uri must be a non-empty Cloud Storage URI.") + if ModelGarden._is_hugging_face_model(model): + raise ValueError( + "export_open_model does not support Hugging Face model IDs " + f"(got {model!r}); only Google publisher open models are " + "exportable via this API." + ) + if config is None: + config = types.ExportOpenModelConfig() + elif isinstance(config, dict): + config = types.ExportOpenModelConfig.model_validate(config) + + publisher_model_name = ModelGarden._reconcile_model_name(model) + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + operation = self._export_publisher_model( + parent=parent, + name=publisher_model_name, + config=types.ExportPublisherModelConfig( + destination=genai_types.GcsDestination( + output_uri_prefix=output_gcs_uri, + ), + ), + ) + if not config.wait_for_completion: + return operation + + operation = _operations_utils.await_operation( + operation_name=operation.name, + get_operation_fn=self.get_export_publisher_model_operation, + poll_interval=( + config.poll_interval_seconds + or ModelGarden._DEFAULT_EXPORT_POLL_INTERVAL_SECONDS + ), + timeout_seconds=( + config.timeout_seconds + or ModelGarden._DEFAULT_EXPORT_TIMEOUT_SECONDS + ), + ) + if operation.error: + raise RuntimeError(f"Export failed: {operation.error}") + if not operation.response or not operation.response.destination_uri: + raise RuntimeError( + f"Export completed but response has no destination_uri: {operation!r}" + ) + return operation.response.destination_uri class AsyncModelGarden(_api_module.BaseModule): - """Model Garden module.""" - - async def _list_publisher_models( - self, - *, - parent: Optional[str] = None, - config: Optional[types.ListPublisherModelsConfigOrDict] = None, - ) -> types.ListPublisherModelsResponse: - """ - Lists publisher models (internal). - """ - - parameter_model = types._ListPublisherModelsRequestParameters( - parent=parent, - 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 = _ListPublisherModelsRequestParameters_to_vertex( - parameter_model - ) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{parent}/models".format_map(request_url_dict) - else: - path = "{parent}/models" - - 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) - - return_value = types.ListPublisherModelsResponse._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 - ), - } + """Model Garden module.""" + + async def _list_publisher_models( + self, + *, + parent: Optional[str] = None, + config: Optional[types.ListPublisherModelsConfigOrDict] = None, + ) -> types.ListPublisherModelsResponse: + """Lists publisher models (internal).""" + + parameter_model = types._ListPublisherModelsRequestParameters( + parent=parent, + 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 = _ListPublisherModelsRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/models".format_map(request_url_dict) + else: + path = "{parent}/models" + + 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) + + return_value = types.ListPublisherModelsResponse._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_publisher_model( - self, *, name: str, config: Optional[types.GetPublisherModelConfigOrDict] = None - ) -> types.PublisherModel: - """ - Gets a publisher model (internal). - """ - - parameter_model = types._GetPublisherModelRequestParameters( - 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 = _GetPublisherModelRequestParameters_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) - - return_value = types.PublisherModel._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_publisher_model( + self, + *, + name: str, + config: Optional[types.GetPublisherModelConfigOrDict] = None, + ) -> types.PublisherModel: + """Gets a publisher model (internal).""" + + parameter_model = types._GetPublisherModelRequestParameters( + 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 = _GetPublisherModelRequestParameters_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) + + return_value = types.PublisherModel._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 _recommend_spec( - self, - *, - parent: str, - gcs_uri: str, - config: Optional[types.RecommendSpecConfigOrDict] = None, - ) -> types.RecommendSpecResponse: - """ - Recommends spec for a custom model (internal). - """ - - parameter_model = types._RecommendSpecRequestParameters( - parent=parent, - gcs_uri=gcs_uri, - config=config, - ) - - request_url_dict: Optional[dict[str, str]] - if not self._api_client.vertexai: - raise ValueError( + } + if getattr(parameter_model, "config", None) + else {} + ), + ) + + self._api_client._verify_response(return_value) + return return_value + + async def _recommend_spec( + self, + *, + parent: str, + gcs_uri: str, + config: Optional[types.RecommendSpecConfigOrDict] = None, + ) -> types.RecommendSpecResponse: + """Recommends spec for a custom model (internal).""" + + parameter_model = types._RecommendSpecRequestParameters( + parent=parent, + gcs_uri=gcs_uri, + 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 = _RecommendSpecRequestParameters_to_vertex(parameter_model) - request_url_dict = request_dict.get("_url") - if request_url_dict: - path = "{parent}:recommendSpec".format_map(request_url_dict) - else: - path = "{parent}:recommendSpec" - - 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 ( + else: + request_dict = _RecommendSpecRequestParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}:recommendSpec".format_map(request_url_dict) + else: + path = "{parent}:recommendSpec" + + 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 + http_options = parameter_model.config.http_options - request_dict = _common.convert_to_dict(request_dict) - request_dict = _common.encode_unserializable_types(request_dict) + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) - response = await self._api_client.async_request( + response = await self._api_client.async_request( "post", path, request_dict, http_options ) - response_dict = {} if not response.body else json.loads(response.body) + response_dict = {} if not response.body else json.loads(response.body) - return_value = types.RecommendSpecResponse._from_response( + return_value = types.RecommendSpecResponse._from_response( response=response_dict, kwargs=( { @@ -1253,14 +1556,168 @@ async def _recommend_spec( ), ) - self._api_client._verify_response(return_value) - return return_value + self._api_client._verify_response(return_value) + return return_value + + async def _export_publisher_model( + self, + *, + parent: str, + name: str, + config: Optional[types.ExportPublisherModelConfigOrDict] = None, + ) -> types.ExportModelOperation: + """Exports a publisher model (internal).""" + + parameter_model = types._ExportPublisherModelRequestParameters( + parent=parent, + 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 = _ExportPublisherModelRequestParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{parent}/{name}:export".format_map(request_url_dict) + else: + path = "{parent}/{name}:export" + + 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) + + return_value = types.ExportModelOperation._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_export_publisher_model_operation( + self, + *, + operation_name: str, + config: Optional[ + types.GetExportPublisherModelOperationConfigOrDict + ] = None, + ) -> types.ExportModelOperation: + """Fetches the status of an in-flight ``export_open_model`` LRO.""" + + parameter_model = types._GetExportPublisherModelOperationParameters( + operation_name=operation_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 = _GetExportPublisherModelOperationParameters_to_vertex( + parameter_model + ) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = "{operationName}".format_map(request_url_dict) + else: + path = "{operationName}" + + 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) + + return_value = types.ExportModelOperation._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, - ) -> list[types.PublisherModel]: - """Fetches all pages of publisher models from the API. + async def _list_all_publisher_models( + self, + api_config: types.ListPublisherModelsConfig, + ) -> list[types.PublisherModel]: + """Fetches all pages of publisher models from the API. Args: api_config: The configuration for the ListPublisherModels API call, @@ -1269,32 +1726,32 @@ async def _list_all_publisher_models( Returns: A list of all ``PublisherModel`` objects across all pages. """ - all_models = [] - page_token = None - while True: - if page_token: - api_config = types.ListPublisherModelsConfig( + all_models = [] + page_token = None + while True: + if page_token: + api_config = types.ListPublisherModelsConfig( filter=api_config.filter, list_all_versions=api_config.list_all_versions, page_token=page_token, ) - response = await self._list_publisher_models( + response = await self._list_publisher_models( parent="publishers/*", config=api_config, ) - all_models.extend(response.publisher_models or []) - page_token = response.next_page_token - if not page_token: - break - return all_models + all_models.extend(response.publisher_models or []) + page_token = response.next_page_token + if not page_token: + break + return all_models - async def _list( + async def _list( self, model_filter: Optional[str], include_hugging_face_models: Optional[bool], deployable_only: bool, ) -> list[str]: - """Shared implementation for listing models. + """Shared implementation for listing models. Args: model_filter: Optional substring to filter models by. @@ -1304,29 +1761,29 @@ async def _list( Returns: A list of formatted model name strings. """ - include_hf = include_hugging_face_models is True + include_hf = include_hugging_face_models is True - filter_str = ModelGarden._build_filter_str( + filter_str = ModelGarden._build_filter_str( model_filter, include_hf, deployable_only=deployable_only ) - api_config = types.ListPublisherModelsConfig( + api_config = types.ListPublisherModelsConfig( filter=filter_str, list_all_versions=True, ) - models = await self._list_all_publisher_models(api_config) + models = await self._list_all_publisher_models(api_config) - if deployable_only: - models = [m for m in models if ModelGarden._has_deploy_config(m)] + if deployable_only: + models = [m for m in models if ModelGarden._has_deploy_config(m)] - return [ModelGarden._format_model_name(m, include_hf) for m in models] + return [ModelGarden._format_model_name(m, include_hf) for m in models] - async def list_deployable_models( + async def list_deployable_models( self, config: Optional[types.ListDeployableModelsConfigOrDict] = None, ) -> list[str]: - """Lists models in Model Garden that support deployment. + """Lists models in Model Garden that support deployment. Returns models that have at least one verified deployment configuration. When ``include_hugging_face_models`` is False (the default), @@ -1342,22 +1799,22 @@ async def list_deployable_models( or ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). """ - if config is None: - config = types.ListDeployableModelsConfig() - if isinstance(config, dict): - config = types.ListDeployableModelsConfig.model_validate(config) + if config is None: + config = types.ListDeployableModelsConfig() + if isinstance(config, dict): + config = types.ListDeployableModelsConfig.model_validate(config) - return await self._list( + return await self._list( config.model_filter, config.include_hugging_face_models, deployable_only=True, ) - async def list_models( + async def list_models( self, config: Optional[types.ListModelGardenModelsConfigOrDict] = None, ) -> list[str]: - """Lists all models available in Model Garden. + """Lists all models available in Model Garden. Returns all models regardless of deployment support. When ``include_hugging_face_models`` is False (the default), HuggingFace @@ -1373,23 +1830,23 @@ async def list_models( or ``'{publisher}/{model}'`` when ``include_hugging_face_models`` is True (e.g. ``'meta-llama/Llama-3.3-70B-Instruct'``). """ - if config is None: - config = types.ListModelGardenModelsConfig() - if isinstance(config, dict): - config = types.ListModelGardenModelsConfig.model_validate(config) + if config is None: + config = types.ListModelGardenModelsConfig() + if isinstance(config, dict): + config = types.ListModelGardenModelsConfig.model_validate(config) - return await self._list( + return await self._list( config.model_filter, config.include_hugging_face_models, deployable_only=False, ) - async def list_publisher_model_deploy_options( + async def list_publisher_model_deploy_options( self, model: str, config: Optional[types.ListPublisherModelDeployOptionsConfigOrDict] = None, ) -> Union[str, list[types.DeployOption]]: - """Lists the verified deploy options for a Model Garden publisher model. + """Lists the verified deploy options for a Model Garden publisher model. Supports Google open models (e.g. ``'google/gemma3@gemma-3-12b-it'``), partner publisher models (e.g. @@ -1416,37 +1873,37 @@ async def list_publisher_model_deploy_options( model does not support deployment, or if no deploy options match the provided filters. """ - if config is None: - config = types.ListPublisherModelDeployOptionsConfig() - if isinstance(config, dict): - config = types.ListPublisherModelDeployOptionsConfig.model_validate(config) + if config is None: + config = types.ListPublisherModelDeployOptionsConfig() + if isinstance(config, dict): + config = types.ListPublisherModelDeployOptionsConfig.model_validate(config) - api_config = types.GetPublisherModelConfig( + api_config = types.GetPublisherModelConfig( is_hugging_face_model=ModelGarden._is_hugging_face_model(model), include_equivalent_model_garden_model_deployment_configs=True, ) - publisher_model = await self._get_publisher_model( + publisher_model = await self._get_publisher_model( name=ModelGarden._reconcile_model_name(model), config=api_config ) - options = ModelGarden._extract_and_filter_deploy_options( + options = ModelGarden._extract_and_filter_deploy_options( publisher_model, machine_type_filter=config.machine_type_filter, accelerator_type_filter=config.accelerator_type_filter, serving_container_image_uri_filter=config.serving_container_image_uri_filter, ) - if config.concise is True: - return ModelGarden._format_concise_deploy_options(options) + if config.concise is True: + return ModelGarden._format_concise_deploy_options(options) - return options + return options - async def list_custom_model_deploy_options( + async def list_custom_model_deploy_options( self, src: str, config: Optional[types.ListCustomModelDeployOptionsConfigOrDict] = None, ) -> str: - """Lists the recommended deploy options for a Model Garden custom model. + """Lists the recommended deploy options for a Model Garden custom model. Args: src: The Google Cloud Storage URI of the custom model, storing the model @@ -1464,51 +1921,109 @@ async def list_custom_model_deploy_options( returned by the API (either because the backend produced none or because the ``filter_by_user_quota`` filter dropped them all). """ - if not src: - raise ValueError("src must be specified.") - if config is None: - config = types.ListCustomModelDeployOptionsConfig() - if isinstance(config, dict): - config = types.ListCustomModelDeployOptionsConfig.model_validate(config) - - parent = ( + if not src: + raise ValueError("src must be specified.") + if config is None: + config = types.ListCustomModelDeployOptionsConfig() + if isinstance(config, dict): + config = types.ListCustomModelDeployOptionsConfig.model_validate(config) + + parent = ( f"projects/{self._api_client.project}/locations/" f"{self._api_client.location}" ) - api_config = types.RecommendSpecConfig( + api_config = types.RecommendSpecConfig( check_machine_availability=config.check_machine_availability, check_user_quota=config.filter_by_user_quota, ) - response = await self._recommend_spec( + response = await self._recommend_spec( parent=parent, gcs_uri=src, config=api_config, ) - options = [] - if response.recommendations: - options = [ + options = [] + if response.recommendations: + options = [ ModelGarden._extract_recommendation(recommendation) for recommendation in response.recommendations if recommendation.spec ] - if config.filter_by_user_quota: - options = [ + if config.filter_by_user_quota: + options = [ option for option in options if option.get("user_quota_state") == types.QuotaState.QUOTA_STATE_USER_HAS_QUOTA.name ] - elif response.specs: - options = [ + elif response.specs: + options = [ ModelGarden._extract_recommend_spec(spec) for spec in response.specs if spec ] - if not options: - raise ValueError("No deploy options found.") - - return ModelGarden._format_custom_deploy_options(options) + if not options: + raise ValueError("No deploy options found.") + + return ModelGarden._format_custom_deploy_options(options) + + async def export_open_model( + self, + *, + model: str, + output_gcs_uri: str, + config: Optional[types.ExportOpenModelConfigOrDict] = None, + ) -> Union[str, types.ExportModelOperation]: + """Async variant of ``ModelGarden.export_open_model``.""" + if not output_gcs_uri: + raise ValueError("output_gcs_uri must be a non-empty Cloud Storage URI.") + if ModelGarden._is_hugging_face_model(model): + raise ValueError( + "export_open_model does not support Hugging Face model IDs " + f"(got {model!r}); only Google publisher open models are " + "exportable via this API." + ) + if config is None: + config = types.ExportOpenModelConfig() + elif isinstance(config, dict): + config = types.ExportOpenModelConfig.model_validate(config) + + publisher_model_name = ModelGarden._reconcile_model_name(model) + parent = ( + f"projects/{self._api_client.project}/locations/" + f"{self._api_client.location}" + ) + operation = await self._export_publisher_model( + parent=parent, + name=publisher_model_name, + config=types.ExportPublisherModelConfig( + destination=genai_types.GcsDestination( + output_uri_prefix=output_gcs_uri, + ), + ), + ) + if not config.wait_for_completion: + return operation + + operation = await _operations_utils.await_operation_async( + operation_name=operation.name, + get_operation_fn=self.get_export_publisher_model_operation, + poll_interval=( + config.poll_interval_seconds + or ModelGarden._DEFAULT_EXPORT_POLL_INTERVAL_SECONDS + ), + timeout_seconds=( + config.timeout_seconds + or ModelGarden._DEFAULT_EXPORT_TIMEOUT_SECONDS + ), + ) + if operation.error: + raise RuntimeError(f"Export failed: {operation.error}") + if not operation.response or not operation.response.destination_uri: + raise RuntimeError( + f"Export completed but response has no destination_uri: {operation!r}" + ) + return operation.response.destination_uri diff --git a/agentplatform/_genai/types/__init__.py b/agentplatform/_genai/types/__init__.py index a731e74a0c..adf2d397f6 100644 --- a/agentplatform/_genai/types/__init__.py +++ b/agentplatform/_genai/types/__init__.py @@ -66,6 +66,7 @@ from .common import _DeleteSkillRequestParameters from .common import _EvaluateInstancesRequestParameters from .common import _ExecuteCodeAgentEngineSandboxRequestParameters +from .common import _ExportPublisherModelRequestParameters from .common import _GenerateAgentEngineMemoriesRequestParameters from .common import _GenerateInstanceRubricsRequest from .common import _GenerateLossClustersParameters @@ -94,6 +95,7 @@ from .common import _GetEvaluationMetricParameters from .common import _GetEvaluationRunParameters from .common import _GetEvaluationSetParameters +from .common import _GetExportPublisherModelOperationParameters from .common import _GetImportFilesOperationParameters from .common import _GetMultimodalDatasetOperationParameters from .common import _GetMultimodalDatasetParameters @@ -616,6 +618,18 @@ from .common import ExecuteSandboxEnvironmentResponse from .common import ExecuteSandboxEnvironmentResponseDict from .common import ExecuteSandboxEnvironmentResponseOrDict +from .common import ExportModelOperation +from .common import ExportModelOperationDict +from .common import ExportModelOperationOrDict +from .common import ExportOpenModelConfig +from .common import ExportOpenModelConfigDict +from .common import ExportOpenModelConfigOrDict +from .common import ExportPublisherModelConfig +from .common import ExportPublisherModelConfigDict +from .common import ExportPublisherModelConfigOrDict +from .common import ExportPublisherModelResponse +from .common import ExportPublisherModelResponseDict +from .common import ExportPublisherModelResponseOrDict from .common import FailedRubric from .common import FailedRubricDict from .common import FailedRubricOrDict @@ -730,6 +744,9 @@ from .common import GetEvaluationSetConfig from .common import GetEvaluationSetConfigDict from .common import GetEvaluationSetConfigOrDict +from .common import GetExportPublisherModelOperationConfig +from .common import GetExportPublisherModelOperationConfigDict +from .common import GetExportPublisherModelOperationConfigOrDict from .common import GetImportFilesOperationConfig from .common import GetImportFilesOperationConfigDict from .common import GetImportFilesOperationConfigOrDict @@ -3469,6 +3486,18 @@ "RecommendSpecResponse", "RecommendSpecResponseDict", "RecommendSpecResponseOrDict", + "ExportPublisherModelConfig", + "ExportPublisherModelConfigDict", + "ExportPublisherModelConfigOrDict", + "ExportPublisherModelResponse", + "ExportPublisherModelResponseDict", + "ExportPublisherModelResponseOrDict", + "ExportModelOperation", + "ExportModelOperationDict", + "ExportModelOperationOrDict", + "GetExportPublisherModelOperationConfig", + "GetExportPublisherModelOperationConfigDict", + "GetExportPublisherModelOperationConfigOrDict", "CreateRuntimeFeedbackEntryConfig", "CreateRuntimeFeedbackEntryConfigDict", "CreateRuntimeFeedbackEntryConfigOrDict", @@ -3619,6 +3648,9 @@ "ListCustomModelDeployOptionsConfig", "ListCustomModelDeployOptionsConfigDict", "ListCustomModelDeployOptionsConfigOrDict", + "ExportOpenModelConfig", + "ExportOpenModelConfigDict", + "ExportOpenModelConfigOrDict", "DeployOption", "DeployOptionDict", "DeployOptionOrDict", @@ -3810,6 +3842,8 @@ "_ListPublisherModelsRequestParameters", "_GetPublisherModelRequestParameters", "_RecommendSpecRequestParameters", + "_ExportPublisherModelRequestParameters", + "_GetExportPublisherModelOperationParameters", "_CreateRuntimeFeedbackEntryRequestParameters", "_DeleteRuntimeFeedbackEntryRequestParameters", "_GetRuntimeFeedbackRequestParameters", diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index dd57c62b1e..34ad4d6643 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -23755,6 +23755,198 @@ class RecommendSpecResponseDict(TypedDict, total=False): RecommendSpecResponseOrDict = Union[RecommendSpecResponse, RecommendSpecResponseDict] +class ExportPublisherModelConfig(_common.BaseModel): + """RPC-level config for ``export_publisher_model``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + destination: Optional[genai_types.GcsDestination] = Field( + default=None, description="""""" + ) + + +class ExportPublisherModelConfigDict(TypedDict, total=False): + """RPC-level config for ``export_publisher_model``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + destination: Optional[genai_types.GcsDestination] + """""" + + +ExportPublisherModelConfigOrDict = Union[ + ExportPublisherModelConfig, ExportPublisherModelConfigDict +] + + +class _ExportPublisherModelRequestParameters(_common.BaseModel): + """Parameters for ``export_publisher_model``.""" + + parent: Optional[str] = Field(default=None, description="""""") + name: Optional[str] = Field(default=None, description="""""") + config: Optional[ExportPublisherModelConfig] = Field( + default=None, description="""""" + ) + + +class _ExportPublisherModelRequestParametersDict(TypedDict, total=False): + """Parameters for ``export_publisher_model``.""" + + parent: Optional[str] + """""" + + name: Optional[str] + """""" + + config: Optional[ExportPublisherModelConfigDict] + """""" + + +_ExportPublisherModelRequestParametersOrDict = Union[ + _ExportPublisherModelRequestParameters, + _ExportPublisherModelRequestParametersDict, +] + + +class ExportPublisherModelResponse(_common.BaseModel): + """Response for the ``ExportPublisherModel`` RPC. + + Fields are re-declared as ``SdkFieldPatch`` (both are already in the + discovery-generated class) so the SDK's dependency on ``destination_uri`` + is visible in one place and proto drift is caught at codegen time + instead of at first user call. + """ + + destination_uri: Optional[str] = Field( + default=None, + description="""Cloud Storage URI where the exported weights were written.""", + ) + publisher_model: Optional[str] = Field( + default=None, + description="""Resource name of the publisher model that was exported.""", + ) + + +class ExportPublisherModelResponseDict(TypedDict, total=False): + """Response for the ``ExportPublisherModel`` RPC. + + Fields are re-declared as ``SdkFieldPatch`` (both are already in the + discovery-generated class) so the SDK's dependency on ``destination_uri`` + is visible in one place and proto drift is caught at codegen time + instead of at first user call. + """ + + destination_uri: Optional[str] + """Cloud Storage URI where the exported weights were written.""" + + publisher_model: Optional[str] + """Resource name of the publisher model that was exported.""" + + +ExportPublisherModelResponseOrDict = Union[ + ExportPublisherModelResponse, ExportPublisherModelResponseDict +] + + +class ExportModelOperation(_common.BaseModel): + """Long-running operation returned by ``ExportPublisherModel``.""" + + name: Optional[str] = Field( + default=None, + description="""The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""", + ) + metadata: Optional[dict[str, Any]] = Field( + default=None, + description="""Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""", + ) + done: Optional[bool] = Field( + default=None, + description="""If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""", + ) + error: Optional[dict[str, Any]] = Field( + default=None, + description="""The error result of the operation in case of failure or cancellation.""", + ) + response: Optional[ExportPublisherModelResponse] = Field( + default=None, description="""""" + ) + + +class ExportModelOperationDict(TypedDict, total=False): + """Long-running operation returned by ``ExportPublisherModel``.""" + + name: Optional[str] + """The server-assigned name, which is only unique within the same service that originally returns it. If you use the default HTTP mapping, the `name` should be a resource name ending with `operations/{unique_id}`.""" + + metadata: Optional[dict[str, Any]] + """Service-specific metadata associated with the operation. It typically contains progress information and common metadata such as create time. Some services might not provide such metadata. Any method that returns a long-running operation should document the metadata type, if any.""" + + done: Optional[bool] + """If the value is `false`, it means the operation is still in progress. If `true`, the operation is completed, and either `error` or `response` is available.""" + + error: Optional[dict[str, Any]] + """The error result of the operation in case of failure or cancellation.""" + + response: Optional[ExportPublisherModelResponseDict] + """""" + + +ExportModelOperationOrDict = Union[ + ExportModelOperation, ExportModelOperationDict +] + + +class GetExportPublisherModelOperationConfig(_common.BaseModel): + """Config for ``get_export_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GetExportPublisherModelOperationConfigDict(TypedDict, total=False): + """Config for ``get_export_publisher_model_operation``.""" + + http_options: Optional[genai_types.HttpOptions] + """Used to override HTTP request options.""" + + +GetExportPublisherModelOperationConfigOrDict = Union[ + GetExportPublisherModelOperationConfig, + GetExportPublisherModelOperationConfigDict, +] + + +class _GetExportPublisherModelOperationParameters(_common.BaseModel): + """Parameters for polling an ``export_publisher_model`` operation.""" + + operation_name: Optional[str] = Field( + default=None, + description="""The server-assigned name for the operation.""", + ) + config: Optional[GetExportPublisherModelOperationConfig] = Field( + default=None, description="""""" + ) + + +class _GetExportPublisherModelOperationParametersDict(TypedDict, total=False): + """Parameters for polling an ``export_publisher_model`` operation.""" + + operation_name: Optional[str] + """The server-assigned name for the operation.""" + + config: Optional[GetExportPublisherModelOperationConfigDict] + """""" + + +_GetExportPublisherModelOperationParametersOrDict = Union[ + _GetExportPublisherModelOperationParameters, + _GetExportPublisherModelOperationParametersDict, +] + + class CreateRuntimeFeedbackEntryConfig(_common.BaseModel): """Config for creating a Feedback Entry.""" @@ -26797,6 +26989,55 @@ class ListCustomModelDeployOptionsConfigDict(TypedDict, total=False): ] +class ExportOpenModelConfig(_common.BaseModel): + """Config for ``export_open_model``.""" + + wait_for_completion: Optional[bool] = Field( + default=True, + description="""Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""", + ) + poll_interval_seconds: Optional[float] = Field( + default=None, + description="""Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""", + ) + timeout_seconds: Optional[float] = Field( + default=None, + description="""Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""", + ) + + +class ExportOpenModelConfigDict(TypedDict, total=False): + """Config for ``export_open_model``.""" + + wait_for_completion: Optional[bool] + """Whether to block on the export long-running operation. When + ``True`` (default), returns the destination URI on completion. When + ``False``, returns the ``ExportModelOperation`` for the caller to + poll.""" + + poll_interval_seconds: Optional[float] + """Seconds between LRO polls when ``wait_for_completion=True``. + Defaults to 30. Ignored when ``wait_for_completion=False``.""" + + timeout_seconds: Optional[float] + """Total wall-clock seconds to wait for the export to complete + when ``wait_for_completion=True``. Defaults to 2 hours to + accommodate large model weights. Ignored when + ``wait_for_completion=False``.""" + + +ExportOpenModelConfigOrDict = Union[ + ExportOpenModelConfig, ExportOpenModelConfigDict +] + + class DeployOption(_common.BaseModel): """A verified deploy option for a model.""" diff --git a/noxfile.py b/noxfile.py index e6c97e2d5f..b43a081d26 100644 --- a/noxfile.py +++ b/noxfile.py @@ -105,11 +105,11 @@ "unit_ray", "unit_langchain", "unit_ag2", - "unit_llama_index", + # "unit_llama_index", "unit_agentplatform_adk", "unit_agentplatform_langchain", "unit_agentplatform_ag2", - "unit_agentplatform_llama_index", + # "unit_agentplatform_llama_index", "unit_agentplatform_a2a", "unit_a2a", "system", 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 828976c517..0702fe1fdf 100644 --- a/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/replays/test_genai_model_garden.py @@ -36,11 +36,11 @@ def test_list_deployable_models(client): def test_list_models(client): """Tests listing all baseline models in Model Garden.""" models = client.model_garden.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False, - model_filter="timesfm", - ) - ) + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=False, + model_filter="timesfm", + ) + ) assert len(models) > 0 assert isinstance(models[0], str) assert "timesfm" in models[0].lower() @@ -83,9 +83,9 @@ def test_list_publisher_model_deploy_options_concise(client): def test_list_publisher_model_deploy_options_hugging_face(client): """Tests deploy options for a Hugging Face model. - Exercises the distinct GetPublisherModel request path where - is_hugging_face_model=True is sent. - """ + Exercises the distinct GetPublisherModel request path where + is_hugging_face_model=True is sent. + """ options = client.model_garden.list_publisher_model_deploy_options( model="codellama/codellama-7b-hf" ) @@ -143,7 +143,9 @@ def test_list_custom_model_deploy_options_no_user_quota_filter(client): """filter_by_user_quota=False -> recommendations returned without quota filtering.""" options = client.model_garden.list_custom_model_deploy_options( src=_CUSTOM_MODEL_SRC, - config=types.ListCustomModelDeployOptionsConfig(filter_by_user_quota=False), + config=types.ListCustomModelDeployOptionsConfig( + filter_by_user_quota=False + ), ) assert isinstance(options, str) assert "[Option 1]" in options @@ -161,6 +163,95 @@ def test_list_custom_model_deploy_options_dict_config(client): assert "region=" not in options +def test_export_open_model_wait_for_completion(client): + """Default wait_for_completion=True path: blocks on the LRO, returns URI. + + End-to-end coverage of the polling loop -- records the initial POST plus + every GET on the operation until it's done, so recording takes as long as + the export itself (~2h for Gemma-2-2B). Once recorded, replaying is + instant. Skip this filter when iterating on unrelated changes. + """ + destination_uri = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it/", + ) + assert isinstance(destination_uri, str) + assert destination_uri.startswith("gs://") + + +def test_export_open_model_no_wait_returns_operation(client): + """wait_for_completion=False returns the ExportModelOperation immediately. + + Records only the initial ExportPublisherModel POST (no LRO polling), so + this is cheap to record (seconds). + """ + operation = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it-no-wait/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + assert isinstance(operation, types.ExportModelOperation) + assert operation.name + assert "/operations/" in operation.name + + +def test_export_open_model_no_wait_dict_config(client): + """The config may be passed as a plain dict; wait_for_completion=False path.""" + operation = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-dict-no-wait/" + ), + config={"wait_for_completion": False}, + ) + assert isinstance(operation, types.ExportModelOperation) + assert operation.name + + +def test_get_export_publisher_model_operation(client): + """Public getter returns an ExportModelOperation for a running export LRO. + + Chains a wait_for_completion=False export with an explicit poll via the + public getter -- exercises the same call users doing their own polling + (custom backoff, cancellation, etc.) will make. + """ + op = client.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + output_gcs_uri="gs://mg-sdk-model-export-testing/gemma-2-2b-it-poll/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + polled = client.model_garden.get_export_publisher_model_operation( + operation_name=op.name, + ) + assert isinstance(polled, types.ExportModelOperation) + assert polled.name == op.name + + +def test_export_open_model_empty_uri_raises(client): + """Empty output_gcs_uri is rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="output_gcs_uri must be a non-empty"): + client.model_garden.export_open_model( + model="google/gemma3@gemma-3-12b-it", output_gcs_uri="" + ) + + +def test_export_open_model_hf_model_id_raises(client): + """Hugging Face model IDs are rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="does not support Hugging Face"): + client.model_garden.export_open_model( + model="meta-llama/Llama-3.3-70B-Instruct", + output_gcs_uri="gs://mg-sdk-model-export-testing/hf-should-not-record/", + ) + + +def test_export_open_model_invalid_name_raises(client): + """Invalid model name is rejected client-side; no RPC needed.""" + with pytest.raises(ValueError, match="not a valid publisher model name"): + client.model_garden.export_open_model( + model="not-a-valid-name", output_gcs_uri="gs://b/out/" + ) + + pytestmark = pytest_helper.setup( file=__file__, globals_for_file=globals(), @@ -220,3 +311,34 @@ async def test_list_custom_model_deploy_options_async(client): assert isinstance(options, str) assert "[Option 1]" in options assert "region=" in options + + +@pytest.mark.asyncio +async def test_export_open_model_async_no_wait_returns_operation(client): + """Async wait_for_completion=False returns the ExportModelOperation immediately.""" + operation = await client.aio.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-async-no-wait/" + ), + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + assert isinstance(operation, types.ExportModelOperation) + assert operation.name + + +@pytest.mark.asyncio +async def test_get_export_publisher_model_operation_async(client): + """Public async getter returns an ExportModelOperation for an in-flight LRO.""" + op = await client.aio.model_garden.export_open_model( + model="google/gemma2@gemma-2-2b-it", + output_gcs_uri=( + "gs://mg-sdk-model-export-testing/gemma-2-2b-it-async-poll/" + ), + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + polled = await client.aio.model_garden.get_export_publisher_model_operation( + operation_name=op.name, + ) + assert isinstance(polled, types.ExportModelOperation) + assert polled.name == op.name diff --git a/tests/unit/agentplatform/genai/test_genai_model_garden.py b/tests/unit/agentplatform/genai/test_genai_model_garden.py index bd39dd7e82..f46752cebb 100644 --- a/tests/unit/agentplatform/genai/test_genai_model_garden.py +++ b/tests/unit/agentplatform/genai/test_genai_model_garden.py @@ -16,13 +16,16 @@ import asyncio from unittest import mock -from agentplatform._genai import model_garden +from agentplatform._genai import ( + model_garden, +) # pylint: disable=unused-import from agentplatform._genai import types from google.genai import client import pytest _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" +_OPEN_MODEL = "google/gemma3@gemma-3-12b-it" @pytest.fixture @@ -219,8 +222,8 @@ def test_list_deployable_models_hf_filter_string(mock_client): def test_list_models_excludes_hf_via_server_filter(mock_client): - """Tests that is_hf_wildcard(false) is sent to exclude HF models server-side.""" - dummy_response = types.ListPublisherModelsResponse( + """Tests that is_hf_wildcard(false) is sent to exclude HF models server-side.""" + dummy_response = types.ListPublisherModelsResponse( publisher_models=[ types.PublisherModel( name="publishers/google/models/gemma-2b", version_id="001" @@ -228,24 +231,24 @@ def test_list_models_excludes_hf_via_server_filter(mock_client): ] ) - with mock.patch.object( + with mock.patch.object( mock_client, "_list_publisher_models", return_value=dummy_response ) as mock_list: - models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False - ) + models = mock_client.list_models( + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=False ) + ) - api_config = mock_list.call_args.kwargs.get("config") - assert "is_hf_wildcard(false)" in api_config.filter - assert len(models) == 1 - assert models[0] == "google/gemma-2b@001" + api_config = mock_list.call_args.kwargs.get("config") + assert "is_hf_wildcard(false)" in api_config.filter + assert len(models) == 1 + assert models[0] == "google/gemma-2b@001" def test_list_models_with_hf(mock_client): - """Tests list_models with include_hugging_face_models=True.""" - dummy_response = types.ListPublisherModelsResponse( + """Tests list_models with include_hugging_face_models=True.""" + dummy_response = types.ListPublisherModelsResponse( publisher_models=[ types.PublisherModel( name="publishers/google/models/gemma-2b", version_id="001" @@ -256,27 +259,27 @@ def test_list_models_with_hf(mock_client): ] ) - with mock.patch.object( + with mock.patch.object( mock_client, "_list_publisher_models", return_value=dummy_response ): - models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=True - ) + models = mock_client.list_models( + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=True ) + ) - assert len(models) == 2 - assert models[0] == "google/gemma-2b" - assert models[1] == "meta/llama-3" + assert len(models) == 2 + assert models[0] == "google/gemma-2b" + assert models[1] == "meta/llama-3" def test_list_models_includes_non_deployable(mock_client): - """Tests that list_models includes models without deploy configs. + """Tests that list_models includes models without deploy configs. Unlike list_deployable_models, list_models should return ALL models regardless of whether they have multi_deploy_vertex configurations. """ - dummy_response = types.ListPublisherModelsResponse( + dummy_response = types.ListPublisherModelsResponse( publisher_models=[ _make_deployable_model("publishers/google/models/gemma-2b"), # This model has no deploy config -- should still be included @@ -286,18 +289,18 @@ def test_list_models_includes_non_deployable(mock_client): ] ) - with mock.patch.object( + with mock.patch.object( mock_client, "_list_publisher_models", return_value=dummy_response ): - models = mock_client.list_models( - config=types.ListModelGardenModelsConfig( - include_hugging_face_models=False - ) + models = mock_client.list_models( + config=types.ListModelGardenModelsConfig( + include_hugging_face_models=False ) + ) - assert len(models) == 2 - assert models[0] == "google/gemma-2b@001" - assert models[1] == "google/bert-base@001" + assert len(models) == 2 + assert models[0] == "google/gemma-2b@001" + assert models[1] == "google/bert-base@001" def test_list_models_default_config(mock_client): @@ -419,8 +422,10 @@ def test_build_filter_str_escapes_special_chars(): """Tests that special regex characters in model_filter are escaped.""" build_filter = model_garden.ModelGarden._build_filter_str result = build_filter( - model_filter="model.v2+", include_hugging_face_models=False, deployable_only=False - ) + model_filter="model.v2+", + include_hugging_face_models=False, + deployable_only=False, + ) # re.escape turns '.' into '\\.' and '+' into '\\+' assert r"model\.v2\+" in result @@ -435,8 +440,8 @@ def _make_deploy_option( accelerator_count=1, image_uri="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/vllm", ): - """Builds a single PublisherModelCallToActionDeploy option.""" - return types.PublisherModelCallToActionDeploy( + """Builds a single PublisherModelCallToActionDeploy option.""" + return types.PublisherModelCallToActionDeploy( deploy_task_name=deploy_task_name, dedicated_resources=types.DedicatedResources( machine_spec=types.MachineSpec( @@ -454,13 +459,13 @@ def _make_model_with_deploy_options( ): """Builds a PublisherModel exposing the given deploy options.""" return types.PublisherModel( - name=name, - supported_actions=types.PublisherModelCallToAction( - multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( - multi_deploy_vertex=options, - ) - ), - ) + name=name, + supported_actions=types.PublisherModelCallToAction( + multi_deploy_vertex=types.PublisherModelCallToActionDeployVertex( + multi_deploy_vertex=options, + ) + ), + ) def test_list_publisher_model_deploy_options_basic(mock_client): @@ -468,8 +473,8 @@ def test_list_publisher_model_deploy_options_basic(mock_client): dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ) as mock_get: + mock_client, "_get_publisher_model", return_value=dummy_model + ) as mock_get: # A simplified name is reconciled to the full publisher resource name. options = mock_client.list_publisher_model_deploy_options( model="google/gemma3@gemma-3-12b-it" @@ -497,12 +502,10 @@ def test_list_publisher_model_deploy_options_basic(mock_client): def test_list_publisher_model_deploy_options_multiple(mock_client): """Tests that all deploy options are returned when no filters are set.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), + _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -516,12 +519,10 @@ def test_list_publisher_model_deploy_options_multiple(mock_client): def test_list_publisher_model_deploy_options_machine_type_filter(mock_client): """Tests machine_type_filter is a case-insensitive substring match.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), + _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -564,12 +565,10 @@ def test_list_publisher_model_deploy_options_accelerator_type_filter( def test_list_publisher_model_deploy_options_image_uri_filter(mock_client): """Tests serving_container_image_uri_filter is a case-insensitive match.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), - _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), + _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -637,13 +636,11 @@ def test_list_publisher_model_deploy_options_accelerator_type_filter_list( def test_list_publisher_model_deploy_options_image_uri_filter_list(mock_client): """Tests a list image-uri filter matches any keyword.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), - _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), - _make_deploy_option(deploy_task_name="sglang", image_uri="docker/sglang"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option(deploy_task_name="vllm", image_uri="docker/vllm"), + _make_deploy_option(deploy_task_name="tgi", image_uri="docker/tgi"), + _make_deploy_option(deploy_task_name="sglang", image_uri="docker/sglang"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -678,12 +675,10 @@ def test_matches_filter(): def test_list_publisher_model_deploy_options_dict_config(mock_client): """Tests config passed as a dict is validated and applied.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), + _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -820,8 +815,8 @@ def test_list_publisher_model_deploy_options_hugging_face_model(mock_client): dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) with mock.patch.object( - mock_client, "_get_publisher_model", return_value=dummy_model - ) as mock_get: + mock_client, "_get_publisher_model", return_value=dummy_model + ) as mock_get: mock_client.list_publisher_model_deploy_options( model="meta-llama/Llama-3.3-70B-Instruct" ) @@ -840,10 +835,10 @@ def test_list_publisher_model_deploy_options_async(mock_async_client): dummy_model = _make_model_with_deploy_options([_make_deploy_option()]) with mock.patch.object( - mock_async_client, - "_get_publisher_model", - new=mock.AsyncMock(return_value=dummy_model), - ): + mock_async_client, + "_get_publisher_model", + new=mock.AsyncMock(return_value=dummy_model), + ): options = asyncio.run( mock_async_client.list_publisher_model_deploy_options( model="publishers/google/models/gemma-2b@001" @@ -862,10 +857,10 @@ def test_list_publisher_model_deploy_options_async_no_deploy_support_raises( dummy_model = types.PublisherModel(name="publishers/google/models/bert-base") with mock.patch.object( - mock_async_client, - "_get_publisher_model", - new=mock.AsyncMock(return_value=dummy_model), - ): + mock_async_client, + "_get_publisher_model", + new=mock.AsyncMock(return_value=dummy_model), + ): with pytest.raises(ValueError, match="does not support deployment"): asyncio.run( mock_async_client.list_publisher_model_deploy_options( @@ -881,12 +876,12 @@ def _make_deploy_option_no_accelerator( ): """Builds a deploy option whose machine has no GPU accelerator (e.g. TPU).""" return types.PublisherModelCallToActionDeploy( - deploy_task_name=deploy_task_name, - dedicated_resources=types.DedicatedResources( - machine_spec=types.MachineSpec(machine_type=machine_type) - ), - container_spec=types.ModelContainerSpec(image_uri=image_uri), - ) + deploy_task_name=deploy_task_name, + dedicated_resources=types.DedicatedResources( + machine_spec=types.MachineSpec(machine_type=machine_type) + ), + container_spec=types.ModelContainerSpec(image_uri=image_uri), + ) def test_list_publisher_model_deploy_options_no_accelerator_defaults( @@ -940,12 +935,10 @@ def test_list_publisher_model_deploy_options_accelerator_filter_excludes_unspeci mock_client, ): """Tests accelerator_type_filter excludes no-accelerator options (legacy parity).""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option_no_accelerator(deploy_task_name="tpu"), - _make_deploy_option(deploy_task_name="gpu", accelerator_type="NVIDIA_L4"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option_no_accelerator(deploy_task_name="tpu"), + _make_deploy_option(deploy_task_name="gpu", accelerator_type="NVIDIA_L4"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -1017,24 +1010,22 @@ def test_list_publisher_model_deploy_options_concise_returns_string( def test_list_publisher_model_deploy_options_concise_multiple(mock_client): """Tests concise formatting of multiple options separated by a blank line.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option( - deploy_task_name="g2", - machine_type="g2-standard-12", - accelerator_type="NVIDIA_L4", - accelerator_count=1, - image_uri="docker/vllm", - ), - _make_deploy_option( - deploy_task_name="a3", - machine_type="a3-highgpu-8g", - accelerator_type="NVIDIA_H100_80GB", - accelerator_count=8, - image_uri="docker/hexllm", - ), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option( + deploy_task_name="g2", + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + image_uri="docker/vllm", + ), + _make_deploy_option( + deploy_task_name="a3", + machine_type="a3-highgpu-8g", + accelerator_type="NVIDIA_H100_80GB", + accelerator_count=8, + image_uri="docker/hexllm", + ), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -1062,12 +1053,10 @@ def test_list_publisher_model_deploy_options_concise_multiple(mock_client): def test_list_publisher_model_deploy_options_concise_with_filter(mock_client): """Tests concise output honors filters before formatting.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), - _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option(deploy_task_name="g2", machine_type="g2-standard-12"), + _make_deploy_option(deploy_task_name="a3", machine_type="a3-highgpu-8g"), + ]) with mock.patch.object( mock_client, "_get_publisher_model", return_value=dummy_model @@ -1087,39 +1076,35 @@ def test_list_publisher_model_deploy_options_concise_with_filter(mock_client): def test_format_concise_deploy_options_omits_none_fields(): """Tests that None fields are omitted and the header has no name if unset.""" options = [ - types.DeployOption( - machine_type="g2-standard-12", - accelerator_count=1, - ) - ] + types.DeployOption( + machine_type="g2-standard-12", + accelerator_count=1, + ) + ] result = model_garden.ModelGarden._format_concise_deploy_options(options) expected = ( - "[Option 1]\n" - ' machine_type="g2-standard-12",\n' - " accelerator_count=1," - ) + '[Option 1]\n machine_type="g2-standard-12",\n accelerator_count=1,' + ) assert result == expected def test_list_publisher_model_deploy_options_concise_async(mock_async_client): """Tests the async client returns a concise string when requested.""" - dummy_model = _make_model_with_deploy_options( - [ - _make_deploy_option( - deploy_task_name="option-1", - machine_type="g2-standard-12", - accelerator_type="NVIDIA_L4", - accelerator_count=1, - image_uri="docker/vllm", - ) - ] - ) + dummy_model = _make_model_with_deploy_options([ + _make_deploy_option( + deploy_task_name="option-1", + machine_type="g2-standard-12", + accelerator_type="NVIDIA_L4", + accelerator_count=1, + image_uri="docker/vllm", + ) + ]) with mock.patch.object( - mock_async_client, - "_get_publisher_model", - new=mock.AsyncMock(return_value=dummy_model), - ): + mock_async_client, + "_get_publisher_model", + new=mock.AsyncMock(return_value=dummy_model), + ): result = asyncio.run( mock_async_client.list_publisher_model_deploy_options( model="publishers/google/models/gemma-2b@001", @@ -1399,3 +1384,389 @@ def test_list_custom_model_deploy_options_async_src_required_raises( """Tests the async client also validates src.""" with pytest.raises(ValueError, match="src must be specified"): asyncio.run(mock_async_client.list_custom_model_deploy_options(src="")) + + +# ===================================================================== +# export_open_model +# ===================================================================== + +_TEST_EXPORT_OPERATION = "projects/p/locations/us-central1/operations/exp-1" +_TEST_EXPORT_URI = "gs://my-bucket/gemma-weights/exported/" + + +def _make_export_operation( + name=_TEST_EXPORT_OPERATION, *, done=False, destination_uri=None, error=None +): + """Builds a fake ExportModelOperation for use in tests.""" + op = types.ExportModelOperation(name=name, done=done) + if destination_uri is not None: + op.response = types.ExportPublisherModelResponse( + destination_uri=destination_uri + ) + if error is not None: + op.error = error + return op + + +def _patch_export_rpc(mock_client, initial_op=None): + """Mocks ``_export_publisher_model`` to return ``initial_op``.""" + return mock.patch.object( + mock_client, + "_export_publisher_model", + return_value=( + initial_op if initial_op is not None else _make_export_operation() + ), + ) + + +def _patch_async_export_rpc(mock_client, initial_op=None): + """Async variant of _patch_export_rpc.""" + return mock.patch.object( + mock_client, + "_export_publisher_model", + new=mock.AsyncMock( + return_value=( + initial_op if initial_op is not None else _make_export_operation() + ) + ), + ) + + +# --- Return type: wait_for_completion=True (default) ------------------- + + +def test_export_open_model_returns_destination_uri_when_waiting(mock_client): + """Default (wait_for_completion=True) blocks on LRO and returns the URI.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + m_await.assert_called_once() + assert result == _TEST_EXPORT_URI + + +def test_export_open_model_polls_via_public_getter(mock_client): + """``get_export_publisher_model_operation`` is used as the polling callback.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + kwargs = m_await.call_args.kwargs + assert ( + kwargs["get_operation_fn"] + == mock_client.get_export_publisher_model_operation + ) + assert kwargs["operation_name"] == _TEST_EXPORT_OPERATION + + +def test_export_open_model_uses_default_poll_and_timeout(mock_client): + """No override in config -> ModelGarden defaults are forwarded to the poller.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + kwargs = m_await.call_args.kwargs + assert ( + kwargs["poll_interval"] == 30 + ) # _DEFAULT_EXPORT_POLL_INTERVAL_SECONDS + assert ( + kwargs["timeout_seconds"] == 2 * 60 * 60 + ) # _DEFAULT_EXPORT_TIMEOUT_SECONDS + + +def test_export_open_model_config_overrides_poll_and_timeout(mock_client): + """Config values override the defaults on the polling call.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ) as m_await, + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig( + poll_interval_seconds=5, + timeout_seconds=600, + ), + ) + kwargs = m_await.call_args.kwargs + assert kwargs["poll_interval"] == 5 + assert kwargs["timeout_seconds"] == 600 + + +def test_export_open_model_raises_on_operation_error(mock_client): + """A backend error on the LRO surfaces as RuntimeError.""" + failed_op = _make_export_operation( + done=True, error={"message": "EULA required"} + ) + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=failed_op, + ), + ): + with pytest.raises(RuntimeError, match="Export failed"): + mock_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + + +def test_export_open_model_raises_when_response_missing_destination_uri( + mock_client, +): + """A done LRO with no destination_uri surfaces as RuntimeError.""" + done_op = _make_export_operation(done=True) # no destination_uri set + with ( + _patch_export_rpc(mock_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation", + return_value=done_op, + ), + ): + with pytest.raises(RuntimeError, match="destination_uri"): + mock_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + + +# --- Return type: wait_for_completion=False ---------------------------- + + +def test_export_open_model_returns_operation_when_not_waiting(mock_client): + """wait_for_completion=False returns the operation immediately (no polling).""" + initial_op = _make_export_operation() + with ( + _patch_export_rpc(mock_client, initial_op=initial_op), + mock.patch.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + m_await.assert_not_called() + assert result is initial_op + + +# --- Input validation ------------------------------------------------- + + +def test_export_open_model_empty_output_gcs_uri_raises(mock_client): + """An empty output_gcs_uri raises before any RPC or polling.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="output_gcs_uri must be a non-empty"): + mock_client.export_open_model(model=_OPEN_MODEL, output_gcs_uri="") + m_rpc.assert_not_called() + + +def test_export_open_model_hf_model_id_raises(mock_client): + """Hugging Face model IDs are rejected client-side with a targeted message.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="does not support Hugging Face"): + mock_client.export_open_model( + model="meta-llama/Llama-3.3-70B-Instruct", + output_gcs_uri="gs://b/out/", + ) + m_rpc.assert_not_called() + + +def test_export_open_model_invalid_name_raises_before_rpc(mock_client): + """Malformed model names raise ValueError; no RPC and no polling.""" + with _patch_export_rpc(mock_client) as m_rpc: + with pytest.raises(ValueError, match="not a valid publisher model name"): + mock_client.export_open_model( + model="not-a-valid-name", output_gcs_uri="gs://b/out/" + ) + m_rpc.assert_not_called() + + +def test_export_open_model_dict_config_validated(mock_client): + """A plain dict is validated into ExportOpenModelConfig.""" + initial_op = _make_export_operation() + with ( + _patch_export_rpc(mock_client, initial_op=initial_op), + mock.patch.object( + model_garden._operations_utils, "await_operation" + ) as m_await, + ): + result = mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config={"wait_for_completion": False}, + ) + m_await.assert_not_called() + assert result is initial_op + + +# --- Model-name resolution --------------------------------------------- + + +def test_export_open_model_simplified_name_reconciled(mock_client): + """Simplified '{publisher}/{model}@{version}' -> full resource name.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch.object(model_garden._operations_utils, "await_operation"), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + kwargs = m_rpc.call_args.kwargs + assert kwargs["name"] == "publishers/google/models/gemma3@gemma-3-12b-it" + assert kwargs["parent"] == ( + f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}" + ) + + +def test_export_open_model_full_resource_name_preserved(mock_client): + """A full 'publishers/.../models/...@version' is passed through unchanged.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch.object(model_garden._operations_utils, "await_operation"), + ): + mock_client.export_open_model( + model="publishers/google/models/gemma3@gemma-3-12b-it", + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + assert ( + m_rpc.call_args.kwargs["name"] + == "publishers/google/models/gemma3@gemma-3-12b-it" + ) + + +# --- Destination forwarded to RPC config -------------------------------- + + +def test_export_open_model_forwards_output_gcs_uri(mock_client): + """Top-level output_gcs_uri becomes destination.output_uri_prefix on the RPC.""" + with ( + _patch_export_rpc(mock_client) as m_rpc, + mock.patch.object(model_garden._operations_utils, "await_operation"), + ): + mock_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://my-bucket/weights/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + assert ( + m_rpc.call_args.kwargs["config"].destination.output_uri_prefix + == "gs://my-bucket/weights/" + ) + + +# --- Async surface parity ---------------------------------------------- + + +def test_export_open_model_async_returns_destination_uri_when_waiting( + mock_async_client, +): + """Async default (wait=True) awaits the LRO and returns the URI string.""" + done_op = _make_export_operation(done=True, destination_uri=_TEST_EXPORT_URI) + with ( + _patch_async_export_rpc(mock_async_client), + mock.patch.object( + model_garden._operations_utils, + "await_operation_async", + new=mock.AsyncMock(return_value=done_op), + ), + ): + result = asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="gs://b/out/" + ) + ) + assert result == _TEST_EXPORT_URI + + +def test_export_open_model_async_returns_operation_when_not_waiting( + mock_async_client, +): + """Async wait=False returns the operation immediately.""" + initial_op = _make_export_operation() + with ( + _patch_async_export_rpc(mock_async_client, initial_op=initial_op), + mock.patch.object( + model_garden._operations_utils, "await_operation_async" + ) as m_await, + ): + result = asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, + output_gcs_uri="gs://b/out/", + config=types.ExportOpenModelConfig(wait_for_completion=False), + ) + ) + m_await.assert_not_called() + assert result is initial_op + + +def test_export_open_model_async_empty_uri_raises(mock_async_client): + """Async missing output_gcs_uri raises ValueError before any RPC or polling.""" + with _patch_async_export_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="output_gcs_uri must be a non-empty"): + asyncio.run( + mock_async_client.export_open_model( + model=_OPEN_MODEL, output_gcs_uri="" + ) + ) + m_rpc.assert_not_called() + + +def test_export_open_model_async_hf_model_id_raises(mock_async_client): + """Async Hugging Face model IDs are rejected client-side.""" + with _patch_async_export_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="does not support Hugging Face"): + asyncio.run( + mock_async_client.export_open_model( + model="meta-llama/Llama-3.3-70B-Instruct", + output_gcs_uri="gs://b/out/", + ) + ) + m_rpc.assert_not_called() + + +def test_export_open_model_async_invalid_name_raises(mock_async_client): + """Async invalid model name raises ValueError before any RPC or polling.""" + with _patch_async_export_rpc(mock_async_client) as m_rpc: + with pytest.raises(ValueError, match="not a valid publisher model name"): + asyncio.run( + mock_async_client.export_open_model( + model="not-a-valid-name", output_gcs_uri="gs://b/out/" + ) + ) + m_rpc.assert_not_called()