diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 072bcc5781..6cfaf48493 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1373,7 +1373,9 @@ def _start_query_ml_ddl( # https://cloud.google.com/bigquery/docs/customer-managed-encryption#encrypt-model job_config.destination_encryption_configuration = None - return bf_io_bigquery.start_query_with_client(self.bqclient, sql, job_config) + return bf_io_bigquery.start_query_with_client( + self.bqclient, sql, job_config, metrics=self._metrics + ) def _export( self, diff --git a/tests/system/large/ml/test_linear_model.py b/tests/system/large/ml/test_linear_model.py index f593ac2983..5c3f7631d9 100644 --- a/tests/system/large/ml/test_linear_model.py +++ b/tests/system/large/ml/test_linear_model.py @@ -128,10 +128,20 @@ def test_unordered_mode_linear_regression_configure_fit_score_predict( ] ] y_train = df[["body_mass_g"]] + + start_execution_count = df._block._expr.session._metrics.execution_count model.fit(X_train, y_train) + end_execution_count = df._block._expr.session._metrics.execution_count + # The fit function initiates two queries: the first generates and caches + # the training data, while the second creates and fits the model. + assert end_execution_count - start_execution_count == 2 # Check score to ensure the model was fitted + start_execution_count = end_execution_count result = model.score(X_train, y_train).to_pandas() + end_execution_count = df._block._expr.session._metrics.execution_count + assert end_execution_count - start_execution_count == 1 + utils.check_pandas_df_schema_and_index( result, columns=utils.ML_REGRESSION_METRICS, index=1 ) @@ -154,7 +164,10 @@ def test_unordered_mode_linear_regression_configure_fit_score_predict( assert reloaded_model.max_iterations == 20 assert reloaded_model.tol == 0.01 + start_execution_count = df._block._expr.session._metrics.execution_count pred = reloaded_model.predict(df) + end_execution_count = df._block._expr.session._metrics.execution_count + assert end_execution_count - start_execution_count == 1 utils.check_pandas_df_schema_and_index( pred, columns=("predicted_body_mass_g",), diff --git a/tests/system/small/ml/test_core.py b/tests/system/small/ml/test_core.py index 6b852e87af..65540e7e81 100644 --- a/tests/system/small/ml/test_core.py +++ b/tests/system/small/ml/test_core.py @@ -383,8 +383,14 @@ def test_model_forecast(time_series_bqml_arima_plus_model: core.BqmlModel): def test_model_register(ephemera_penguins_bqml_linear_model: core.BqmlModel): model = ephemera_penguins_bqml_linear_model + + start_execution_count = model.session._metrics.execution_count + model.register() + end_execution_count = model.session._metrics.execution_count + assert end_execution_count - start_execution_count == 1 + assert model.model.model_id is not None model_name = "bigframes_" + model.model.model_id # Only registered model contains the field, and the field includes project/dataset. Here only check model_id. diff --git a/tests/unit/ml/test_golden_sql.py b/tests/unit/ml/test_golden_sql.py index 65f079852e..ce05011546 100644 --- a/tests/unit/ml/test_golden_sql.py +++ b/tests/unit/ml/test_golden_sql.py @@ -36,6 +36,7 @@ def mock_session(): TEMP_MODEL_ID.project, TEMP_MODEL_ID.dataset_id ) mock_session._bq_kms_key_name = None + mock_session._metrics = None query_job = mock.create_autospec(bigquery.QueryJob) type(query_job).destination = mock.PropertyMock(