feat: temporary resources no longer use BigQuery Sessions (#194)

* feat: temporary resources no longer use BigQuery Sessions

This allows multiple queries to run in parallel.

* use anonymous dataset for models

* remove reference to bq session

* avoid 'model already exists' error

* remove global _model_id

Author: tswast

bigframes/ml/core.py

@@ -16,12 +16,14 @@
 
 from __future__ import annotations
 
+import datetime
 from typing import Callable, cast, Iterable, Mapping, Optional, Union
 import uuid
 
 from google.cloud import bigquery
 
 import bigframes
+import bigframes.constants as constants
 from bigframes.ml import sql as ml_sql
 import bigframes.pandas as bpd
 
@@ -192,24 +194,27 @@ def register(self, vertex_ai_model_id: Optional[str] = None) -> BqmlModel:
 
 class BqmlModelFactory:
     def __init__(self):
-        model_id = self._create_temp_model_id()
-        self._model_creation_sql_generator = ml_sql.ModelCreationSqlGenerator(model_id)
+        self._model_creation_sql_generator = ml_sql.ModelCreationSqlGenerator()
 
-    def _create_temp_model_id(self) -> str:
-        return uuid.uuid4().hex
-
-    def _reset_model_id(self):
-        self._model_creation_sql_generator._model_id = self._create_temp_model_id()
+    def _create_model_ref(
+        self, dataset: bigquery.DatasetReference
+    ) -> bigquery.ModelReference:
+        return bigquery.ModelReference.from_string(
+            f"{dataset.project}.{dataset.dataset_id}.{uuid.uuid4().hex}"
+        )
 
     def _create_model_with_sql(self, session: bigframes.Session, sql: str) -> BqmlModel:
         # fit the model, synchronously
         _, job = session._start_query(sql)
 
         # real model path in the session specific hidden dataset and table prefix
-        model_name_full = f"{job.destination.dataset_id}.{job.destination.table_id}"
-        model = session.bqclient.get_model(model_name_full)
+        model_name_full = f"{job.destination.project}.{job.destination.dataset_id}.{job.destination.table_id}"
+        model = bigquery.Model(model_name_full)
+        model.expires = (
+            datetime.datetime.now(datetime.timezone.utc) + constants.DEFAULT_EXPIRATION
+        )
+        model = session.bqclient.update_model(model, ["expires"])
 
-        self._reset_model_id()
         return BqmlModel(session, model)
 
     def create_model(
@@ -219,7 +224,7 @@ def create_model(
         transforms: Optional[Iterable[str]] = None,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> BqmlModel:
-        """Create a session-temporary BQML model with the CREATE MODEL statement
+        """Create a session-temporary BQML model with the CREATE OR REPLACE MODEL statement
 
         Args:
             X_train: features columns for training
@@ -241,9 +246,11 @@ def create_model(
             options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()})
 
         session = X_train._session
+        model_ref = self._create_model_ref(session._anonymous_dataset)
 
         sql = self._model_creation_sql_generator.create_model(
             source_df=input_data,
+            model_ref=model_ref,
             transforms=transforms,
             options=options,
         )
@@ -272,9 +279,11 @@ def create_time_series_model(
         options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]})
 
         session = X_train._session
+        model_ref = self._create_model_ref(session._anonymous_dataset)
 
         sql = self._model_creation_sql_generator.create_model(
             source_df=input_data,
+            model_ref=model_ref,
             transforms=transforms,
             options=options,
         )
@@ -287,7 +296,7 @@ def create_remote_model(
         connection_name: str,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> BqmlModel:
-        """Create a session-temporary BQML remote model with the CREATE MODEL statement
+        """Create a session-temporary BQML remote model with the CREATE OR REPLACE MODEL statement
 
         Args:
             connection_name:
@@ -298,8 +307,10 @@ def create_remote_model(
         Returns:
             BqmlModel: a BqmlModel wrapping a trained model in BigQuery
         """
+        model_ref = self._create_model_ref(session._anonymous_dataset)
         sql = self._model_creation_sql_generator.create_remote_model(
             connection_name=connection_name,
+            model_ref=model_ref,
             options=options,
         )
 
@@ -310,15 +321,17 @@ def create_imported_model(
         session: bigframes.Session,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> BqmlModel:
-        """Create a session-temporary BQML imported model with the CREATE MODEL statement
+        """Create a session-temporary BQML imported model with the CREATE OR REPLACE MODEL statement
 
         Args:
             options: a dict of options to configure the model. Generates a BQML OPTIONS
                 clause
 
         Returns: a BqmlModel, wrapping a trained model in BigQuery
         """
+        model_ref = self._create_model_ref(session._anonymous_dataset)
         sql = self._model_creation_sql_generator.create_imported_model(
+            model_ref=model_ref,
             options=options,
         )
 

bigframes/ml/sql.py

@@ -18,6 +18,8 @@
 
 from typing import Iterable, Mapping, Optional, Union
 
+import google.cloud.bigquery
+
 import bigframes.constants as constants
 import bigframes.pandas as bpd
 
@@ -121,22 +123,26 @@ def ml_label_encoder(
 class ModelCreationSqlGenerator(BaseSqlGenerator):
     """Sql generator for creating a model entity. Model id is the standalone id without project id and dataset id."""
 
-    def __init__(self, model_id: str):
-        self._model_id = model_id
+    def _model_id_sql(
+        self,
+        model_ref: google.cloud.bigquery.ModelReference,
+    ):
+        return f"`{model_ref.project}`.`{model_ref.dataset_id}`.`{model_ref.model_id}`"
 
     # Model create and alter
     def create_model(
         self,
         source_df: bpd.DataFrame,
+        model_ref: google.cloud.bigquery.ModelReference,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
         transforms: Optional[Iterable[str]] = None,
     ) -> str:
-        """Encode the CREATE TEMP MODEL statement for BQML"""
+        """Encode the CREATE OR REPLACE MODEL statement for BQML"""
         source_sql = source_df.sql
         transform_sql = self.transform(*transforms) if transforms is not None else None
         options_sql = self.options(**options)
 
-        parts = [f"CREATE TEMP MODEL `{self._model_id}`"]
+        parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
         if transform_sql:
             parts.append(transform_sql)
         if options_sql:
@@ -147,25 +153,27 @@ def create_model(
     def create_remote_model(
         self,
         connection_name: str,
+        model_ref: google.cloud.bigquery.ModelReference,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> str:
-        """Encode the CREATE TEMP MODEL statement for BQML remote model."""
+        """Encode the CREATE OR REPLACE MODEL statement for BQML remote model."""
         options_sql = self.options(**options)
 
-        parts = [f"CREATE TEMP MODEL `{self._model_id}`"]
+        parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
         parts.append(self.connection(connection_name))
         if options_sql:
             parts.append(options_sql)
         return "\n".join(parts)
 
     def create_imported_model(
         self,
+        model_ref: google.cloud.bigquery.ModelReference,
         options: Mapping[str, Union[str, int, float, Iterable[str]]] = {},
     ) -> str:
-        """Encode the CREATE TEMP MODEL statement for BQML remote model."""
+        """Encode the CREATE OR REPLACE MODEL statement for BQML remote model."""
         options_sql = self.options(**options)
 
-        parts = [f"CREATE TEMP MODEL `{self._model_id}`"]
+        parts = [f"CREATE OR REPLACE MODEL {self._model_id_sql(model_ref)}"]
         if options_sql:
             parts.append(options_sql)
         return "\n".join(parts)

bigframes/session/__init__.py

@@ -161,7 +161,7 @@ def __init__(
                 application_name=context.application_name,
             )
 
-        self._create_and_bind_bq_session()
+        self._create_bq_datasets()
         self.ibis_client = typing.cast(
             ibis_bigquery.Backend,
             ibis.bigquery.connect(
@@ -210,19 +210,12 @@ def _project(self):
 
     def __hash__(self):
         # Stable hash needed to use in expression tree
-        return hash(self._session_id)
-
-    def _create_and_bind_bq_session(self):
-        """Create a BQ session and bind the session id with clients to capture BQ activities:
-        go/bigframes-transient-data"""
-        job_config = bigquery.QueryJobConfig(create_session=True)
-        # Make sure the session is a new one, not one associated with another query.
-        job_config.use_query_cache = False
-        query_job = self.bqclient.query(
-            "SELECT 1", job_config=job_config, location=self._location
-        )
+        return hash(str(self._anonymous_dataset))
+
+    def _create_bq_datasets(self):
+        """Create and identify dataset(s) for temporary BQ resources."""
+        query_job = self.bqclient.query("SELECT 1", location=self._location)
         query_job.result()  # blocks until finished
-        self._session_id = query_job.session_info.session_id
 
         # The anonymous dataset is used by BigQuery to write query results and
         # session tables. BigQuery DataFrames also writes temp tables directly
@@ -235,17 +228,6 @@ def _create_and_bind_bq_session(self):
             query_destination.dataset_id,
         )
 
-        self.bqclient.default_query_job_config = bigquery.QueryJobConfig(
-            connection_properties=[
-                bigquery.ConnectionProperty("session_id", self._session_id)
-            ]
-        )
-        self.bqclient.default_load_job_config = bigquery.LoadJobConfig(
-            connection_properties=[
-                bigquery.ConnectionProperty("session_id", self._session_id)
-            ]
-        )
-
         # Dataset for storing remote functions, which don't yet
         # support proper session temporary storage yet
         self._session_dataset = bigquery.Dataset(
@@ -254,28 +236,7 @@ def _create_and_bind_bq_session(self):
         self._session_dataset.location = self._location
 
     def close(self):
-        """Terminated the BQ session, otherwises the session will be terminated automatically after
-        24 hours of inactivity or after 7 days."""
-        if self._session_id is not None and self.bqclient is not None:
-            abort_session_query = "CALL BQ.ABORT_SESSION('{}')".format(self._session_id)
-            try:
-                query_job = self.bqclient.query(abort_session_query)
-                query_job.result()  # blocks until finished
-            except google.api_core.exceptions.BadRequest as exc:
-                # Ignore the exception when the BQ session itself has expired
-                # https://cloud.google.com/bigquery/docs/sessions-terminating#auto-terminate_a_session
-                if not exc.message.startswith(
-                    f"Session {self._session_id} has expired and is no longer available."
-                ):
-                    raise
-            except google.auth.exceptions.RefreshError:
-                # The refresh token may itself have been invalidated or expired
-                # https://developers.google.com/identity/protocols/oauth2#expiration
-                # Don't raise the exception in this case while closing the
-                # BigFrames session, so that the end user has a path for getting
-                # out of a bad session due to unusable credentials.
-                pass
-            self._session_id = None
+        """No-op. Temporary resources are deleted after 7 days."""
 
     def read_gbq(
         self,

tests/system/small/test_progress_bar.py

@@ -90,9 +90,6 @@ def assert_loading_msg_exist(capystOut: str, pattern=job_load_message_regex):
 
 def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame):
     bf.options.display.progress_bar = "terminal"
-    penguins_df_default_index._block._expr.session.bqclient.default_query_job_config.use_query_cache = (
-        False
-    )
     penguins_df_default_index.to_pandas()
     query_job_repr = formatting_helpers.repr_query_job_html(
         penguins_df_default_index.query_job
@@ -109,9 +106,6 @@ def test_query_job_repr_html(penguins_df_default_index: bf.dataframe.DataFrame):
 
 
 def test_query_job_repr(penguins_df_default_index: bf.dataframe.DataFrame):
-    penguins_df_default_index._block._expr.session.bqclient.default_query_job_config.use_query_cache = (
-        False
-    )
     penguins_df_default_index.to_pandas()
     query_job_repr = formatting_helpers.repr_query_job(
         penguins_df_default_index.query_job

tests/system/small/test_session.py

@@ -19,7 +19,6 @@
 import typing
 from typing import List
 
-import google.api_core.exceptions
 import google.cloud.bigquery as bigquery
 import numpy as np
 import pandas as pd
@@ -985,26 +984,3 @@ def test_read_json_gcs_default_engine(session, scalars_dfs, gcs_folder):
 
     assert df.shape[0] == scalars_df.shape[0]
     pd.testing.assert_series_equal(df.dtypes, scalars_df.dtypes)
-
-
-def test_session_id(session):
-    assert session._session_id is not None
-
-    # BQ client always runs query within the opened session.
-    query_job = session.bqclient.query("SELECT 1")
-    assert query_job.session_info.session_id == session._session_id
-
-    # TODO(chelsealin): Verify the session id can be binded with a load job.
-
-
-@pytest.mark.flaky(retries=2)
-def test_to_close_session():
-    session = bigframes.Session()
-    assert session._session_id is not None
-    session.close()
-    assert session._session_id is None
-
-    # Session has expired and is no longer available.
-    with pytest.raises(google.api_core.exceptions.BadRequest):
-        query_job = session.bqclient.query("SELECT 1")
-        query_job.result()  # blocks until finished