feat: enhance training model functionality and reporting

- Added `evidently` to requirements for improved model evaluation.
- Introduced `TrainModelResult` class with a `to_dict` method for better result handling.
- Updated `train_model` method to return a comprehensive training result, including run details.
- Enhanced `cleanup_run_directory` method in `DataManagerRepository` for improved resource management.
- Adjusted type hints in `TrainModel` for clarity and consistency.
This commit is contained in:
vitor-aignosi
2026-04-06 11:40:08 -03:00
parent a8b926649a
commit 1352d1ac8f
5 changed files with 62 additions and 22 deletions

View File

@@ -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,
{