feat(simple_metrics): thread model_type from model_config to activity

Extracts model_type from model_config (optional, defaults to None when
absent) and passes it to calculate_simple_metrics. Required for the r2
lock - is_r2_supported() needs model_type to decide whether to compute r2.

Updates workflow tests to expect the new key in the activity-call dict.
This commit is contained in:
PedroHMCosme
2026-08-19 11:08:07 -03:00
parent 3b55072ab4
commit 33b4ae8406
2 changed files with 4 additions and 0 deletions

View File

@@ -31,6 +31,7 @@ class SimpleMetrics:
model_config = input_data['model_config']
target_name = model_config['target']
model_type = model_config.get('model_type')
query = f"""
select p."timestamp", p.prediction, ld.value as "target"
@@ -70,6 +71,7 @@ class SimpleMetrics:
'target_data': target_data,
'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
'interval_minutes': interval_minutes,
'model_type': model_type,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),

View File

@@ -103,6 +103,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
'target_data': target_data,
'metrics': input_data['metrics'],
'interval_minutes': input_data['interval_minutes'],
'model_type': None,
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -209,6 +210,7 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
'target_data': target_data,
'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value
'interval_minutes': input_data['interval_minutes'],
'model_type': None,
},
retry_policy=ANY,
start_to_close_timeout=ANY,