From e2b2f702f812fa9c57256c2cd61137d6ef7c313f Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 16 Apr 2026 14:51:41 -0300 Subject: [PATCH] feat: integrate SientiaModel wrapper and enhance logging in Training class - Imported SientiaModel to standardize the wrapper type in the Training class. - Added conditional logging to capture training process details when a logger is provided, improving traceability during model training. --- model_manager/activities/training.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index 12302e5..fb14f98 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -6,6 +6,7 @@ The activity extends BaseActivity and receives pre-downloaded files and raises `ModelTrainingError` when training fails. """ +from sientia_model.wrappers.sientia_model import SientiaModel from temporalio import activity, workflow with workflow.unsafe.imports_passed_through(): @@ -228,6 +229,9 @@ class Training(SientiaMonitoring): metadata=metadata, ) + if self.logger is not None: + wrapper.logger = self.logger.base_logger + self.info(f'Training model for {train_params.model_type}', metadata) train_data = train_result.train_data val_data = train_result.val_data @@ -317,7 +321,7 @@ class Training(SientiaMonitoring): self, train_result: TrainModelResult, train_params: TrainModelParams, - wrapper: Any, + wrapper: SientiaModel, metadata: dict[str, Any] | None, ) -> None: self.info(f'Generating report for {train_params.model_type}', metadata)