fix: model.fit metric not collected issue. (#1085)

* fix: model.fit metric not collected issue.

* update

* fix unit test

* update code and test

* update code

* update comment

* update test

Author: Genesis929

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,

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",),

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.

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(