|
| 1 | +import json |
| 2 | +from collections import defaultdict |
| 3 | +from contextlib import contextmanager |
| 4 | +from functools import wraps |
| 5 | +from typing import Generator, Optional, Union, Callable, Iterable |
| 6 | +from weakref import WeakKeyDictionary |
| 7 | + |
| 8 | +import google.auth |
| 9 | +import vertexai |
| 10 | +from dagster import ( |
| 11 | + AssetExecutionContext, |
| 12 | + AssetKey, |
| 13 | + ConfigurableResource, |
| 14 | + DagsterInvariantViolationError, |
| 15 | + InitResourceContext, |
| 16 | + OpExecutionContext, |
| 17 | +) |
| 18 | +from dagster._annotations import public |
| 19 | +from google.auth.credentials import Credentials as GoogleCredentials |
| 20 | +from pydantic import Field, PrivateAttr |
| 21 | +from vertexai.generative_models import GenerationResponse, GenerativeModel |
| 22 | + |
| 23 | +from dagster._annotations import preview |
| 24 | + |
| 25 | + |
| 26 | +# A WeakKeyDictionary is used to track usage counters for each execution context. |
| 27 | +# This ensures that counters are automatically garbage-collected when the context is no longer in use, |
| 28 | +# preventing memory leaks in long-running Dagster processes. |
| 29 | +context_to_counters = WeakKeyDictionary() |
| 30 | + |
| 31 | + |
| 32 | +def _add_to_asset_metadata( |
| 33 | + context: AssetExecutionContext, usage_metadata: dict, output_name: Optional[str] |
| 34 | +) -> None: |
| 35 | + """Adds and aggregates usage metadata to the current asset materialization.""" |
| 36 | + if context not in context_to_counters: |
| 37 | + context_to_counters[context] = defaultdict(lambda: 0) |
| 38 | + counters = context_to_counters[context] |
| 39 | + |
| 40 | + for metadata_key, delta in usage_metadata.items(): |
| 41 | + counters[metadata_key] += delta |
| 42 | + context.add_output_metadata(dict(counters), output_name) |
| 43 | + |
| 44 | + |
| 45 | +def _with_usage_metadata( |
| 46 | + context: AssetExecutionContext, output_name: Optional[str], func: Callable |
| 47 | +) -> Callable: |
| 48 | + """A wrapper for the `generate_content` method to handle both streaming and non-streaming |
| 49 | + responses, extracting and logging usage metadata correctly for each case. |
| 50 | + """ |
| 51 | + |
| 52 | + @wraps(func) |
| 53 | + def wrapper( |
| 54 | + *args, **kwargs |
| 55 | + ) -> Union[GenerationResponse, Iterable[GenerationResponse]]: |
| 56 | + is_streaming = kwargs.get("stream", False) |
| 57 | + |
| 58 | + if not is_streaming: |
| 59 | + # Standard non-streaming request |
| 60 | + response = func(*args, **kwargs) |
| 61 | + if hasattr(response, "usage_metadata"): |
| 62 | + usage = response.usage_metadata |
| 63 | + usage_metadata = { |
| 64 | + "vertexai.calls": 1, |
| 65 | + "vertexai.candidates_token_count": usage.candidates_token_count, |
| 66 | + "vertexai.prompt_token_count": usage.prompt_token_count, |
| 67 | + "vertexai.total_token_count": usage.total_token_count, # Fixed: Use correct attribute |
| 68 | + } |
| 69 | + _add_to_asset_metadata(context, usage_metadata, output_name) |
| 70 | + return response |
| 71 | + |
| 72 | + # Streaming request: must return a generator |
| 73 | + def streaming_generator() -> Iterable[GenerationResponse]: |
| 74 | + # The original function call returns an iterator |
| 75 | + response_iterator = func(*args, **kwargs) |
| 76 | + yield from response_iterator |
| 77 | + |
| 78 | + # After the stream is exhausted, the iterator itself has the usage_metadata |
| 79 | + if hasattr(response_iterator, "usage_metadata"): |
| 80 | + usage = response_iterator.usage_metadata |
| 81 | + usage_metadata = { |
| 82 | + "vertexai.calls": 1, |
| 83 | + "vertexai.candidates_token_count": usage.candidates_token_count, |
| 84 | + "vertexai.prompt_token_count": usage.prompt_token_count, |
| 85 | + "vertexai.total_token_count": usage.total_token_count, # Fixed: Use correct attribute |
| 86 | + } |
| 87 | + _add_to_asset_metadata(context, usage_metadata, output_name) |
| 88 | + |
| 89 | + return streaming_generator() |
| 90 | + |
| 91 | + return wrapper |
| 92 | + |
| 93 | + |
| 94 | +@public |
| 95 | +@preview |
| 96 | +class VertexAIResource(ConfigurableResource): |
| 97 | + """A Dagster resource for interacting with Google Cloud Vertex AI's generative models. |
| 98 | +
|
| 99 | + This resource provides a `vertexai.generative_models.GenerativeModel` instance, configured |
| 100 | + for use in ops and assets. It automatically handles the initialization of the Vertex AI SDK |
| 101 | + and logs API usage (token counts and call counts) as asset metadata. |
| 102 | +
|
| 103 | + **Authentication:** |
| 104 | + Authentication is handled via `Application Default Credentials (ADC) <https://cloud.google.com/docs/authentication/application-default-credentials>`_ |
| 105 | + by default. To use a service account, provide the service account key as a JSON string via |
| 106 | + the ``google_credentials`` configuration field. |
| 107 | +
|
| 108 | + **Dependencies:** |
| 109 | + This resource requires the ``google-cloud-aiplatform`` library to be installed. |
| 110 | + ``pip install google-cloud-aiplatform`` |
| 111 | +
|
| 112 | + Examples: |
| 113 | + .. code-block:: python |
| 114 | +
|
| 115 | + from dagster import AssetExecutionContext, Definitions, EnvVar, asset |
| 116 | + from dagster_vertexai import VertexAIResource |
| 117 | +
|
| 118 | + @asset(compute_kind="vertexai") |
| 119 | + def my_report(context: AssetExecutionContext, vertex_ai: VertexAIResource): |
| 120 | + with vertex_ai.get_model(context) as model: |
| 121 | + response = model.generate_content( |
| 122 | + "Generate a business summary for a new marketing campaign." |
| 123 | + ) |
| 124 | + return response.text |
| 125 | +
|
| 126 | + defs = Definitions( |
| 127 | + assets=[my_report], |
| 128 | + resources={ |
| 129 | + "vertex_ai": VertexAIResource( |
| 130 | + project_id="my-gcp-project", |
| 131 | + location="us-central1", |
| 132 | + generative_model_name="gemini-1.5-flash-001" |
| 133 | + ), |
| 134 | + }, |
| 135 | + ) |
| 136 | + """ |
| 137 | + |
| 138 | + project_id: str = Field(description="The Google Cloud project ID.") |
| 139 | + location: str = Field( |
| 140 | + description="The Google Cloud location (e.g., 'us-central1')." |
| 141 | + ) |
| 142 | + |
| 143 | + google_credentials: Optional[str] = Field( |
| 144 | + default=None, |
| 145 | + description=( |
| 146 | + "A JSON string of Google Cloud service account credentials. If not provided, " |
| 147 | + "Application Default Credentials (ADC) will be used." |
| 148 | + ), |
| 149 | + ) |
| 150 | + api_endpoint: Optional[str] = Field( |
| 151 | + default=None, |
| 152 | + description=( |
| 153 | + "The regional API endpoint for Vertex AI. If not set, the SDK will determine it " |
| 154 | + "based on the `location`. Example: 'us-central1-aiplatform.googleapis.com'" |
| 155 | + ), |
| 156 | + ) |
| 157 | + generative_model_name: str = Field( |
| 158 | + description="The name of the generative model on Vertex AI (e.g., 'gemini-1.5-flash-001')." |
| 159 | + ) |
| 160 | + |
| 161 | + _generative_model: GenerativeModel = PrivateAttr() |
| 162 | + |
| 163 | + @classmethod |
| 164 | + def _is_dagster_maintained(cls) -> bool: |
| 165 | + return False |
| 166 | + |
| 167 | + def setup_for_execution(self, context: InitResourceContext) -> None: |
| 168 | + """Initializes the Vertex AI SDK and the generative model instance.""" |
| 169 | + creds: Optional[GoogleCredentials] = None |
| 170 | + if self.google_credentials: |
| 171 | + try: |
| 172 | + creds_dict = json.loads(self.google_credentials) |
| 173 | + creds, _ = google.auth.load_credentials_from_dict(creds_dict) |
| 174 | + except (json.JSONDecodeError, TypeError) as e: |
| 175 | + raise DagsterInvariantViolationError( |
| 176 | + "Failed to load Google credentials from 'google_credentials'. " |
| 177 | + f"Please ensure it is a valid JSON string. Error: {e}" |
| 178 | + ) |
| 179 | + |
| 180 | + vertexai.init( |
| 181 | + project=self.project_id, |
| 182 | + location=self.location, |
| 183 | + credentials=creds, |
| 184 | + api_endpoint=self.api_endpoint, |
| 185 | + ) |
| 186 | + self._generative_model = GenerativeModel(model_name=self.generative_model_name) |
| 187 | + |
| 188 | + @public |
| 189 | + @contextmanager |
| 190 | + def get_model( |
| 191 | + self, context: Union[AssetExecutionContext, OpExecutionContext] |
| 192 | + ) -> Generator[GenerativeModel, None, None]: |
| 193 | + """Yields a ``vertexai.generative_models.GenerativeModel`` for interacting with Vertex AI. |
| 194 | +
|
| 195 | + When used within an asset's compute function, this method wraps the ``model.generate_content`` |
| 196 | + method to automatically log its usage to the asset's materialization metadata. |
| 197 | +
|
| 198 | + Args: |
| 199 | + context (Union[AssetExecutionContext, OpExecutionContext]): The Dagster execution context. |
| 200 | + """ |
| 201 | + with self._get_model(context=context, asset_key=None) as model: |
| 202 | + yield model |
| 203 | + |
| 204 | + def _wrap_for_usage_tracking( |
| 205 | + self, context: AssetExecutionContext, output_name: Optional[str] |
| 206 | + ): |
| 207 | + """Patches the `generate_content` method on the model instance for the current context.""" |
| 208 | + original_generate_content = self._generative_model.generate_content |
| 209 | + self._generative_model.generate_content = _with_usage_metadata( |
| 210 | + context, output_name, func=original_generate_content |
| 211 | + ) |
| 212 | + |
| 213 | + @public |
| 214 | + @contextmanager |
| 215 | + def get_model_for_asset( |
| 216 | + self, context: AssetExecutionContext, asset_key: AssetKey |
| 217 | + ) -> Generator[GenerativeModel, None, None]: |
| 218 | + """Yields a ``vertexai.generative_models.GenerativeModel`` for a specific asset. |
| 219 | +
|
| 220 | + This method is for use in ``@multi_asset``s. It ensures that any usage metadata |
| 221 | + is logged to the materialization of the asset specified by ``asset_key``. |
| 222 | +
|
| 223 | + Args: |
| 224 | + context (AssetExecutionContext): The Dagster asset execution context. |
| 225 | + asset_key (AssetKey): The key of the asset to associate the metadata with. |
| 226 | + """ |
| 227 | + with self._get_model(context=context, asset_key=asset_key) as model: |
| 228 | + yield model |
| 229 | + |
| 230 | + @contextmanager |
| 231 | + def _get_model( |
| 232 | + self, |
| 233 | + context: Union[AssetExecutionContext, OpExecutionContext], |
| 234 | + asset_key: Optional[AssetKey] = None, |
| 235 | + ) -> Generator[GenerativeModel, None, None]: |
| 236 | + """Internal method to provide the model, applying usage tracking if in an asset context.""" |
| 237 | + # Store original method to restore later |
| 238 | + original_generate_content = None |
| 239 | + |
| 240 | + try: |
| 241 | + if isinstance(context, AssetExecutionContext): |
| 242 | + if asset_key is None: |
| 243 | + # For single assets or multi-assets with one output, determine the asset key from context. |
| 244 | + if len(context.assets_def.keys_by_output_name) > 1: |
| 245 | + raise DagsterInvariantViolationError( |
| 246 | + "The `asset_key` argument must be specified when calling `get_model` " |
| 247 | + "from a @multi_asset with more than one output. To specify the asset, " |
| 248 | + "use `get_model_for_asset` instead." |
| 249 | + ) |
| 250 | + asset_key = context.asset_key |
| 251 | + |
| 252 | + output_name = context.output_for_asset_key(asset_key) |
| 253 | + # Store original method before wrapping |
| 254 | + original_generate_content = self._generative_model.generate_content |
| 255 | + self._wrap_for_usage_tracking(context=context, output_name=output_name) |
| 256 | + |
| 257 | + # Yield the model for both Asset and Op contexts |
| 258 | + yield self._generative_model |
| 259 | + |
| 260 | + finally: |
| 261 | + # Restore original method if we wrapped it |
| 262 | + if original_generate_content is not None: |
| 263 | + self._generative_model.generate_content = original_generate_content |
| 264 | + |
| 265 | + def teardown_after_execution(self, context: InitResourceContext) -> None: |
| 266 | + """No explicit teardown is required for the Vertex AI SDK.""" |
| 267 | + pass |
0 commit comments