SIENTIAPDE-1273

Update version and enhance metrics calculation in Laborious system

- Updated image tag in values.yaml from 1.1.0 to 1.1.1.
- Modified GITHUB_BRANCH environment variable for consistency.
- Added a new method `calculate_simple_metrics` in model_metrics.py to compute various model performance metrics including RMSE, MSE, MAE, and R2.
- Integrated the new metrics calculation into the worker setup, allowing for concurrent processing of simple metrics.
- Updated tests to cover the new metrics calculation functionality, ensuring comprehensive validation of the implementation.
This commit is contained in:
vitor-aignosi
2025-11-14 15:43:28 -03:00
parent 66193cea15
commit 64f0747e8e
8 changed files with 692 additions and 11 deletions

View File

@@ -639,3 +639,255 @@ async def test_get_drift_metrics_univariate_error(
else:
raise AssertionError('Expected Exception')
@mark.asyncio
async def test_calculate_simple_metrics_success_all_metrics(
model_metrics_activity
):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['rmse', 'mse', 'mae', 'r2'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 4
assert 'rmse' in result['metric'].values
assert 'mse' in result['metric'].values
assert 'mae' in result['metric'].values
assert 'r2' in result['metric'].values
assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']',
metadata['metadata']
)
model_metrics_activity.debug.assert_called_once()
@mark.asyncio
async def test_calculate_simple_metrics_success_rmse_only(
model_metrics_activity
):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['rmse'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'rmse'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\']',
metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_mse_only(
model_metrics_activity
):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['mse'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'mse'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'mse\']',
metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_mae_only(
model_metrics_activity
):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['mae'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'mae'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'mae\']',
metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_r2_only(
model_metrics_activity
):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['r2'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'r2'
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'r2\']',
metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_r2_zero_ss_tot(
model_metrics_activity
):
# Arrange
# All target values are the same, so ss_tot will be 0
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 1.0],
'prediction': [1.1, 1.1],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['r2'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 1
assert result['metric'].values[0] == 'r2'
assert result['value'].values[0] == 0.0 # Should return 0.0 when ss_tot == 0
assert result['model_id'].values[0] == 'test_model_id'
assert result['timestamp'].values[0] == '2023-05-26 11:12:28'
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'r2\']',
metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_multiple_metrics_subset(
model_metrics_activity
):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
})
input_data = {
**metadata,
'model_id': 'test_model_id',
'target_data': target_data.to_dict(),
'metrics': ['rmse', 'mae'],
'interval_minutes': 5,
}
# Act
result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data))
# Assert
assert len(result['metric']) == 2
assert 'rmse' in result['metric'].values
assert 'mae' in result['metric'].values
assert all(model_id == 'test_model_id' for model_id in result['model_id'].values)
assert all(timestamp == '2023-05-26 11:12:29' for timestamp in result['timestamp'].values)
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']',
metadata['metadata']
)

View File

@@ -39,6 +39,8 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
'chunk_period': 'hour',
}
target_name = input_data['model_config']['target']
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
@@ -96,7 +98,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': input_data['target_name'],
'target_name': target_name,
'drift_metrics': input_data['drift_metrics'],
'chunk_period': input_data['chunk_period'],
},
@@ -131,7 +133,7 @@ async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
}
@@ -163,10 +165,12 @@ async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
}
target_name = input_data['model_config']['target']
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = None
@@ -188,7 +192,7 @@ async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': input_data['target_name'],
'target_name': target_name,
'drift_metrics': input_data['drift_metrics'],
'chunk_period': input_data.get('chunk_period', 'min'),
},
@@ -211,11 +215,13 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
'source_table_name': 'test_source_table',
'target_table_name': 'test_target_table',
'interval': 60,
'target_name': 'test_target',
'model_config': {'target': 'test_target'},
'drift_metrics': ['psi', 'ks'],
# chunk_period not provided, should default to 'min'
}
target_name = input_data['model_config']['target']
target_data = {'data': 'test_target_data'}
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
@@ -237,7 +243,7 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': input_data['target_name'],
'target_name': target_name,
'drift_metrics': input_data['drift_metrics'],
'chunk_period': 'min', # Default value
},

View File

@@ -0,0 +1,228 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from laborious.activities.activities import Activities
from laborious.workflows.simple_metrics import SimpleMetrics
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
def simple_metrics() -> SimpleMetrics:
return SimpleMetrics()
metadata = {
'metadata': {
'model_id': 'test_model_id',
'model_name': 'test_model',
'workflow_name': 'simple_metrics',
'schedule_name': 'test_schedule',
},
}
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
'metrics': ['rmse', 'mse', 'mae', 'r2'],
}
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_activity_method = AsyncMock()
# Act
await simple_metrics.run(input_data)
# Assert - Check load_custom_query call
expected_query = f"""
select p."timestamp", p.prediction, ld.value as "target"
from {input_data['schema']}.{input_data['predictions_table_name']} p
inner join {input_data['schema']}.{input_data['data_table_name']} ld
on p."timestamp" = ld."timestamp"
where
p.model_id = {input_data['model_id']} and
p.prediction is not null and
ld.variable = '{input_data['model_config']['target']}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
order by
p."timestamp" desc;
"""
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
{
**metadata,
'query': expected_query,
'datetime_columns': ['timestamp'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.calculate_simple_metrics,
{
**metadata,
'model_id': input_data['model_id'],
'target_data': target_data,
'metrics': input_data['metrics'],
'interval_minutes': input_data['interval_minutes'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)
# Assert - Check export_data_to_postgres call
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.export_data_to_postgres,
{
**metadata,
'data': simple_metrics_data,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=ANY,
start_to_close_timeout=ANY,
)
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
'metrics': ['rmse', 'mse'],
}
target_data = None
workflow_mock.execute_local_activity_method.return_value = target_data
workflow_mock.execute_activity_method = AsyncMock()
# Act
await simple_metrics.run(input_data)
# Assert - Should not call calculate_simple_metrics or export
assert workflow_mock.execute_local_activity_method.call_count == 1
workflow_mock.execute_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
'metrics': ['rmse', 'mse'],
}
target_data = {'data': 'test_target_data'}
simple_metrics_data = None
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_activity_method = AsyncMock()
# Act
await simple_metrics.run(input_data)
# Assert - Should call calculate_simple_metrics but not export
assert workflow_mock.execute_local_activity_method.call_count == 2
workflow_mock.execute_activity_method.assert_not_called()
@mark.asyncio
@patch('laborious.workflows.simple_metrics.workflow', new_callable=AsyncMock)
async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
# Arrange
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
'model_id': 'test_model_id',
'interval_minutes': 60,
'model_config': {'target': 'test_target'},
'schema': 'test_schema',
'predictions_table_name': 'test_predictions_table',
'data_table_name': 'test_data_table',
'target_table_name': 'test_target_table',
# metrics not provided, should default to ['rmse', 'mse', 'mae', 'r2']
}
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_activity_method = AsyncMock()
# Act
await simple_metrics.run(input_data)
# Assert - Check calculate_simple_metrics call with default metrics
workflow_mock.execute_local_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
ANY,
retry_policy=ANY,
start_to_close_timeout=ANY,
),
call(
Activities.calculate_simple_metrics,
{
**metadata,
'model_id': input_data['model_id'],
'target_data': target_data,
'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value
'interval_minutes': input_data['interval_minutes'],
},
retry_policy=ANY,
start_to_close_timeout=ANY,
),
]
)