feat(model-engine): add GcpPubSubQueueEndpointResourceDelegate for async endpoints on GCP #827
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,133 @@ | ||
| from typing import Any, Dict, Optional | ||
|
|
||
| from google.api_core import exceptions as gcp_exceptions | ||
| from google.cloud import pubsub_v1 | ||
| from google.protobuf import field_mask_pb2 | ||
| from model_engine_server.core.loggers import logger_name, make_logger | ||
| from model_engine_server.domain.exceptions import EndpointResourceInfraException | ||
| from model_engine_server.infra.gateways.resources.queue_endpoint_resource_delegate import ( | ||
| QueueEndpointResourceDelegate, | ||
| QueueInfo, | ||
| ) | ||
|
|
||
| logger = make_logger(logger_name()) | ||
|
|
||
| GCP_PUBSUB_MAX_ACK_DEADLINE_SECONDS = 600 # Pub/Sub hard limit | ||
|
|
||
|
|
||
| class GcpPubSubQueueEndpointResourceDelegate(QueueEndpointResourceDelegate): | ||
| """ | ||
| Using GCP Pub/Sub (topic + subscription per endpoint). | ||
|
|
||
| topic_prefix and subscription_prefix control the GCP resource name prefix. | ||
| The logical queue_name returned to callers always uses the canonical | ||
| QueueEndpointResourceDelegate.endpoint_id_to_queue_name format, independent | ||
| of these prefixes. | ||
| """ | ||
|
|
||
| def __init__( | ||
| self, | ||
| project_id: str, | ||
| topic_prefix: str = "launch-endpoint-id-", | ||
| subscription_prefix: str = "launch-endpoint-id-", | ||
| ) -> None: | ||
| if not project_id: | ||
| raise ValueError( | ||
| "GcpPubSubQueueEndpointResourceDelegate requires a non-empty project_id; " | ||
| "set infra.gcp_project_id in the service config." | ||
| ) | ||
| self.project_id = project_id | ||
| self.topic_prefix = topic_prefix | ||
| self.subscription_prefix = subscription_prefix | ||
| self._publisher = pubsub_v1.PublisherClient() | ||
| self._subscriber = pubsub_v1.SubscriberClient() | ||
|
|
||
| def _topic_id(self, endpoint_id: str) -> str: | ||
| return f"{self.topic_prefix}{endpoint_id}" | ||
|
|
||
| def _subscription_id(self, endpoint_id: str) -> str: | ||
| return f"{self.subscription_prefix}{endpoint_id}" | ||
|
|
||
| async def create_queue_if_not_exists( | ||
| self, | ||
| endpoint_id: str, | ||
| endpoint_name: str, | ||
| endpoint_created_by: str, | ||
| endpoint_labels: Dict[str, Any], | ||
| queue_message_timeout_seconds: Optional[int] = None, | ||
| ) -> QueueInfo: | ||
| queue_name = QueueEndpointResourceDelegate.endpoint_id_to_queue_name( | ||
| endpoint_id | ||
| ) | ||
| topic_path = f"projects/{self.project_id}/topics/{self._topic_id(endpoint_id)}" | ||
| subscription_path = f"projects/{self.project_id}/subscriptions/{self._subscription_id(endpoint_id)}" | ||
| ack_deadline = min( | ||
| queue_message_timeout_seconds or 60, GCP_PUBSUB_MAX_ACK_DEADLINE_SECONDS | ||
| ) | ||
|
|
||
| try: | ||
| self._publisher.create_topic(name=topic_path) | ||
| except gcp_exceptions.AlreadyExists: | ||
| pass | ||
|
|
||
| try: | ||
| self._subscriber.create_subscription( | ||
| name=subscription_path, | ||
| topic=topic_path, | ||
| ack_deadline_seconds=ack_deadline, | ||
| ) | ||
| except gcp_exceptions.AlreadyExists: | ||
| try: | ||
| self._subscriber.update_subscription( | ||
| subscription=pubsub_v1.types.Subscription( | ||
| name=subscription_path, | ||
| ack_deadline_seconds=ack_deadline, | ||
| ), | ||
| update_mask=field_mask_pb2.FieldMask( | ||
| paths=["ack_deadline_seconds"] | ||
| ), | ||
| ) | ||
| except gcp_exceptions.GoogleAPIError as e: | ||
| logger.warning( | ||
| f"Failed to update ack_deadline for Pub/Sub subscription {subscription_path}: {e}" | ||
| ) | ||
|
|
||
| # Pub/Sub has no URL concept analogous to SQS queue URLs | ||
| return QueueInfo(queue_name, queue_url=None) | ||
|
|
||
| async def delete_queue(self, endpoint_id: str) -> None: | ||
| subscription_path = f"projects/{self.project_id}/subscriptions/{self._subscription_id(endpoint_id)}" | ||
| topic_path = f"projects/{self.project_id}/topics/{self._topic_id(endpoint_id)}" | ||
|
|
||
| try: | ||
| self._subscriber.delete_subscription(subscription=subscription_path) | ||
| except gcp_exceptions.NotFound: | ||
| logger.info( | ||
| f"Could not find Pub/Sub subscription {subscription_path} for endpoint {endpoint_id}" | ||
| ) | ||
| except gcp_exceptions.GoogleAPIError as e: | ||
| raise EndpointResourceInfraException( | ||
| f"Failed to delete Pub/Sub subscription {subscription_path} for endpoint {endpoint_id}: {e}" | ||
| ) from e | ||
|
|
||
| try: | ||
| self._publisher.delete_topic(topic=topic_path) | ||
| except gcp_exceptions.NotFound: | ||
| logger.info( | ||
| f"Could not find Pub/Sub topic {topic_path} for endpoint {endpoint_id}" | ||
| ) | ||
| except gcp_exceptions.GoogleAPIError as e: | ||
| raise EndpointResourceInfraException( | ||
| f"Failed to delete Pub/Sub topic {topic_path} for endpoint {endpoint_id}: {e}" | ||
| ) from e | ||
|
|
||
| async def get_queue_attributes(self, endpoint_id: str) -> Dict[str, Any]: | ||
| queue_name = QueueEndpointResourceDelegate.endpoint_id_to_queue_name( | ||
| endpoint_id | ||
| ) | ||
| return { | ||
| "name": queue_name, | ||
| # Pub/Sub does not expose a synchronous undelivered message count; | ||
| # real observability requires the Cloud Monitoring API as a separate concern. | ||
| "num_undelivered_messages": -1, | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
GCP_PROJECT_IDenv var is dead config — delegate will always raiseValueErroron GCP clustersThe Helm chart injects
GCP_PROJECT_IDfrom.Values.gcp.project_id, but the delegate is instantiated withproject_id=infra_config().gcp_project_id, which is loaded from theinfra_service_configYAML ConfigMap. That ConfigMap is rendered from.Values.config.values.infrainservice_config_map.yaml— not from.Values.gcp. Since the newgcp.project_idkey is in a separate Helm section andgcp_project_idis never injected intoconfig.values.infra,infra_config().gcp_project_idwill always beNoneat runtime. The delegate's ownif not project_id: raise ValueError(...)guard will fire on every startup on any GCP cluster that follows the sample values. The same broken path is repeated ink8s_cache.pyandstart_batch_job_orchestration.py. Either reados.getenv("GCP_PROJECT_ID")directly in the delegate (consistent with howASBQueueEndpointResourceDelegatereadsos.getenv("SERVICEBUS_NAMESPACE")), or addgcp_project_idtoconfig.values.infrain the chart.Prompt To Fix With AI