diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py index 399f9bb..03b31ec 100644 --- a/model_manager/utils/repository/data_manager_repository.py +++ b/model_manager/utils/repository/data_manager_repository.py @@ -527,6 +527,32 @@ class DataManagerRepository(SientiaMonitoring): self.error(error_msg, metadata) raise OSError(error_msg) from e + def _save_data_csv( + self, frame: pd.DataFrame, file_path: str, metadata: dict[str, Any] | None = None + ) -> None: + """ + Write a train/test frame to CSV with its time index as a ``timestamp`` column. + + laborious drift reads ``test_data.csv`` from MLflow as its reference and needs + ``timestamp`` (SIENTIAPDE-2201). If a column already has that name, the frame is + written without the index and that model's drift keeps falling back to the target + window: training never fails over drift. + + Args: + frame: Frame indexed by the UTC training timestamps. + file_path: Destination CSV path. + metadata: Workflow metadata for logging. + """ + if 'timestamp' in frame.columns: + self.warning( + f"Column 'timestamp' already exists; {path.basename(file_path)} saved without " + 'the time column (drift reference falls back to the target window)', + metadata, + ) + frame.to_csv(file_path, index=False) + return + frame.to_csv(file_path, index_label='timestamp') + def generate_report( self, data: TrainModelResult, metadata: dict[str, Any] | None = None ) -> TrainModelResult: @@ -598,11 +624,11 @@ class DataManagerRepository(SientiaMonitoring): # Save training / validation CSVs using the same frames as the report (includes # literal `target` alias for Evidently, plus predictions and float-cast features). data.train_data_path = path.join(data.run_dir, 'train_data.csv') - reference_data_float.to_csv(data.train_data_path, index=False) + self._save_data_csv(reference_data_float, data.train_data_path, metadata) # Save test data CSV data.test_data_path = path.join(data.run_dir, 'test_data.csv') - current_data_float.to_csv(data.test_data_path, index=False) + self._save_data_csv(current_data_float, data.test_data_path, metadata) # Save equation as JSON if data.equation is not None and data.params.model_type == 'linear_regression': diff --git a/tests/utils/repository/test_data_manager_repository.py b/tests/utils/repository/test_data_manager_repository.py index a7c478c..09c52b3 100644 --- a/tests/utils/repository/test_data_manager_repository.py +++ b/tests/utils/repository/test_data_manager_repository.py @@ -514,6 +514,58 @@ def test_generate_report_adds_target_alias_for_reports(tmp_path): assert test_csv['target'].equals(test_csv['t']) +_UTC_INDEX = pd.DatetimeIndex(['2024-01-01 02:00', '2024-01-01 01:00'], tz='UTC') + + +def _dated_result(**train_columns: list[float]) -> TrainModelResult: + """Frames indexed by UTC time, like prepare_training_data + predict produce them.""" + columns = train_columns or {'v1': [1.0, 2.0], 't': [1.0, 2.0]} + return TrainModelResult( + params=TrainModelParams.from_dict(_minimal_dict_for_prepare()), + train_data=pd.DataFrame(columns, index=_UTC_INDEX), + val_data=pd.DataFrame(columns, index=_UTC_INDEX), + y_train_pred=pd.DataFrame({'t': [1.0, 2.0]}, index=_UTC_INDEX), + y_pred=pd.DataFrame({'t': [1.0, 2.0]}, index=_UTC_INDEX), + run_name='testrun', + ) + + +def _generate_report(repo, tmr: TrainModelResult, tmp_path) -> TrainModelResult: + with ( + patch.object(repo, '_get_reports_directory', return_value=str(tmp_path)), + patch('model_manager.utils.repository.data_manager_repository.Reports'), + ): + return repo.generate_report(tmr, {}) + + +def test_generate_report_csvs_keep_the_time_index_as_timestamp(tmp_path): + """laborious drift reads test_data.csv as its reference and needs ``timestamp``.""" + repo = dmr.DataManagerRepository(MagicMock()) + + out = _generate_report(repo, _dated_result(), tmp_path) + + for csv_path in (out.train_data_path, out.test_data_path): + csv = pd.read_csv(csv_path) + assert list(csv.columns) == ['timestamp', 'prediction', 'v1', 't', 'target'] + assert pd.to_datetime(csv['timestamp']).tolist() == list(_UTC_INDEX) + + +def test_generate_report_csvs_skip_timestamp_when_a_column_already_has_the_name(tmp_path): + """A variable literally named ``timestamp`` wins; training must not fail over drift.""" + repo = dmr.DataManagerRepository(MagicMock()) + tmr = _dated_result(timestamp=[5.0, 6.0], v1=[1.0, 2.0], t=[1.0, 2.0]) + + with patch.object(repo, 'warning') as warning: + out = _generate_report(repo, tmr, tmp_path) + + for csv_path in (out.train_data_path, out.test_data_path): + csv = pd.read_csv(csv_path) + assert list(csv.columns) == ['prediction', 'timestamp', 'v1', 't', 'target'] + assert csv['timestamp'].tolist() == [5.0, 6.0] + assert warning.call_count == 2 + assert "Column 'timestamp' already exists" in warning.call_args.args[0] + + def test_generate_report_skips_equation_file_when_not_linear(tmp_path): repo = dmr.DataManagerRepository(MagicMock()) p = TrainModelParams.from_dict(_minimal_dict_for_prepare())