diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index b56feb6..c2f80ff 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -23,6 +23,9 @@ with workflow.unsafe.imports_passed_through(): from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.repository.data_manager_repository import DataManagerRepository + from model_manager.utils.models.train_model_result import TrainModelResult + + import mlflow class Training(SientiaMonitoring): @@ -146,7 +149,7 @@ class Training(SientiaMonitoring): raise @activity.defn(name='train_model') - async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]: + async def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]: """ Train a machine learning model. @@ -246,22 +249,24 @@ class Training(SientiaMonitoring): tags=None, metadata=metadata, ) as run_info: + train_result.run_name = run_info.run_name + train_result.run_id = run_info.run_id + + train_result = self.data_manager_repository.generate_report( + train_result, + metadata=metadata, + ) + + if train_result.report_path is None or train_result.train_data_path is None or train_result.test_data_path is None: + raise ValueError('Report path, train data path, or test data path is not set') + wrapper.store_model(name=train_params.model_name) - train_result.run_name = run_info.run_name - train_result.run_id = run_info.run_id + mlflow.log_artifact(train_result.report_path) + mlflow.log_artifact(train_result.train_data_path) + mlflow.log_artifact(train_result.test_data_path) - train_result = self.data_manager_repository.generate_report( - train_result, - metadata=metadata, - ) - - model_saved = True - - return { - 'run_name': run_info.run_name, - 'run_id': run_info.run_id, - } + return train_result.to_dict() except Exception as e: # noqa: BLE001 error_msg = f'Error training model - error: {str(e)}' @@ -296,7 +301,7 @@ class Training(SientiaMonitoring): run_dir = input_data.get('run_dir', '') try: - self.model_repository.cleanup_run_directory(run_dir) + self.data_manager_repository.cleanup_run_directory(run_dir, metadata) except Exception as e: # noqa: BLE001 error_msg = f'Error cleaning up resources - Run directory: {run_dir}, Error: {str(e)}' diff --git a/model_manager/utils/models/train_model_result.py b/model_manager/utils/models/train_model_result.py index 98edc56..dc006b9 100644 --- a/model_manager/utils/models/train_model_result.py +++ b/model_manager/utils/models/train_model_result.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from typing import Any import pandas as pd @@ -50,3 +51,9 @@ class TrainModelResult: test_data_path: str | None = None run_dir: str | None = None + + def to_dict(self) -> dict[str, Any]: + """ + Convert TrainModelResult to a dictionary. + """ + return self.__dict__ \ No newline at end of file diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py index 90061d2..e0eab11 100644 --- a/model_manager/utils/repository/data_manager_repository.py +++ b/model_manager/utils/repository/data_manager_repository.py @@ -15,8 +15,10 @@ and this repository focuses solely on preparing data structures for them. from datetime import datetime from io import BytesIO +import json from os import makedirs, path from typing import Any +from shutil import rmtree import numpy as np import pandas as pd @@ -433,3 +435,24 @@ class DataManagerRepository(SientiaMonitoring): json.dump(data.equation, f, indent=2, ensure_ascii=False) return data + + def cleanup_run_directory(self, run_dir: str, metadata: dict[str, Any] | None = None) -> None: + """ + Clean up temporary run directory after model training. + + This activity deletes the temporary directory created during model training + and artifact generation. It implements idempotent cleanup to handle cases + where the directory may have already been deleted. + + Args: + run_dir (str): Path to the run directory to delete + """ + if not run_dir: + self.info('No run directory specified, skipping cleanup') + return + + if path.exists(run_dir): + rmtree(run_dir) + self.info(f'Run directory deleted successfully: {run_dir}') + else: + self.info(f'Run directory already deleted: {run_dir}') diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index 2960fc6..4f5e255 100644 --- a/model_manager/workflows/train_model.py +++ b/model_manager/workflows/train_model.py @@ -71,7 +71,7 @@ class TrainModel: """ @workflow.run - async def run(self, input_data: dict[str, Any]) -> dict[str, str | None]: + async def run(self, input_data: dict[str, Any]) -> dict[str, str | None] | None: """ Execute the complete model training workflow. @@ -112,7 +112,7 @@ class TrainModel: ) training_succeeded = False - train_result: dict[str, str | None] + train_result: dict[str, str | None] | None = None try: train_result = await self._train_model( @@ -123,10 +123,11 @@ class TrainModel: training_succeeded = True finally: try: - await self._cleanup_resources( - run_dir=train_result.get('run_dir'), - metadata=metadata, - ) + if train_result is not None: + await self._cleanup_resources( + run_dir=train_result.get('run_dir'), + metadata=metadata, + ) except Exception: if training_succeeded: raise @@ -290,7 +291,7 @@ class TrainModel: async def _cleanup_resources( self, - run_dir: str, + run_dir: str | None, metadata: dict[str, Any], ) -> None: """ @@ -302,6 +303,9 @@ class TrainModel: run_dir: Temporary directory to remove metadata: Workflow execution metadata """ + if run_dir is None: + return + await workflow.execute_activity_method( Activities.cleanup_resources, { diff --git a/requirements.txt b/requirements.txt index 3e48f0d..dd2264c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ botocore==1.40.55 git+https://github.com/Aignosi/sientia-dataops-library.git@v1.10.1 prometheus-client==0.23.1 beautifulsoup4==4.12.3 +evidently \ No newline at end of file