14 Commits

Author SHA1 Message Date
vitor-aignosi
d4902d1f7c Add .gitignore
Some checks failed
Quality gate / quality-gate (push) Has been cancelled
2026-08-26 16:15:14 -03:00
vitor-aignosi
295b8a0e81 Add .github 2026-08-26 16:15:14 -03:00
vitor-aignosi
20e855707c Add .env.example 2026-08-26 16:15:14 -03:00
vitor-aignosi
41efa8525a Add sonar-project.properties 2026-08-26 16:15:14 -03:00
vitor-aignosi
65a282e0bc Add scripts 2026-08-26 16:15:14 -03:00
vitor-aignosi
6737604d6b Add requirements.txt 2026-08-26 16:15:14 -03:00
vitor-aignosi
7e5443501a Add requirements-local.txt 2026-08-26 16:15:14 -03:00
vitor-aignosi
fb6acce8f6 Add requirements-light.txt 2026-08-26 16:15:14 -03:00
vitor-aignosi
26e15c5944 Add requirements-dev.txt 2026-08-26 16:15:14 -03:00
vitor-aignosi
c78ddb1e2b Add README.md 2026-08-26 16:15:14 -03:00
vitor-aignosi
ef593009fe Add pyproject.toml 2026-08-26 16:15:14 -03:00
vitor-aignosi
a89e40ba51 Add Makefile 2026-08-26 16:15:14 -03:00
vitor-aignosi
44ce1655d3 Add laborious 2026-08-26 16:15:14 -03:00
vitor-aignosi
6eb7c49aa7 Add git-requirements-mapping.txt 2026-08-26 16:15:14 -03:00
55 changed files with 13088 additions and 0 deletions

73
.env.example Normal file
View File

@@ -0,0 +1,73 @@
POSTGRES_HOST="paradedb-rw.paradedb.svc.cluster.local"
POSTGRES_PORT="5432"
POSTGRES_USER="sientia"
POSTGRES_PASSWORD="password"
POSTGRES_DBNAME="sientia"
POSTGRES_MIN_CONNECTIONS="10"
POSTGRES_MAX_CONNECTIONS="30"
MLFLOW_HOST="http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local"
MLFLOW_PORT="80"
MLFLOW_USERNAME="aignosi"
MLFLOW_PASSWORD="mlflow_password"
OPC_ID="1"
OPC_URL="opc.tcp://sientia-opc-simulator-opc.sientia.svc.cluster.local:4840"
LOG_LEVEL="DEBUG"
HTTP_METRICS_PORT="9090"
HTTP_SDK_METRICS_PORT="9091"
PROJECT_NAME="sientia-laborious"
TEMPORAL_HOST="temporal-frontend.temporal.svc.cluster.local:7233"
TEMPORAL_NAMESPACE="laborious"
MONGODB_USERNAME="mongo_user"
MONGODB_PASSWORD="mongo_db_password"
MONGODB_URL="my-release-mongodb.mongodb.svc.cluster.local:27017"
MONGODB_DATABASE="sientia"
MONGODB_TTL_INDEX_HOURS="1"
MINIO_ENDPOINT_URL="http://localhost:9000"
MINIO_ACCESS_KEY="sientia"
MINIO_SECRET_KEY="sientia"
MINIO_REGION_NAME="sa-east-1"
MINIO_DEFAULT_BUCKET="sientia"
# Suffix for every Temporal task queue name; required, non-empty.
# The import worker listens on "import_model-${RUNTIME}-queue".
RUNTIME="local"
# --- Model import (.sientia bundles) -----------------------------------------
# The import worker is registered only when IMPORT_MODEL_ENABLED is truthy
# (1/true/yes/on). Do not turn it on before the V12 migration is applied to
# public.experiment_run in the BFF database.
IMPORT_MODEL_ENABLED="false"
# base64 of the 32-byte AES-256-GCM key that wraps bundle passwords.
# Generate a throwaway one with: openssl rand -base64 32
IMPORT_PASSWORD_KEY="replace-with-base64-of-32-random-bytes"
# Where the frontend uploads the bundle. The object key must sit under the prefix.
IMPORT_BUNDLE_BUCKET="sientia"
IMPORT_BUNDLE_PREFIX="imported_models/"
# MongoDB collection holding the model listing the import writes its document into.
IMPORT_MODELS_COLLECTION="models"
# The import log lives in the Spring Boot BFF database on the POSTGRES_* server
# above: same host, port, user and password, another database name. Only that
# name is configuration here; the schema (public) and the table (experiment_run)
# are constants of the code. Set IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} only
# if the log really is on a different server than POSTGRES_* — they override it.
IMPORT_STATUS_DB_NAME="sientia-core-mlops-bff"
# --- scripts/import_model_manual_run.py only ---------------------------------
# Inputs for the manual run that stands in for the frontend. The script writes
# one object and one row to whatever the variables above name.
IMPORT_SCRIPT_BUNDLE_PATH="tests/fixtures/bundle/golden_model_v3.sientia"
IMPORT_SCRIPT_BUNDLE_PASSWORD="the-password-the-bundle-was-written-with"
IMPORT_SCRIPT_USERNAME="manual.run@example.com"
# Cell 2's guard: restate POSTGRES_HOST here (or in the cell) after reading what cell 1
# printed. Left empty on purpose — the script refuses to write until it matches.
IMPORT_SCRIPT_CONFIRM_TARGET=""

31
.github/workflows/quality-gate.yml vendored Normal file
View File

@@ -0,0 +1,31 @@
name: Quality gate
on:
push:
branches:
- main
- 'release/**'
- 'feature/**'
pull_request:
branches:
- main
- 'release/**'
- 'feature/**'
types: [ opened, synchronize, reopened ]
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
jobs:
quality-gate:
uses: Aignosi/github_workflow_templates/.github/workflows/python-quality-gate.yml@main
permissions:
contents: read
pull-requests: write
issues: write
with:
project_name: 'laborious'
repositories: 'sientia-dataops-library, sientia-mlops-library'
requirements_file: 'requirements-light.txt'
secrets: inherit

25
.github/workflows/release.yml vendored Normal file
View File

@@ -0,0 +1,25 @@
name: Create Release on Merge to Main
on:
pull_request:
types: [closed]
branches:
- main
workflow_dispatch:
inputs:
version:
description: 'Version to release'
required: false
type: string
jobs:
release:
if: |
(github.event_name == 'pull_request' && github.event.pull_request.merged == true) ||
github.event_name == 'workflow_dispatch'
uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-release.yml@main
permissions: write-all
with:
project_name: 'laborious'
release_version: ${{ github.event.inputs.version || '' }}
secrets: inherit

59
.gitignore vendored Normal file
View File

@@ -0,0 +1,59 @@
# Ignorar volumes do Docker
docker-compose.override.yml
**/db_data/
**/kafka-volume/
**/zookeeper-volume/
**/mage_data/
**/minio_data/
**/venv/
**/certs/*.pem
**/certs/*.der
**/certs/*.csr
**/deploy/*.yaml
scouter/.file_versions/
scouter/pipelines/**/triggers.yaml
**/postgres_data/**
# Ignorar arquivos e diretórios de cache do Python
__pycache__/
*.pyc
*.pyo
*.pyd
# Ignorar logs
*.log
# Ignorar arquivos de configuração locais
.vscode/
.pytest_cache/
.idea/
*.swp
# Ignorar arquivos temporários
*.tmp
*.bak
*.old
.secret
# Ignorar coverage
htmlcov/
.coverage
coverage.xml
# git keys
git_key*
git_log
.env
tmp/
catboost_info/
.ruff_cache/
.mypy_cache/
mlruns/
relatorio*
openspec/
.cursor/
.claude/

7
Makefile Normal file
View File

@@ -0,0 +1,7 @@
VERSION = 1.0.8
name = sientia-laborious
# ENVIRONMENT = production
docker-hub:
@docker build --no-cache -t aignosi.azurecr.io/$(name):$(VERSION) .
@docker push aignosi.azurecr.io/$(name):$(VERSION)

1423
README.md Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,2 @@
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git:sientia-do
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git:sientia

0
laborious/__init__.py Normal file
View File

View File

View File

@@ -0,0 +1,215 @@
from sientia_do.observability.metrics_controller import MetricsController
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.repository.minio_repository import MinioRepository
from sientia_do.repository.mongodb_repository import MongoDBRepository
from laborious.activities.api import API
from laborious.activities.gates import Gates
from laborious.activities.mlflow import MLFlow
from laborious.activities.model_import import ModelImport
from laborious.activities.model_metrics import ModelMetrics
from laborious.activities.opc import OPC
from laborious.activities.storage import Storage
from laborious.utils.connectors_config import (
build_import_config,
build_import_status_config,
)
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
"""
Main activities orchestrator for the Laborious system.
This class combines functionality from multiple activity classes to provide
a unified interface for all workflow operations. It manages database connections,
MLFlow model interactions, data quality validation, and OPC server communications.
The class implements multiple inheritance to combine specialized functionality:
- Storage: Database operations and data persistence
- MLFlow: Model inference and transformation operations
- Gates: Data quality validation and filtering mechanisms
- OPC: Real-time data export to OPC servers
- ModelMetrics: Model performance metrics and drift detection
- API: PI Web API export operations for industrial systems
Attributes:
postgres_config (dict): PostgreSQL connection configuration
mlflow_config (dict): MLFlow server configuration
opc_config (dict): OPC server configuration
pi_web_api_config (dict): PI Web API server configuration
logger (Logger): Logging and observability instance
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(
self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
minio_config: dict[str, Any],
opc_config: dict[str, Any],
pi_web_api_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler,
mongo_config: dict[str, Any] | None = None,
import_config: dict[str, Any] | None = None,
import_status_config: dict[str, Any] | None = None,
):
"""
Initialize the Activities orchestrator with all required configurations.
This constructor initializes all parent classes with their respective
configurations and sets up the foundation for all activity operations.
Args:
postgres_config: PostgreSQL connection configuration dictionary
Required keys: host, port, user, password, dbname, min_connections, max_connections
mlflow_config: MLFlow server configuration dictionary
Required keys: host, port, username, password
opc_config: OPC server configuration dictionary
Can contain multiple server configurations
pi_web_api_config: PI Web API server configuration dictionary
Required keys: base_url, auth_type, auth_token
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
mongo_config: MongoDB configuration for the model listing the import writes. When
absent, the import activities that need it refuse rather than guess — the
notification handler keeps its own client either way.
import_config: `build_import_config()`; read from the environment when absent
import_status_config: `build_import_status_config()`; read from the environment when
absent. This is the BFF database holding the import log: another database on the
`POSTGRES_*` server, and a second connection that is never the one Storage holds.
Raises:
Exception: If any parent class initialization fails
"""
metrics_controller = MetricsController(logger=logger)
minio_repository = MinioRepository(
endpoint=minio_config['endpoint_url'],
access_key=minio_config['access_key'],
secret_key=minio_config['secret_key'],
bucket=minio_config['default_bucket'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
secure=minio_config['secure'],
)
# Initialize parent classes
Storage.__init__(
self,
host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
retention_hours=minio_config['retention_hours'],
minio_repository=minio_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
MLFlow.__init__(
self,
mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
minio_repository=minio_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
Gates.__init__(
self,
minio_repository=minio_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
OPC.__init__(
self,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
ModelMetrics.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
API.__init__(
self,
base_url=pi_web_api_config['base_url'],
auth_type=pi_web_api_config['auth_type'],
auth_token=pi_web_api_config['auth_token'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
mongo_repository = (
MongoDBRepository(
connection_string=mongo_config['connection_string'],
database_name=mongo_config['database_name'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
if mongo_config
else None
)
ModelImport.__init__(
self,
status_config=import_status_config or build_import_status_config(),
import_config=import_config or build_import_config(),
minio_repository=minio_repository,
# The MLflow repository is shared rather than rebuilt: setting the tracking URI is a
# process-wide side effect, so two instances would be two chances to disagree. Read
# through `getattr` because `MLFlow.__init__` is what creates it — if that ever stops
# running before this line, the import activities refuse with "MLflow repository not
# initialized" at the point of use rather than silently provisioning nothing.
mlflow_repository=getattr(self, 'model_monitoring_repository', None),
mongo_repository=mongo_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
async def shutdown(self):
"""
Gracefully shutdown all activities and clean up resources.
This method ensures proper cleanup of all resources including:
- PostgreSQL connection pools
- OPC server connections
- PI Web API client connections
- MLFlow model repositories
- Any other resources that need explicit cleanup
The method should be called before the application terminates to ensure
proper resource cleanup and prevent resource leaks.
"""
Storage.close(self)
MLFlow.close(self)
Gates.close(self)
await OPC.aclose(self)
ModelMetrics.close(self)
API.close(self)
ModelImport.close(self)

305
laborious/activities/api.py Normal file
View File

@@ -0,0 +1,305 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import json
import traceback
from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.pi_web_api_client import PIWebAPIClient
from laborious import metrics
PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13
class API(SientiaMonitoring):
"""
PI Web API operations for writing prediction data to PI Web API.
This class provides Temporal activities for interacting with the PI Web API
to write prediction and confidence values to industrial systems. It handles
error scenarios gracefully by setting error confidence values and sending
notifications when write operations fail.
The class implements comprehensive error handling for both prediction and
confidence value writes, ensuring that partial failures are properly
reported and handled.
"""
def __init__(
self,
base_url: str,
auth_type: str,
auth_token: str,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
) -> None:
"""
Initialize API activity with PI Web API client.
Args:
base_url (str): Base URL of the PI Web API server
auth_type (str): Authentication type ('basic' or 'bearer')
auth_token (str): Authentication token
logger (Logger): Logger instance for operation logging
notification_handler (NotificationHandler): Handler for system notifications
metrics_controller (MetricsController): Controller for metrics collection
"""
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.pi_web_api_client = PIWebAPIClient(
base_url=base_url,
auth_config={
'type': auth_type,
'token': auth_token,
},
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
headers_config={
'Content-Type': 'application/json',
'Accept': 'application/json',
'x-requested-with': 'piwebapistreams',
'User-Agent': 'Aig-Laborious-Agent/1.0',
},
)
def get_pi_web_api_core_labels(
self,
metadata: dict[str, Any],
operation_type: str = 'write_pi_web_api_data',
) -> dict[str, Any]:
"""
Generate core labels for PI Web API metrics.
PI Web API metrics in laborious use the shared ``CORE_LABELS`` from
``sientia_do``, which includes ``operation_type``. For this reason,
operation_type must always be present in emitted labels.
Args:
- metadata (dict[str, Any]): Workflow execution metadata used to derive labels.
- operation_type (str): Operation type label for metric cardinality.
Return:
dict[str, Any]: Core labels dictionary including operation_type.
"""
return super().get_core_labels(
metadata=metadata,
operation_type=operation_type,
)
def close(self) -> None:
"""
Close the PI Web API client and shutdown monitoring services.
This method properly closes all connections and resources associated
with the PI Web API client and monitoring services.
"""
self.pi_web_api_client.close()
SientiaMonitoring.shutdown(self)
async def process_pi_web_api_response(
self,
response_data: list[dict[str, Any]],
tags: dict[str, str],
core_labels: dict[str, str],
metadata: dict[str, Any],
) -> tuple[int, str]:
"""
Process the response data from PI Web API write operation.
Validates that all tags were successfully written, emits metrics for each tag
(success or error), and returns the appropriate prediction confidence value.
Sets error confidence if any tag write fails or if the number of written tags
doesn't match the expected count.
Args:
- response_data (dict[str, Any]): The response data from the PI Web API write operation.
- tags (dict[str, str]): The tags that were written to the PI Web API.
- core_labels (dict[str, str]): The core labels of the workflow execution.
- metadata (dict[str, Any]): The metadata of the workflow execution.
Returns:
int: Prediction confidence value (0 for success, 13 for errors)
"""
# Convert tags from name:webid to webid:name
tags = {w: t for t, w in tags.items()}
tag_names = list[str](tags.values())
confidence = 0
message = ''
# Evaluate response for each tag
written_tags = []
for item in response_data:
web_id = item.get('WebId')
if not web_id:
self.error('The response did not contain some WebIds', metadata)
continue
errors = item.get('Errors', [])
tag_name = tags.get(web_id)
if not tag_name:
self.error(
f'The response did not contain the tag name for WebId {web_id}', metadata
)
continue
if errors:
self.error(
f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata
)
await self.emit_metric(
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT,
tags={
**core_labels,
'tag_name': tag_name,
},
)
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
else:
await self.emit_metric(
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT,
tags={
**core_labels,
'tag_name': tag_name,
},
)
written_tags.append(tag_name)
if len(written_tags) != len(tag_names):
message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.'
self.error(
f'{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
metadata,
)
await self.send_notification_async(
metadata=metadata,
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
)
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
return confidence, message
@activity.defn(name='write_pi_web_api_data')
async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Write prediction and confidence data to PI Web API.
Writes prediction values and confidence scores to PI Web API using configured
web IDs. Processes responses to validate writes and emit metrics. Handles errors
gracefully by setting error confidence values when writes fail and sending
notifications for both prediction and confidence write errors.
Args:
input_data (dict[str, Any]): The input data containing:
- metadata (dict[str, Any]): Workflow execution metadata
- pi_web_api_output_config (dict[str, Any]): PI Web API configuration with:
- endpoint (str): PI Web API endpoint URL
- prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions
- confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence
- data (dict[str, Any]): Prediction data, its a dataframe converted to dict.
Returns:
dict[Any, Any]: Data dictionary with potentially modified confidence values
If prediction write fails, prediction_confidence is set to error value (13)
"""
metadata = input_data['metadata']
data = DataFrame(input_data['data'])
pi_web_api_output_config = input_data['pi_web_api_output_config']
self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata)
raw_prediction_tags = pi_web_api_output_config['prediction_tags']
raw_confidence_tags = pi_web_api_output_config['confidence_tags']
prediction_tags = list[str](raw_prediction_tags.values())
confidence_tags = list(raw_confidence_tags.values())
core_labels = self.get_pi_web_api_core_labels(metadata)
prediction_value = data.head(1)['prediction'].values[0]
confidence_value = data.head(1)['prediction_confidence'].values[0]
try:
prediction_response = await self.pi_web_api_client.write_value(
web_ids=prediction_tags,
value={
'Timestamp': data.head(1)['timestamp'].values[0],
'Value': prediction_value,
},
metadata=metadata,
)
confidence, message = await self.process_pi_web_api_response(
response_data=prediction_response,
tags=raw_prediction_tags,
core_labels=core_labels,
metadata=metadata,
)
# Preserve incoming confidence/comments on successful PI writes.
# Only downgrade confidence or override comments when PI response
# explicitly reports a problem (e.g. partial write mismatch).
if confidence != 0:
data['prediction_confidence'] = confidence
if message:
data['comments'] = message
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata)
data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
data['comments'] = str(e)
return data.to_dict()
try:
confidence_response = await self.pi_web_api_client.write_value(
web_ids=confidence_tags,
value={
'Timestamp': data.head(1)['timestamp'].values[0],
'Value': float(confidence_value),
},
metadata=metadata,
)
await self.process_pi_web_api_response(
response_data=confidence_response,
tags=raw_confidence_tags,
core_labels=core_labels,
metadata=metadata,
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
block='write_pi_web_api_data',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
return data.to_dict()

View File

@@ -0,0 +1,804 @@
from sientia_do.repository.minio_repository import MinioRepository
from temporalio import activity, workflow
from laborious.utils.repository.minio_manager import MinioManager
with workflow.unsafe.imports_passed_through():
import traceback
from collections.abc import Callable, Mapping
from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.utils.formatters import create_sample_dict
from laborious import metrics
from laborious.utils.dataframe_debug import build_dataframe_debug_message
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
)
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
# Strongly-typed filter function signatures
InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
# Input filter function mappings
input_filter_functions: dict[str, InputFilterFunc] = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data,
}
# Confidence mappings kept separate from function maps to avoid Union types
input_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 2,
'REPEAT': -1,
}
# MLFlow response filter function mappings
mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
'API_ERROR': api_error_filter,
}
mlflow_response_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 10,
'REPEAT': -1,
}
# MLFlow content filter function mappings
mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
'NAN_VALUES': nan_values_filter,
'EMPTY_DATA': filter_empty_data,
}
mlflow_content_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 18,
'REPEAT': -1,
}
class Gates(MinioManager):
"""
Data quality gates and filtering activities for the Laborious system.
This class implements comprehensive data quality validation and filtering
mechanisms that can be applied at different stages of the prediction pipeline.
It provides configurable filters with policy-based decision making to ensure
data integrity and quality throughout the ML workflow.
The class supports multiple filter types and implements a flexible policy
system that can be configured for different validation requirements. Each
filter returns a path decision (STOP, CONTINUE, REPEAT) along with confidence
scores and detailed comments for monitoring and debugging.
Attributes:
input_filter_functions (dict): Mapping of input filter names to functions
mlflow_response_filter_functions (dict): Mapping of MLFlow response filter names to functions
mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
"""
minio_repository: MinioRepository | None = None
_MAX_DEBUG_DATAFRAME_ROWS = 100
def __init__(
self,
minio_repository: MinioRepository | None = None,
logger: Logger | None = None,
notification_handler: NotificationHandler | None = None,
metrics_controller: MetricsController | None = None,
):
"""
Initialize data quality gates with logging and notification capabilities.
Args:
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If BaseActivity initialization fails
"""
MinioManager.__init__(
self, minio_repository, logger, notification_handler, metrics_controller
)
def close(self) -> None:
"""
Close the gates activity and clean up resources.
"""
MinioManager.close(self)
def __del__(self):
self.close()
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
"""
Log dataframe content only when row count is below the configured threshold
Args:
- message (str): Base log message to identify the dataframe in logs
- data (Any): Dataframe-like payload to be logged
- metadata (dict[str, Any]): Workflow metadata for contextual logging
"""
self.debug(
build_dataframe_debug_message(
message=message,
data=data,
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
),
metadata,
)
@staticmethod
def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]:
"""
Read filter policy/config keys in a case-insensitive way.
Args:
config (dict[str, Any]): Filter configuration dictionary.
Return:
tuple[str, dict[str, Any]]: Parsed policy and config payload.
"""
normalized = {str(key).upper(): value for key, value in config.items()}
policy = normalized['POLICY']
filter_config = normalized.get('CONFIG', {})
return policy, filter_config
@activity.defn(name='input_gate')
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Apply input data quality filters and validation.
This activity validates input data quality using configurable filters
before proceeding with ML operations. It applies multiple filter types
and returns a path decision based on the filter results and configured
policies.
The method implements a comprehensive filtering system that:
1. Applies configured filters to input data
2. Evaluates filter results against policy configurations
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
4. Provides confidence scores and detailed comments
5. Handles errors gracefully with notification integration
Args:
input_data: Configuration and data for input validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Filter configuration and policies
- data (dict): Input data to validate
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If filter execution fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing input gate...', metadata)
filters = input_data['filters']
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
path_priority = input_data['path_priority']
filter_output = []
self._debug_dataframe('Input data:', data, metadata)
self.debug(f'Filters: {filters}', metadata)
# Apply each configured filter
for fil, config in filters.items():
if fil not in input_filter_functions:
self.error(f'Filter {fil} not found', metadata)
continue
policy, filter_config = self._read_filter_entry(config)
try:
if input_filter_functions[fil](data, filter_config):
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
filter_output.append(policy)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id=f'INTPUT_GATE_ERROR__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='input_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Input gate result: {path_flag}', metadata)
return path_flag, input_path_confidence[path_flag], 'Input data with bad quality'
self.info('Nothing was filtered by the input gate', metadata)
del data
return None, 0, ''
@activity.defn(name='mlflow_response_gate')
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow API response quality and integrity.
This activity validates MLFlow API responses to ensure they meet quality
standards before proceeding with further processing. It applies response-specific
filters and determines appropriate path decisions based on response quality.
The method implements response validation that:
1. Applies MLFlow response-specific filters
2. Evaluates API response quality and integrity
3. Determines path decisions based on response validation results
4. Provides confidence scores and detailed validation comments
5. Handles API errors and response validation failures
Args:
input_data: Configuration and data for response validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Response filter configuration and policies
- data (dict): MLFlow API response data to validate
- type (str): Type of MLFlow operation (transform, predict)
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If response validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing mlflow response gate...', metadata)
raw_data = input_data['data']
filters = input_data['filters']
self.debug(
f'Input data: \n {create_sample_dict(raw_data, max_items=5, max_depth=5)}', metadata
)
self.debug(f'Filters: {filters}', metadata)
payload = MinioDataFramePayload.from_dict(raw_data)
data = await payload.retrieve(self.minio_repository, metadata)
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
comments = []
status = payload.status or {}
for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
continue
policy, filter_config = self._read_filter_entry(config)
try:
if mlflow_response_filter_functions[fil](status, filter_config):
filter_output.append(policy)
comments.append(status.get('message', 'Unknown MLFlow API error'))
await self.send_notification_async(
metadata=metadata,
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
message=status.get('message', 'Unknown MLFlow API error'),
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=status.get('traceback'),
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Mlflow response gate result: {path_flag}', metadata)
return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments)
self.info('Nothing was filtered by the mlflow response gate', metadata)
del data
return None, 0, ''
@activity.defn(name='mlflow_content_gate')
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow prediction content quality and integrity.
This activity validates the content of MLFlow predictions to ensure they
meet quality standards before export and persistence. It applies content-specific
filters and determines appropriate path decisions based on content quality.
The method implements content validation that:
1. Applies MLFlow content-specific filters
2. Evaluates prediction content quality and integrity
3. Determines path decisions based on content validation results
4. Provides confidence scores and detailed validation comments
5. Handles content validation failures and quality issues
Args:
input_data: Configuration and data for content validation
Required keys:
- metadata (dict): Workflow execution metadata
- filters (dict): Content filter configuration and policies
- data (dict): MLFlow prediction content to validate
- type (str): Type of MLFlow operation (transform, predict)
- path_priority (list[str]): Priority order for path decisions
Returns:
tuple: (path_flag, confidence, comment)
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
- confidence (int): Confidence score for the decision
- comment (str): Detailed explanation of the decision
Raises:
Exception: If content validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters']
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self._debug_dataframe('Input data:', data, metadata)
self.debug(f'Filters: \n {filters}', metadata)
for fil, config in filters.items():
if fil not in mlflow_content_filter_functions:
continue
policy, filter_config = self._read_filter_entry(config)
try:
if mlflow_content_filter_functions[fil](data, filter_config):
filter_output.append(policy)
await self.send_notification_async(
metadata=metadata,
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
message=f'Data not passed the content filter {fil}:{config}',
block='mlflow_gate',
level=NotificationLevel.WARNING,
attachment_content=data.to_string(),
)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
message=f'Error in filter {fil}:{config}: \n {e}',
block='mlflow_gate',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
for path_flag in path_priority:
if path_flag in filter_output:
self.info(f'Mlflow content gate result: {path_flag}', metadata)
return (
path_flag,
mlflow_content_path_confidence[path_flag],
'Transformed data not passed the content filter',
)
self.info('Nothing was filtered by the mlflow content gate', metadata)
del data
return None, 0, ''
def get_prediction_store_policy(
self, prediction_store_policy: str, metadata: dict[str, Any]
) -> tuple[str, int]:
"""
Parse and validate prediction store policy configuration.
This method parses prediction store policy strings in the format 'type:value'
and validates them against allowed policy types and values. It provides
sensible defaults for invalid configurations and logs policy validation
failures for operational monitoring.
Supported Policy Types:
- 'lts': Latest timestamp - sorts data by timestamp descending
- 'erl': Earliest timestamp - sorts data by timestamp ascending
Args:
prediction_store_policy (str): Policy string in format 'type:value'
metadata (dict[str, Any]): Context metadata for logging and notifications
Returns:
tuple[str, int]: (policy_type, policy_value)
- policy_type (str): Validated policy type ('lts' or 'erl')
- policy_value (int): Number of rows to retain
"""
policy_elements = prediction_store_policy.split(':')
if len(policy_elements) < 2:
self.error(
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
policy_type = policy_elements[0]
policy_value = policy_elements[1]
# If the policy_type is not lts or erl, we use the default policy
# If the policty_value is not a number or 0, we use the default policy
if (
policy_type not in ['lts', 'erl']
or not policy_value.isdigit()
or int(policy_value) == 0
):
self.error(
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
return policy_type, int(policy_value)
@activity.defn(name='format_transformed_data')
async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
"""
Format transformed data for storage and export operations.
This method formats transformed data from MLFlow model transformations
into a standardized format suitable for database storage. It converts
wide-format data (columns as variables) into long-format (melted)
with proper timestamp handling and model identification.
The formatting process includes:
1. Converting input data dictionary to DataFrame
2. Extracting timestamps from DataFrame index
3. Resetting index to create sequential row numbers
4. Melting data from wide format to long format (variable-value pairs)
5. Adding model_id for data lineage tracking
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict)
- model_id (str): Unique identifier for the ML model
Returns:
dict: Formatted data dictionary with keys:
- timestamp (dict): Timestamp values indexed by row number
- variable (dict): Variable names indexed by row number
- value (dict): Variable values indexed by row number
- model_id (dict): Model identifiers indexed by row number
"""
metadata = input_data['metadata']
model_id = input_data['model_id']
self.info('Formatting transformed data...', metadata)
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
data['timestamp'] = data.index
data = data.reset_index(drop=True)
data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
data['model_id'] = model_id
return await MinioDataFramePayload.from_dataframe(
dataframe=data,
minio_repo=self.minio_repository,
model_name=input_data['model_name'],
operation='transform',
workflow_metadata=metadata,
last_timestamp=payload.last_timestamp,
logger=self.logger,
)
@activity.defn(name='format_prediction')
async def format_prediction(self, input_data: dict[str, Any]) -> dict:
"""
Format prediction data according to configured storage policies.
This method formats prediction data for storage and export operations.
It applies timestamp-based sorting policies, adds metadata fields,
and ensures data consistency before persistence. The method supports
multiple storage policies for flexible data retention strategies.
If only one row is present, we use the last timestamp as the timestamp
Storage Policies:
- 'lts:N': Latest timestamp - retains N most recent predictions
- 'erl:N': Earliest timestamp - retains N oldest predictions
Args:
input_data (dict): Input data containing:
- data (dict[str, Any]): Raw prediction data to format
- timestamp (str): Timestamp of the data
- model_id (str): Unique identifier for the ML model
- prediction_confidence (float): Confidence score for the prediction
- prediction_store_policy (str): Storage policy in format 'type:value'
Returns:
dict: Formatted prediction data ready for storage and export
"""
metadata = input_data['metadata']
last_timestamp = input_data['timestamp']
prediction_store_policy = input_data['prediction_store_policy']
self.info('Formatting prediction...', metadata)
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
# Create timestamp column from index and reset index
data['timestamp'] = data.index
data = data.reset_index(drop=True)
self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
self._debug_dataframe('Prediction data:', data, metadata)
policy_type, policy_value = self.get_prediction_store_policy(
prediction_store_policy, metadata
)
# If data has no timestamp, we use the default timestamp and not sort the data
self.info(
f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
)
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
if policy_type == 'lts':
self.debug('Sorting data by timestamp descending', metadata)
data = data.sort_values(by='timestamp', ascending=False)
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
elif policy_type == 'erl':
self.debug('Sorting data by timestamp ascending', metadata)
data = data.sort_values(by='timestamp', ascending=True)
else:
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
raise ValueError(f'Invalid policy type: {policy_type}')
int_policy_value = int(policy_value)
data = data.head(int_policy_value)
if int_policy_value == 1:
data['timestamp'] = last_timestamp
data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good'
data['comments'] = ''
data = data.sort_values(by='timestamp', ascending=False)
data = data.reset_index(drop=True)
self.info(f'Prediction formatted: {len(data)} rows', metadata)
self._debug_dataframe('Prediction data:', data, metadata)
return data.to_dict()
@activity.defn(name='format_default_prediction')
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
"""
Create and format default prediction data for error conditions.
This method generates default prediction data when the main prediction
pipeline encounters errors or quality issues. It creates a standardized
data structure with zero values for predictions and useful metadata
for operational monitoring and debugging.
The default prediction serves as a fallback mechanism to:
1. Maintain data pipeline continuity during failures
2. Provide operational visibility into prediction quality issues
3. Enable downstream systems to handle error conditions gracefully
4. Support debugging and troubleshooting efforts
Args:
input_data (dict): Input data containing:
- timestamp (str): Timestamp for the default prediction
- model_id (str): Unique identifier for the ML model
- prediction_confidence (float): Confidence score (typically low for errors)
- comment (str): Error description or operational comment
Returns:
dict: Formatted default prediction data with error indicators
"""
metadata = input_data['metadata']
self.debug('Formatting default prediction...', metadata)
data = DataFrame(
{
'prediction': [0],
'response_time': [0],
'timestamp': [input_data['timestamp']],
'model_id': [input_data['model_id']],
'prediction_confidence': [input_data['prediction_confidence']],
'prediction_status': ['Bad'],
'comments': [input_data['comment']],
}
)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
@activity.defn(name='format_retrain_report')
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
"""
Format retrain report data for storage and audit trail maintenance.
This method formats model retraining operation results into a standardized
report format suitable for database storage and operational monitoring.
It captures retraining status, timestamps, and model version information
for comprehensive audit trails and operational visibility.
The formatting process includes:
1. Extracting retraining experiment response data
2. Capturing model update report information (version, MLflow IDs)
3. Formatting timestamps and status information
4. Conditionally including version information for successful retrains
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- experiment_response (dict): Retraining experiment response containing:
- success (bool): Retraining operation success status
- timestamp (str): Timestamp of the retraining operation
- message (str): Status message or error description
- update_report (dict): Model update report containing:
- version (str): New model version identifier
- mlflow_run_id (str): MLflow run identifier
- mlflow_experiment_id (str): MLflow experiment identifier
- model_id (str): Unique identifier for the ML model
- model_name (str): Name of the ML model
Returns:
dict: Formatted retrain report dictionary with keys:
- model_id (dict): Model identifiers indexed by row number
- model_name (dict): Model names indexed by row number
- timestamp (dict): Retraining timestamps indexed by row number
- status (dict): Retraining status messages indexed by row number
- version (dict, optional): Model versions indexed by row number
Only included if experiment_response['success'] is True
- mlflow_run_id (dict, optional): MLflow run IDs indexed by row number
Only included if experiment_response['success'] is True
- mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number
Only included if experiment_response['success'] is True
"""
metadata = input_data['metadata']
self.info('Formatting retrain report...', metadata)
experiment_response = input_data['experiment_response']
update_report = input_data['update_report']
model_id = input_data['model_id']
model_name = input_data['model_name']
report = DataFrame(
{
'model_id': [model_id],
'model_name': [model_name],
'timestamp': [experiment_response['timestamp']],
'status': [experiment_response['message']],
}
)
if experiment_response['success']:
# Retrain was successfull
report['version'] = update_report['version']
report['mlflow_run_id'] = update_report['mlflow_run_id']
report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
self._debug_dataframe('Retrain report:', report, metadata)
return report.to_dict()
@activity.defn(name='write_metrics')
async def write_metrics(self, input_data: dict[str, Any]):
"""
Write prediction performance metrics to Prometheus monitoring system.
This method records comprehensive metrics for prediction operations,
enabling operational monitoring, performance analysis, and alerting.
It tracks prediction counts, confidence levels, and response times
for each model and pipeline combination.
Metrics Recorded:
1. Prediction Count: Incremental counter for successful predictions
2. Confidence Monitor: Current confidence level for predictions
3. Response Time Monitor: Histogram of prediction response times
Args:
input_data (dict): Input data containing:
- metadata (dict[str, Any]): Workflow execution metadata
- prediction (dict[str, Any]): Prediction data with metrics
Raises:
Exception: If metrics writing fails or configuration is invalid
"""
metadata = input_data['metadata']
prediction = DataFrame(input_data['prediction'])
prediction_confidence = prediction['prediction_confidence'].values[0]
response_time = prediction['response_time'].values[0]
opc_metrics = input_data['opc_metrics']
self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
core_tags = {
'pod_id': self.pod_id,
'runtime': self.runtime,
'operation_type': 'predict',
'model_name': metadata['model_name'],
'workflow_name': metadata['workflow_name'],
}
await self.emit_metric(
metric_object=metrics.PREDICTIONS_WRITTEN_COUNT,
tags=core_tags,
)
await self.emit_metric(
metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
method='set',
tags=core_tags,
value=prediction_confidence,
)
await self.emit_metric(
metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
method='observe',
tags=core_tags,
value=response_time,
)
for server_id, tags in opc_metrics.items():
for tag, response_time in tags.items():
if response_time is not None:
await self.emit_metric(
metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
method='observe',
tags={
**core_tags,
'opc_server_id': server_id,
'tag': tag,
},
value=response_time,
)
await self.emit_metric(
metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
tags={
**core_tags,
'opc_server_id': server_id,
'tag': tag,
},
)
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

View File

@@ -0,0 +1,528 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from typing import Any
import numpy as np
from pandas import to_datetime
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.repository.minio_repository import MinioRepository
from sientia_do.temporal.constants import (
DATETIME_FORMAT,
DATETIME_FORMAT_MS_WITH_TZ,
DATETIME_FORMAT_WITH_TZ,
now,
)
from sientia_do.utils.formatters import create_sample_dict
from laborious.utils.dataframe_debug import build_dataframe_debug_message
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
from laborious.utils.repository.minio_manager import MinioManager
from laborious.utils.repository.model_repository import MLFlowRepository
class MLFlow(MinioManager):
"""
MLFlow integration activities for model inference operations.
This class provides activities for interacting with MLFlow models, including
data transformation and prediction operations. It handles authentication,
data preprocessing, and model management with configurable retention policies.
The class implements comprehensive error handling and logging for all
MLFlow operations, ensuring reliable model inference in production environments.
Attributes:
mlflow_host (str): MLFlow server hostname
mlflow_port (int): MLFlow server port
mlflow_username (str): MLFlow authentication username
mlflow_password (str): MLFlow authentication password
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
"""
_MAX_DEBUG_DATAFRAME_ROWS = 100
def __init__(
self,
mlflow_host: str,
mlflow_port: int,
mlflow_username: str,
mlflow_password: str,
minio_repository: MinioRepository | None = None,
logger: Logger | None = None,
notification_handler: NotificationHandler | None = None,
metrics_controller: MetricsController | None = None,
):
"""
Initialize MLFlow activities with server configuration.
Args:
mlflow_host: MLFlow server hostname or IP address
mlflow_port: MLFlow server port number
mlflow_username: Username for MLFlow authentication
mlflow_password: Password for MLFlow authentication
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
Raises:
Exception: If MLFlowRepository initialization fails
"""
MinioManager.__init__(
self, minio_repository, logger, notification_handler, metrics_controller
)
self.mlflow_host = mlflow_host
self.mlflow_port = mlflow_port
self.mlflow_username = mlflow_username
self.mlflow_password = mlflow_password
self.model_monitoring_repository = MLFlowRepository(
f'{mlflow_host}:{mlflow_port}',
mlflow_username,
mlflow_password,
logger,
notification_handler,
metrics_controller,
)
def close(self) -> None:
"""
Close the MLFlow activity and clean up resources.
"""
MinioManager.close(self)
def __del__(self):
self.close()
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
"""
Log dataframe content only when row count is below the configured threshold
Args:
- message (str): Base log message to identify the dataframe in logs
- data (Any): Dataframe-like object expected to expose shape and to_csv
- metadata (dict[str, Any]): Workflow metadata for contextual logging
"""
self.debug(
build_dataframe_debug_message(
message=message,
data=data,
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
),
metadata,
)
@activity.defn(name='request_transform')
async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
"""
Transform input data using MLFlow models.
This activity processes input data through MLFlow model transformation,
including data preprocessing, format conversion, and validation. It handles
data deduplication, pivoting, and cleanup to ensure optimal model performance.
The transformation process includes:
1. Data deduplication based on variable and timestamp
2. Data pivoting for model input format
3. Null value handling and cleanup
4. MLFlow model transformation request
5. Response validation and logging
Args:
input_data: Configuration and data for transformation
Required keys:
- metadata (dict): Workflow execution metadata
- data (dict): Input data for transformation
- model_name (str): Name of the MLFlow model to use
- model_retention (int): Model retention period in minutes
Returns:
dict: Transformed data from MLFlow model
Raises:
Exception: If transformation fails or MLFlow model is unavailable
"""
metadata = input_data['metadata']
self.info('Transforming data...', metadata)
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self._debug_dataframe('Raw input data:', data, metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
data = data.sort_values('created_at', ascending=False).drop_duplicates(
subset=['variable', 'timestamp'], keep='first'
)
# Pivot data for model input format
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
data.columns.name = None
data.index.name = None
data['timestamp'] = data.index
self._debug_dataframe('Processed input data:', data, metadata)
# Request transformation from MLFlow model
response_data = await self.model_monitoring_repository.transform(
model_name, data, model_config, metadata
)
self.debug(
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.debug(
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info('Data transformed successfully', metadata)
if not response_data.get('success', False):
return await MinioDataFramePayload.from_dataframe(
dataframe=None,
minio_repo=self.minio_repository,
model_name=model_name,
operation='transform',
status=response_data,
workflow_metadata=metadata,
last_timestamp=payload.last_timestamp,
logger=self.logger,
)
return await MinioDataFramePayload.from_dataframe(
dataframe=response_data['content'],
minio_repo=self.minio_repository,
model_name=model_name,
operation='transform',
workflow_metadata=metadata,
status={
'success': True,
},
last_timestamp=payload.last_timestamp,
logger=self.logger,
)
@activity.defn(name='request_predict')
async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
"""
Execute predictions using MLFlow models.
This activity performs ML model inference using MLFlow models with the
transformed data. It handles data format conversion, null value processing,
and model prediction requests with comprehensive error handling.
The prediction process includes:
1. Data format validation and cleanup
2. Null value handling for model compatibility
3. MLFlow model prediction request
4. Response validation and logging
5. Performance monitoring and metrics
Args:
input_data: Configuration and data for prediction
Required keys:
- metadata (dict): Workflow execution metadata
- data (dict): Transformed data for prediction
- model_name (str): Name of the MLFlow model to use
- model_retention (int): Model retention period in minutes
Returns:
dict: Prediction results from MLFlow model
Raises:
Exception: If prediction fails or MLFlow model is unavailable
"""
metadata = input_data['metadata']
self.info('Predicting data...', metadata)
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self._debug_dataframe('Input data for prediction:', data, metadata)
# Convert numpy.nan to None for model compatibility
data.replace(np.nan, None, inplace=True)
data['timestamp'] = data.index
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
# Request prediction from MLFlow model
response_data = await self.model_monitoring_repository.predict(
model_name, data, model_config, metadata
)
self.debug(
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info('Data predicted successfully', metadata)
if not response_data.get('success', False):
return await MinioDataFramePayload.from_dataframe(
dataframe=None,
minio_repo=self.minio_repository,
model_name=model_name,
operation='predict',
status=response_data,
workflow_metadata=metadata,
last_timestamp=payload.last_timestamp,
logger=self.logger,
)
return await MinioDataFramePayload.from_dataframe(
dataframe=response_data['content'],
minio_repo=self.minio_repository,
model_name=model_name,
operation='predict',
workflow_metadata=metadata,
status={
'success': True,
},
last_timestamp=payload.last_timestamp,
logger=self.logger,
)
@activity.defn(name='retrain_model')
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Retrain MLFlow models with updated training data.
This activity orchestrates the complete model retraining process,
including data preparation, model retraining execution, and result
validation. It handles data preprocessing, column cleanup, and
comprehensive error handling for production model management.
The retraining process includes:
1. Data timestamp extraction and validation
2. Column cleanup and data preparation
3. Data pivoting for model input format
4. MLFlow model retraining execution
5. Result validation and error handling
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- data (dict[str, Any]): Training data for model retraining
- model_name (str): Name of the MLFlow model to retrain
Returns:
dict: Retraining results containing:
- status (str): Retraining operation status
- timestamp (str): Timestamp of the retraining operation
- experiment (str): MLFlow experiment identifier
Raises:
Exception: If retraining fails or encounters critical errors
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
metadata = input_data['metadata']
try:
# Payload-based retrain input (inline dict or MinIO offloaded).
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='ERROR_LOADING_RETRAIN_DATA',
message=f'Error loading retrain data: {e}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata)
return {
'success': False,
'message': f'Error loading retrain data: {e}',
'traceback': trace,
'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
}
self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
self.info(f'Retraining model {model_name}...', metadata)
timestamp = data['timestamp'].max()
self.debug(f'Timestamp: {timestamp}', metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
if 'created_at' in data.columns:
data = data.sort_values('created_at', ascending=False).drop_duplicates(
subset=['variable', 'timestamp'], keep='first'
)
else:
data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')
data.drop(columns=['model_id'], inplace=True, errors='ignore')
data.drop(columns=['created_at'], inplace=True, errors='ignore')
# Pivot data for model input format
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True)
data.columns.name = None
data['timestamp'] = data.index
data['timestamp'] = to_datetime(
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
).dt.strftime(DATETIME_FORMAT)
data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
data.columns.name = None
retrain_output = await self.model_monitoring_repository.retrain_model(
data=data, model_name=model_name, model_config=model_config, metadata=metadata
)
if not retrain_output['success']:
trace = retrain_output['traceback']
await self.send_notification_async(
metadata=metadata,
notification_id='RETRAIN_MODEL_ERROR',
message=f'Error retraining model {model_name}: {retrain_output["message"]}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata=metadata)
return {**retrain_output, 'timestamp': timestamp}
@activity.defn(name='update_production_model')
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Update production model with newly trained model version.
This activity manages the critical process of updating production
models with newly trained versions. It handles model deployment,
status tracking, and comprehensive reporting for operational
visibility and audit trails.
The update process includes:
1. Production model update execution
2. Status and metadata tracking
3. Comprehensive reporting and logging
4. Error handling and notification
5. Audit trail maintenance
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- model_name (str): Name of the MLFlow model to update
- experiment (str): MLFlow experiment identifier
- model_id (str): Unique identifier for the model version
- timestamp (str): Timestamp of the update operation
- status (str): Current status of the model update
Returns:
dict[Any, Any]: Comprehensive update report containing:
- model_id (str): Model version identifier
- model_name (str): Name of the updated model
- timestamp (str): Update operation timestamp
- status (str): Update operation status
- Additional MLFlow response metadata
Raises:
Exception: If production model update fails
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
experiment = input_data['experiment']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata
)
try:
response = await self.model_monitoring_repository.update_production_model(
experiment=experiment, model_name=model_name, metadata=metadata
)
self.info(f'Production model {model_name} updated successfully', metadata)
return response
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
message=f'Error updating production model {model_name}: {e}',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
@activity.defn(name='get_reference_data')
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
"""
Get reference data from the MLflow Model Registry.
This method retrieves evaluation reference data stored as artifacts in the
MLflow Model Registry. The reference data is typically used for model
drift detection, performance comparison, and quality validation. The method
loads the data from a CSV artifact file and formats timestamps for
consistent processing.
The method handles:
1. Loading evaluation data artifact from MLflow Model Registry
2. Timestamp parsing and formatting for consistency
3. Data conversion to dictionary format for workflow consumption
4. Graceful handling of missing reference data
Args:
input_data (dict): Input data containing:
- metadata (dict): Workflow execution metadata
- model_name (str): Name of the MLFlow model to get reference data from
Returns:
list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
as a list of dictionaries. Returns None if reference data is not found
or if the artifact does not exist.
Raises:
Exception: If artifact loading fails or encounters errors during processing
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
artifact = 'evaluation_data.csv'
reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
model_name=model_name, artifact_path=artifact, metadata=metadata
)
if reference_data is None:
self.warning(f'Reference data not found for model {model_name}', metadata)
return None
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
return reference_data.to_dict(orient='records')

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,348 @@
"""The import's failure vocabulary: a closed catalog, typed errors and one reporting channel.
Three things live here and nowhere else:
- **the catalog** — one stable `error_code` and one plain-language `error_reason` per pipeline step.
The reason is written for a person who does not know this system exists: what happened, what state
the model is in, what to do next. It never carries a path, a table name, an object key, a Python
exception class or a traceback, and it is never built from `str(exception)`.
- **the typed errors** — the two failures that must not be retried, because repeating them cannot
change the answer: a record this workflow does not own, and a statement the database refuses.
- **the channel** — `report_import_failure`, the single funnel every failure passes through. It is
the only place in the import path that logs an error or sends a notification, which a test
asserts, so no failure is reported twice or half-reported.
"""
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import traceback
from collections.abc import Mapping
from dataclasses import dataclass
from enum import StrEnum
from types import MappingProxyType
from typing import Any, Final, Protocol
from temporalio.exceptions import ApplicationError
from laborious import metrics
from laborious.utils.bundle.steps import ImportStep
class ImportErrorCode(StrEnum):
"""The machine-readable half of a failure: the frontend's i18n key and support's triage key.
Written under `import.code` in the record's JSON detail — not in a column of its own, so the
import adds nothing to a table it shares with the training flow.
"""
REQUEST_INVALID = 'IMPORT_REQUEST_INVALID'
FILE_UNREADABLE = 'IMPORT_FILE_UNREADABLE'
BUNDLE_CANNOT_BE_OPENED = 'IMPORT_BUNDLE_CANNOT_BE_OPENED'
BUNDLE_REJECTED = 'IMPORT_BUNDLE_REJECTED'
BUNDLE_UNPACK_FAILED = 'IMPORT_BUNDLE_UNPACK_FAILED'
BUNDLE_INCOMPLETE = 'IMPORT_BUNDLE_INCOMPLETE'
BUNDLE_UNEXPECTED_CONTENT = 'IMPORT_BUNDLE_UNEXPECTED_CONTENT'
STORAGE_PREPARATION_FAILED = 'IMPORT_STORAGE_PREPARATION_FAILED'
MODEL_FILES_NOT_STORED = 'IMPORT_MODEL_FILES_NOT_STORED'
MODEL_NOT_REGISTERED = 'IMPORT_MODEL_NOT_REGISTERED'
MODEL_SETTINGS_NOT_SAVED = 'IMPORT_MODEL_SETTINGS_NOT_SAVED'
UNEXPECTED_ERROR = 'IMPORT_UNEXPECTED_ERROR'
# Codes that never reach a record's `error_message`, and are therefore absent from the catalog:
# one names the failure to own a row, one the failure to write one, one the temporary files left
# on the worker after an import that otherwise succeeded. All three live on the channel — log,
# metric, notification, Temporal error — and the third is additionally flagged in the record's
# detail as `cleanup_failed`, without a code and without changing the status.
LOG_ROW_NOT_CLAIMABLE = 'IMPORT_LOG_ROW_NOT_CLAIMABLE'
STATUS_NOT_RECORDED = 'IMPORT_STATUS_NOT_RECORDED'
CLEANUP_INCOMPLETE = 'IMPORT_CLEANUP_INCOMPLETE'
@dataclass(frozen=True)
class ImportFailure:
"""One classification of one failure: the step, its code and the sentence it stores."""
step: ImportStep | None
code: ImportErrorCode
reason: str
_CONTACT_SUPPORT = "Contact support with this import's identifier."
_RE_EXPORT = 'Export the model again from the origin platform.'
IMPORT_FAILURE_CATALOG: Final[Mapping[ImportStep, tuple[ImportErrorCode, str]]] = MappingProxyType(
{
ImportStep.RECEIVED: (
ImportErrorCode.REQUEST_INVALID,
'The import could not be started because the request was incomplete or '
'contradictory. Nothing was created. Submit the import again.',
),
ImportStep.DOWNLOAD: (
ImportErrorCode.FILE_UNREADABLE,
'The uploaded file could not be read, or it is not the file that was sent. '
'Upload it again.',
),
ImportStep.DECRYPTION: (
ImportErrorCode.BUNDLE_CANNOT_BE_OPENED,
'This file could not be opened: the password is wrong, or the file is damaged or '
'was altered. Confirm the password with whoever exported the model, then try '
'again.',
),
ImportStep.ARCHIVE_INSPECTION: (
ImportErrorCode.BUNDLE_REJECTED,
"The file's contents did not pass the platform's safety checks, so it was not "
f'opened. {_RE_EXPORT}',
),
ImportStep.EXTRACTION: (
ImportErrorCode.BUNDLE_UNPACK_FAILED,
f'The file could not be unpacked safely and nothing from it was kept. {_RE_EXPORT}',
),
ImportStep.STRUCTURE_VALIDATION: (
ImportErrorCode.BUNDLE_INCOMPLETE,
'This file is not a complete model export — part of what the platform needs is '
f'missing from it. {_RE_EXPORT}',
),
ImportStep.CONTENT_POLICY: (
ImportErrorCode.BUNDLE_UNEXPECTED_CONTENT,
'The file contains something a model export should not contain, so it was '
f'rejected. {_RE_EXPORT}',
),
ImportStep.EXPERIMENT_CREATION: (
ImportErrorCode.STORAGE_PREPARATION_FAILED,
'The platform could not prepare a place to keep this model. The import stopped '
f'and no model was created. {_CONTACT_SUPPORT}',
),
ImportStep.ARTIFACT_UPLOAD: (
ImportErrorCode.MODEL_FILES_NOT_STORED,
"The model's files could not be stored on this platform. The import stopped and "
f'the model is not available. {_CONTACT_SUPPORT}',
),
ImportStep.REGISTRATION: (
ImportErrorCode.MODEL_NOT_REGISTERED,
"The model's files were stored, but the model itself could not be registered, so "
f'it cannot be used. {_CONTACT_SUPPORT}',
),
ImportStep.MODEL_DOCUMENT: (
ImportErrorCode.MODEL_SETTINGS_NOT_SAVED,
'The model was registered, but the settings that tell the platform how to run it '
f'could not be saved. {_CONTACT_SUPPORT}',
),
}
)
UNCLASSIFIED_FAILURE: Final[tuple[ImportErrorCode, str]] = (
ImportErrorCode.UNEXPECTED_ERROR,
f'The import stopped for an unexpected reason and no model was created. {_CONTACT_SUPPORT}',
)
# `cleanup` deletes temporary files on the worker. It is never the reported failure: when every
# other step succeeded and only cleanup failed, the import is a success and the leak is a log line.
# So it deliberately has no user-facing sentence, and the catalog-completeness test knows it.
STEPS_WITHOUT_A_USER_FACING_REASON: Final[frozenset[ImportStep]] = frozenset({ImportStep.CLEANUP})
def classify_import_failure(
step: ImportStep | None, exception: BaseException | None = None
) -> ImportFailure:
"""Turn a failure into the one code and the one sentence that will be recorded.
The exception's text is **ignored entirely** — it is what carries paths, table names and type
names, and it belongs in the log. What decides the outcome is the step.
Args:
step: the step that was being attempted, or `None` when none had been entered.
exception: the failure, when there is one to hand. It is accepted so callers read as
"classify this failure" and never absent because the classification needs it: the
catalog is keyed by step, and an exception cannot cross an activity boundary with its
type intact anyway — the terminal write classifies from the step alone.
Returns:
ImportFailure: step, code and the reason to store.
"""
del exception # the catalog is keyed by step; the exception's text never reaches the record
if step is None or step not in IMPORT_FAILURE_CATALOG:
code, reason = UNCLASSIFIED_FAILURE
return ImportFailure(step=step, code=code, reason=reason)
code, reason = IMPORT_FAILURE_CATALOG[step]
return ImportFailure(step=step, code=code, reason=reason)
class ImportLogRowNotClaimableError(ApplicationError):
"""The record named by `import_run_id` is not this workflow's to write.
Raised when the claim affects zero rows: no such row, a `TRAIN` row, or a row already running
under another workflow. Non-retryable on purpose — a second attempt would read the same rows and
reach the same conclusion — and loud everywhere except the row itself, which belongs to someone
else.
"""
def __init__(self, message: str, *, import_run_id: Any = None, observed: str = '') -> None:
super().__init__(
message,
type='ImportLogRowNotClaimableError',
non_retryable=True,
)
self.import_run_id = import_run_id
self.observed = observed
class ImportStatusRejectedError(ApplicationError):
"""The database accepted the connection and refused the statement.
A `CHECK` violation, an undefined column, a value too long for its column. `V12` not being
applied yet lands here, at the terminal write, and it must fail on the first attempt instead of
spending the whole retry envelope to say so.
"""
def __init__(self, message: str) -> None:
super().__init__(message, type='ImportStatusRejectedError', non_retryable=True)
class ImportInputError(ApplicationError):
"""The request is incomplete or contradictory. Fails at `received`, on the claimed row."""
def __init__(self, message: str) -> None:
super().__init__(
message,
{'step': ImportStep.RECEIVED.value, 'gate': None},
type='ImportInputError',
non_retryable=True,
)
class ImportGateRejectedError(ApplicationError):
"""A gate rejected the bundle, and the error says which one.
Gates 2 to 7 all run inside one activity, so the step the workflow was *attempting* is not
precise enough to record. An exception loses its Python type crossing an activity boundary but
keeps its `details`, so the step and the gate travel there and the terminal write can name them.
"""
def __init__(
self,
message: str,
*,
step: ImportStep,
gate: int | None = None,
cause: str | None = None,
) -> None:
super().__init__(
message,
{'step': step.value, 'gate': gate, 'cause': cause},
type='ImportGateRejectedError',
non_retryable=True,
)
self.step = step
self.gate = gate
def failure_location(exception: BaseException | None) -> tuple[str | None, int | None]:
"""Read the step and gate an activity's failure carries, if it carried any.
Walks the `__cause__` chain because Temporal wraps an activity failure in an `ActivityError`.
Returns `(None, None)` when the failure did not say where it happened, and the caller then falls
back to the step it was attempting.
"""
seen = 0
candidate: BaseException | None = exception
while candidate is not None and seen < 10:
details = getattr(candidate, 'details', None)
if details and isinstance(details[0], dict) and 'step' in details[0]:
location = details[0]
gate = location.get('gate')
return location.get('step'), int(gate) if isinstance(gate, int) else None
candidate = candidate.__cause__
seen += 1
return None, None
class FailureSinks(Protocol):
"""What the channel needs from its caller: the repository's existing observability surface."""
def error(self, message: str, metadata: dict | None = None) -> None: ...
def get_core_labels(self, metadata: dict | None = None, operation_type: str = '-') -> dict: ...
async def emit_metric(
self,
metric_object: Any,
tags: dict[str, Any],
method: str = 'inc',
value: float | None = None,
) -> None: ...
async def send_notification_async(
self,
metadata: dict,
notification_id: str,
message: str,
block: str,
level: str = 'INFO',
attachment_content: str | None = None,
) -> None: ...
async def report_import_failure(
sinks: FailureSinks,
step: ImportStep | None,
exception: BaseException | None,
*,
import_run_id: Any,
workflow_id: str,
metadata: dict[str, Any] | None = None,
code: ImportErrorCode | None = None,
marker: str | None = None,
verdict: str | None = None,
detail_text: str | None = None,
) -> ImportFailure:
"""Report one failure through every sink, exactly once.
Sinks, in order: a structured log carrying the ids, the step, the code **and** the exception
with its traceback (the only sink that gets the exception text); the import error metric, with
the code in `operation_type`; a notification through the existing handler. The exception itself
reaches Temporal by being raised, and the record is the fifth, optional sink — the terminal
write, when the row can be written at all.
Args:
sinks: the activity providing logger, metrics and notifications.
step: the pipeline step, or `None` for a failure that is not one (see `marker`).
exception: the failure.
import_run_id: the record's id, for correlation.
workflow_id: the identifier a user quotes to support.
metadata: workflow metadata for labels and the notification.
code: an explicit code, for the two failures that are not pipeline steps.
marker: `not_claimable` or `status_write` — what a step name would have been.
verdict: for a failed terminal write, the verdict it could not record, so the outcome
survives in text even when it cannot survive in the row.
detail_text: the failure as text, for the callers whose exception object is already gone —
anything that crossed an activity boundary. Used in place of a traceback.
Returns:
ImportFailure: the classification that was reported.
"""
failure = classify_import_failure(step, exception)
if code is not None:
failure = ImportFailure(step=step, code=code, reason=failure.reason)
where = marker or (step.value if step else 'unknown')
if exception is not None:
detail = ''.join(traceback.format_exception(exception)).strip()
else:
detail = detail_text or 'no exception object available'
verdict_line = f' verdict_not_recorded={verdict}' if verdict else ''
sinks.error(
f'model import failed: step={where} code={failure.code.value} '
f'import_run_id={import_run_id} workflow_id={workflow_id}{verdict_line}\n{detail}',
metadata or {},
)
labels = sinks.get_core_labels(metadata or {}, operation_type=failure.code.value)
await sinks.emit_metric(metric_object=metrics.MODEL_IMPORT_ERROR_COUNT, tags=labels)
await sinks.send_notification_async(
metadata=metadata or {},
notification_id=failure.code.value,
message=failure.reason,
block=where,
level='ERROR',
)
return failure

View File

@@ -0,0 +1,364 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import time
import traceback
import warnings
from typing import Any
import numpy as np
from pandas import DataFrame, Index, to_datetime
from sientia.ModelAnalysis import ModelAnalysis
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from laborious import metrics
from laborious.utils.dataframe_debug import build_dataframe_debug_message
warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0')
warnings.filterwarnings(
'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
)
class ModelMetrics(SientiaMonitoring):
"""
Metrics activities for the Laborious system.
This class provides activities for writing metrics to the Prometheus monitoring system.
"""
_MAX_DEBUG_DATAFRAME_ROWS = 100
def __init__(
self,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
def close(self) -> None:
"""
Close the model metrics activity and clean up resources.
"""
SientiaMonitoring.shutdown(self)
def __del__(self):
self.close()
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
"""
Log dataframe content only when row count is below the configured threshold
Args:
- message (str): Base log message to identify the dataframe in logs
- data (Any): Dataframe-like payload to be logged
- metadata (dict[str, Any]): Workflow metadata for contextual logging
"""
self.debug(
build_dataframe_debug_message(
message=message,
data=data,
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
),
metadata,
)
async def get_drift_metrics(
self,
reference_data: DataFrame,
target_data: DataFrame,
target_name: str,
reference_columns: Index,
drift_metrics: list[str],
chunk_period: str,
metadata: dict[str, Any],
) -> DataFrame:
"""
Calculate univariate drift metrics for a model.
Args:
model_analysis (ModelAnalysis): Model analysis object
reference_data (DataFrame): Reference data
target_data (DataFrame): Target data
reference_columns (list[str]): Reference columns
drift_metrics (list[str]): Drift metrics
metadata (dict[str, Any]): Workflow execution metadata
"""
config = {
'target': target_name,
'prediction': 'prediction',
'timestamp': 'timestamp',
'features': reference_columns,
}
model_analysis = ModelAnalysis(config=config)
self._debug_dataframe(
f'Reference data: Size {reference_data.shape}', reference_data, metadata
)
self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata)
core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
start_time = time.time()
try:
univariate_drift = model_analysis.detect_univariate_drift(
reference_df=reference_data,
analysis_df=target_data,
features=reference_columns,
timestamp_col=config['timestamp'],
methods=drift_metrics,
chunk_period=chunk_period,
)
except Exception as e:
self.error(f'Error detecting univariate drift: {e}', metadata)
await self.emit_metric(
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
start_time = time.time()
try:
multivariate_drift = model_analysis.detect_multivariate_drift(
reference_df=reference_data,
analysis_df=target_data,
features=reference_columns,
timestamp_col=config['timestamp'],
chunk_period=chunk_period,
)
except Exception as e:
self.error(f'Error detecting multivariate drift: {e}', metadata)
await self.emit_metric(
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
start_time = time.time()
core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe')
try:
drift_df = model_analysis.get_drift_metrics_dataframe(
univariate_drift=univariate_drift,
multivariate_drift=multivariate_drift,
)
except Exception as e:
self.error(f'Error getting drift metrics: {e}', metadata)
await self.emit_metric(
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
)
raise e
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
return drift_df
@activity.defn(name='calculate_drift')
async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
"""
Calculate drift metrics for a model.
Args:
input_data (dict[str, Any]): Input data containing:
- metadata (dict): Workflow execution metadata
- model_name (str): Name of the MLFlow model to calculate drift for
- reference_data (pd.DataFrame): Reference data for the model
- target_data (pd.DataFrame): Target data for calculating drift
- target_name (str): Name of the target column
- drift_metrics (list[str]): List of drift metrics to calculate
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
model_id = input_data['model_id']
reference_raw_data = input_data['reference_data']
target_data = DataFrame(input_data['target_data'])
target_name = input_data['target_name']
drift_metrics = input_data['drift_metrics']
chunk_period = input_data['chunk_period']
if chunk_period not in ['min', 's']:
self.error(f'Invalid chunk period: {chunk_period}', metadata)
raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"')
self.info(f'Calculating drift for model {model_name}', metadata)
target_data = target_data.pivot(index='timestamp', columns='variable', values='value')
target_data['timestamp'] = target_data.index
target_data['timestamp'] = to_datetime(target_data['timestamp'])
target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT)
target_data = target_data.reset_index(drop=True)
target_data.dropna(inplace=True)
if reference_raw_data is not None:
self.info('Using reference data', metadata)
reference_data = DataFrame(reference_raw_data)
accurate = True
else:
# Get 30% first rows of target_data
self.warning('Using 30% first rows of target data as reference data', metadata)
target_data.sort_values(by='timestamp', ascending=True, inplace=True)
reference_data = target_data.head(int(len(target_data) * 0.3))
accurate = False
await self.send_notification_async(
metadata=metadata,
notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
message='Using 30% first rows of target data as reference data',
block='model_metrics',
level=NotificationLevel.WARNING,
attachment_content=reference_data.to_csv(),
)
reference_columns = reference_data.drop(
columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
).columns
try:
drift_df = await self.get_drift_metrics(
reference_data=reference_data,
target_data=target_data,
target_name=target_name,
reference_columns=reference_columns,
drift_metrics=drift_metrics,
chunk_period=chunk_period,
metadata=metadata,
)
except Exception as e:
self.error(f'Error getting drift metrics: {e}', metadata)
await self.send_notification_async(
metadata=metadata,
notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
message=f'Error getting drift metrics: {e}',
block='model_metrics',
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc(),
)
return []
if drift_df.empty:
self.warning('No drift metrics found', metadata)
return []
# Drop unnecessary columns
drift_df.drop(columns=['p_value'], inplace=True)
# Extract timestamps only until minutes
if chunk_period == 'min':
target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
else:
target_timestamps = target_data['timestamp']
# Drop rows where timestamp is not in target data, to avoid save drift from reference
drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
if drift_df.empty:
self.warning(
'No drift metrics found after dropping rows where timestamp is not in target data',
metadata,
)
return []
# Rename columns to match database columns
drift_df.rename(
columns={
'metric': 'method',
'statistic': 'value',
},
inplace=True,
)
# Drop duplicates
drift_df.drop_duplicates(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
drift_df['model_id'] = model_id
drift_df['accurate'] = accurate
drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
return drift_df.to_dict(orient='records')
@activity.defn(name='calculate_simple_metrics')
async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
"""
Calculate simple metrics for a model. Metrics available are:
- rmse
- mse
- mae
- r2
- accuracy
- precision
- recall
- f1
Args:
input_data (dict[str, Any]): Input data containing:
- metadata (dict): Workflow execution metadata
- model_id (str): ID of the MLFlow model
- target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns
- metrics (list[str]): List of metrics to calculate
Returns:
dict[Hashable, Any]: Dictionary containing the calculated metrics
"""
metadata = input_data['metadata']
model_id = input_data['model_id']
target_data = DataFrame(input_data['target_data'])
metrics = input_data['metrics']
interval_minutes = input_data['interval_minutes']
data_size = target_data.shape[0]
output_data = []
diff = target_data['target'] - target_data['prediction']
diff_squared = diff**2
self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
for metric in metrics:
if metric == 'rmse':
output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
elif metric == 'mse':
output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
elif metric == 'mae':
output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
elif metric == 'r2':
y_true = target_data['target']
y_mean = np.mean(y_true)
ss_res = np.sum(diff_squared)
ss_tot = np.sum((y_true - y_mean) ** 2)
# Evita divisão por zero
if ss_tot == 0:
r2_score = 0.0
else:
r2_score = 1 - (ss_res / ss_tot)
output_data.append({'metric': 'r2', 'value': r2_score})
data = DataFrame(output_data)
data['model_id'] = model_id
data['timestamp'] = target_data['timestamp'].max()
data['data_size'] = data_size
data['interval_minutes'] = interval_minutes
self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata)
return data.to_dict(orient='records')

515
laborious/activities/opc.py Normal file
View File

@@ -0,0 +1,515 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import traceback
from collections.abc import Hashable
from typing import Any
from pandas import DataFrame
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from laborious.utils.repository.opc_repository import OpcRepository
OPC_WRITTING_ERROR_CONFIDENCE = 12
OPC_SESSION_BAD_CONFIDENCE = 14
OPC_SESSION_BAD_COMMENT_PREFIX = 'OPC UA session/channel error:'
OPC_WRITTING_ERROR_MESSAGE = 'Some data could not be written to OPC servers'
OPC_RECONNECT_IN_PROGRESS_COMMENT = 'OPC UA reconnect in progress'
OPC_COMMENT_SEPARATOR = ' | '
def _opc_session_bad_comment(opc_status: str | None) -> str:
status = opc_status or 'Unknown'
return f'{OPC_SESSION_BAD_COMMENT_PREFIX} {status}'
def _apply_opc_write_error(
error_info: dict[str, Any] | None,
session_bad_seen: bool,
session_bad_status: str | None,
reconnect_in_progress_seen: bool,
) -> tuple[bool, str | None, bool]:
"""
Update session/reconnect flags from an OPC write error payload.
Args:
error_info: Repository error details, or None when the write succeeded.
session_bad_seen: Whether a session_bad error was seen so far.
session_bad_status: Last known OPC status for session errors.
reconnect_in_progress_seen: Whether reconnect_in_progress was seen so far.
Return:
Updated (session_bad_seen, session_bad_status, reconnect_in_progress_seen).
"""
if not error_info:
return session_bad_seen, session_bad_status, reconnect_in_progress_seen
kind = error_info.get('opc_error_kind')
if kind == 'session_bad':
return True, error_info.get('opc_status', session_bad_status), reconnect_in_progress_seen
if kind == 'reconnect_in_progress':
return session_bad_seen, session_bad_status, True
return session_bad_seen, session_bad_status, reconnect_in_progress_seen
class OPC(SientiaMonitoring):
"""
OPC server integration activities for real-time data export.
This class provides comprehensive OPC UA client functionality for connecting
to multiple OPC servers and writing prediction data in real-time. It implements
secure communication with certificate-based authentication and automatic
reconnection capabilities.
The class supports multiple OPC servers with individual configurations and
provides robust error handling and monitoring for production environments.
Attributes:
opc_servers (dict): Configuration for multiple OPC servers
opc_repository (dict): Active OPC repository connections
logger (Logger): Logging and observability instance
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(
self,
opc_servers: dict[str, dict[str, Any]],
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
self.logger = logger
self.notification_handler = notification_handler
self.opc_servers = opc_servers
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.opc_repository: dict[str, OpcRepository] = {}
async def init_opc(self):
"""
Initialize OPC server connections and establish communication channels.
This method iterates through all configured OPC servers and attempts to
establish secure connections using certificate-based authentication.
Each server connection is managed independently, and connection failures
are reported through the notification system.
The method performs the following operations:
1. Creates OpcRepository instances for each configured server
2. Establishes secure connections with certificate validation
3. Reports connection success/failure through notifications
4. Logs connection status for operational visibility
Raises:
Exception: If OPC repository initialization fails or connection
establishment encounters critical errors
Note:
Connection failures are logged and reported but do not prevent
the initialization of other OPC servers. Each server is handled
independently to ensure maximum availability.
"""
self.logger.info('Initializing OPC servers...')
for opc_id, server in self.opc_servers.items():
self.opc_repository[opc_id] = OpcRepository(
opc_id=server['id'],
server_name=server['server_name'],
url=server['url'],
logger=self.logger,
server_uri=server['server_uri'],
cert_path=server['cert_path'],
private_key_path=server['private_key_path'],
server_cert_path=server['server_cert_path'],
notification_handler=self.notification_handler,
reconnection_interval=server['reconnection_interval'],
metrics_controller=self.metrics_controller,
)
is_connected, error_data = await self.opc_repository[opc_id].connect()
if not is_connected:
await self.send_notification_async(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION',
},
notification_id=error_data['notification_id'],
message=error_data['message'],
block=error_data['block'],
level=error_data.get('level', NotificationLevel.ERROR),
attachment_content=error_data.get('attachment_content', None),
)
else:
self.logger.info(
f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
)
async def write_data(
self,
server_id: str,
tag: str,
data: Any,
data_type: str,
tag_type: str,
metadata: dict[str, Any],
) -> tuple[float | None, dict[str, Any] | None]:
"""
Write data to a specific OPC server tag with comprehensive error handling.
Return:
tuple[float | None, dict[str, Any] | None]: Response time on success, or
(None, error info_data) on repository failure.
"""
try:
is_success, info_data = await self.opc_repository[server_id].write_data(
tag, data, data_type, metadata
)
if not is_success:
await self.send_notification_async(
metadata=metadata,
notification_id=info_data['notification_id'],
message=info_data['message'],
block=info_data['block'],
level=info_data.get('level', NotificationLevel.ERROR),
attachment_content=info_data.get('attachment_content', None),
)
return None, info_data
return info_data['response_time'], None
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
message=f'Error writing data to OPC server: {e}',
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise e
async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
"""
Validate that an OPC server is available and configured for write operations.
This method checks if the specified OPC server exists in the active
repository and is available for data writing operations. It provides
immediate feedback for server availability and logs validation failures
for operational monitoring.
Args:
server_id (str): Unique identifier for the OPC server to validate
metadata (dict[str, Any]): Context metadata for logging and notifications
Returns:
bool: True if server is available, False otherwise
Note:
Server validation failures are automatically reported through the
notification system with detailed information about available servers.
This helps operators quickly identify configuration issues.
"""
if self.opc_repository.get(server_id) is None:
message = f'OPC server {server_id} not found to perform write operation.'
await self.send_notification_async(
metadata=metadata,
notification_id='OPC_SERVER_NOT_FOUND',
message=message,
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
)
return False
return True
async def _write_tags_from_config(
self,
server_id: str,
tags_config: dict[str, dict[str, Any]],
data: DataFrame,
data_column: str,
tag_type: str,
log_label: str,
metadata: dict[str, Any],
) -> tuple[dict[str, float | None], bool, str | None, bool]:
"""
Write a group of OPC tags and collect response times and error flags.
Args:
server_id: Target OPC server identifier.
tags_config: Tag name to configuration mapping.
data: DataFrame with prediction/confidence columns.
data_column: Column name whose first row value is written.
tag_type: Tag category passed to write_data ('prediction' or 'confidence').
log_label: Human-readable label for success logs.
metadata: Context metadata for logging and notifications.
Return:
(response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen)
"""
response_times: dict[str, float | None] = {}
session_bad_seen = False
session_bad_status: str | None = None
reconnect_in_progress_seen = False
for tag, tag_config in tags_config.items():
response_time, error_info = await self.write_data(
server_id=server_id,
tag=tag,
data=data.head(1)[data_column].values[0],
data_type=tag_config['data_type'],
tag_type=tag_type,
metadata=metadata,
)
session_bad_seen, session_bad_status, reconnect_in_progress_seen = (
_apply_opc_write_error(
error_info,
session_bad_seen,
session_bad_status,
reconnect_in_progress_seen,
)
)
if response_time is not None:
self.info(
f'{log_label} written to OPC server {server_id} for tag {tag}.',
metadata,
)
response_times[tag] = response_time
return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen
async def manage_output_tags(
self,
server_id: str,
config: dict[str, Any],
data: DataFrame,
metadata: dict[str, Any],
) -> tuple[bool, dict[str, float | None], bool, str | None, bool]:
"""
Manage the writing of prediction and confidence data to OPC server tags.
This method orchestrates the writing of multiple data types to OPC servers
based on configuration. It handles both prediction data and confidence
values independently, allowing for flexible tag configuration and
comprehensive error handling.
The method supports two main tag types:
1. Prediction tags: Write actual prediction values to configured OPC tags
2. Confidence tags: Write confidence scores to separate OPC tags
Args:
server_id (str): Unique identifier for the target OPC server
config (dict[str, Any]): OPC tag configuration containing:
- prediction_tags (dict, optional): Prediction tag configurations
- confidence_tags (dict, optional): Confidence tag configurations
data (DataFrame): DataFrame containing prediction and confidence data
metadata (dict[str, Any]): Context metadata for logging and notifications
success (bool): Current success status to maintain across operations
Returns:
tuple[bool, int]: (overall_success, total_tags_written)
- overall_success: True if all configured tags were written successfully
- total_tags_written: Count of successfully written tags
"""
response_times: dict[str, float | None] = {}
session_bad_seen = False
session_bad_status: str | None = None
reconnect_in_progress_seen = False
tag_groups = (
('prediction_tags', 'prediction', 'prediction', 'Prediction data'),
('confidence_tags', 'prediction_confidence', 'confidence', 'Confidence data'),
)
for config_key, data_column, tag_type, log_label in tag_groups:
if config_key not in config:
continue
(
group_times,
group_session_bad,
group_status,
group_reconnect,
) = await self._write_tags_from_config(
server_id=server_id,
tags_config=config[config_key],
data=data,
data_column=data_column,
tag_type=tag_type,
log_label=log_label,
metadata=metadata,
)
response_times.update(group_times)
if group_session_bad:
session_bad_seen = True
session_bad_status = group_status or session_bad_status
if group_reconnect:
reconnect_in_progress_seen = True
success = None not in response_times.values()
return (
success,
response_times,
session_bad_seen,
session_bad_status,
reconnect_in_progress_seen,
)
@activity.defn(name='write_opc_data')
async def write_opc_data(
self, input_data: dict[str, Any]
) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
"""
Write prediction and confidence data to OPC servers. The two writing
operations are optional and independent of each other.
Args:
- input_data(dict[str, Any]): The input data. Contains the following keys:
- data(dict[str, Any]): The dataframe that contains the data to write
to the OPC servers.
- opc_output_config(dict[str, Any]): The OPC writing configuration.
The keys are the OPC server names and the values contain:
- prediction_tags(dict[str, Any]): The tags to write to the OPC servers.
- confidence_tags(dict[str, Any]): The tags to write to the OPC servers.
Returns:
- dict[Any, Any]: The data that was written to the OPC servers.
"""
metadata = input_data['metadata']
self.info('Writing data to OPC servers...', metadata)
data = DataFrame(input_data['data'])
opc_output_config = input_data['opc_output_config']
self.info(f'Data to write: {data.size} rows', metadata)
success = True
session_bad_seen = False
session_bad_status: str | None = None
reconnect_in_progress_seen = False
metrics: dict[str, dict[str, float | None]] = {}
for server_id, config in opc_output_config.items():
if not await self.validate_server(server_id, metadata):
success = False
continue
(
local_success,
local_response_times,
local_session_bad,
local_status,
local_reconnect_in_progress,
) = await self.manage_output_tags(server_id, config, data, metadata)
metrics[server_id] = local_response_times
local_count = len(local_response_times)
success = success and local_success
if local_session_bad:
session_bad_seen = True
session_bad_status = local_status or session_bad_status
if local_reconnect_in_progress:
reconnect_in_progress_seen = True
self.info(
f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
metadata,
)
return (
self.process_confidence(
data,
success,
metadata,
session_bad=session_bad_seen,
opc_status=session_bad_status,
reconnect_in_progress=reconnect_in_progress_seen,
),
metrics,
)
def process_confidence(
self,
data: DataFrame,
success: bool,
metadata: dict[str, Any],
*,
session_bad: bool = False,
opc_status: str | None = None,
reconnect_in_progress: bool = False,
) -> dict[Hashable, Any]:
"""
Process prediction confidence based on OPC write operation success.
This method updates the prediction confidence values in the DataFrame
based on the success status of OPC server write operations. If any
write operations failed, it sets the confidence to a predefined error
value to indicate data quality issues.
The method implements a confidence degradation strategy:
- Success: Maintains original confidence values
- Failure: Sets confidence to error value for operational awareness
Args:
data (DataFrame): DataFrame containing prediction and confidence data
success (bool): Overall success status of OPC write operations
metadata (dict[str, Any]): Context metadata for logging and notifications
Returns:
dict[Any, Any]: Processed data as a dictionary with updated confidence values
Note:
The error confidence value (OPC_WRITTING_ERROR_CONFIDENCE = 12) is
used to indicate that data was not successfully exported to OPC servers.
This allows downstream systems to handle data quality appropriately.
"""
if not success:
comment_parts: list[str] = []
confidence = OPC_WRITTING_ERROR_CONFIDENCE
if session_bad:
comment_parts.append(_opc_session_bad_comment(opc_status))
confidence = OPC_SESSION_BAD_CONFIDENCE
if reconnect_in_progress:
comment_parts.append(OPC_RECONNECT_IN_PROGRESS_COMMENT)
confidence = OPC_SESSION_BAD_CONFIDENCE
if not comment_parts:
comment_parts.append(OPC_WRITTING_ERROR_MESSAGE)
comments = OPC_COMMENT_SEPARATOR.join(comment_parts)
data['prediction_confidence'] = confidence
data['comments'] = comments
self.debug(
f'OPC write issues, confidence={confidence}, comments={comments}',
metadata,
)
else:
self.debug('Data written to OPC servers successfully.', metadata)
return data.to_dict()
async def aclose(self):
"""
Gracefully shutdown all OPC server connections and cleanup resources.
This method ensures proper cleanup of all active OPC server connections
by calling the disconnect method on each repository instance. It's
designed to be called during application shutdown to prevent resource
leaks and ensure clean termination.
The method performs the following cleanup operations:
1. Iterates through all active OPC repository connections
2. Calls disconnect() on each repository instance
3. Allows for graceful connection termination
4. Prevents resource leaks and connection hanging
Note:
This method should be called during application shutdown to ensure
proper cleanup. It handles all active connections regardless of
their current state and provides a clean shutdown experience.
"""
for opc in self.opc_repository.values():
await opc.disconnect()

View File

@@ -0,0 +1,210 @@
from temporalio import activity, workflow
from laborious.utils.repository.minio_manager import MinioManager
with workflow.unsafe.imports_passed_through():
# Extend the Temporal Postgres activities for convenient query -> MinIO export
import traceback
from datetime import timedelta
from typing import Any
import pandas as pd
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.repository.minio_repository import MinioRepository
from sientia_do.temporal.activities.postgres import Postgres
from sientia_do.temporal.constants import now
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
class Storage(Postgres, MinioManager):
"""
Extensions for Postgres activities with a helper to export query results
directly to MinIO as Parquet and return the object name.
"""
minio_repository: MinioRepository | None = None
def __init__(
self,
host: str,
port: int,
user: str,
password: str,
dbname: str,
min_connections: int,
max_connections: int,
retention_hours: int = 24,
minio_repository: MinioRepository | None = None,
logger: Logger | None = None,
notification_handler: NotificationHandler | None = None,
metrics_controller: MetricsController | None = None,
):
self.retention_hours = retention_hours
Postgres.__init__(
self,
host=host,
port=port,
user=user,
password=password,
dbname=dbname,
min_connections=min_connections,
max_connections=max_connections,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
MinioManager.__init__(
self, minio_repository, logger, notification_handler, metrics_controller
)
@activity.defn(name='load_query_with_minio_offload')
async def load_query_with_minio_offload(
self, input_data: dict[str, Any]
) -> MinioDataFramePayload:
"""
Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
Args (input_data):
metadata (dict): Workflow metadata (same as load_custom_query).
query (str): SQL query.
datetime_columns (list[str], optional): Datetime column names.
model_name (str): Model name for object key basename.
key_prefix (str, optional): Directory prefix inside the bucket.
size_threshold_bytes (int, optional): Override env offload threshold.
Returns:
dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
metadata: dict = input_data.get('metadata', {})
model_name = input_data['model_name']
rows = await self.load_custom_query(
input_data,
)
if not rows:
self.error(
'load_query_with_minio_offload failed: No data returned from query', metadata
)
dataframe = None
else:
dataframe = pd.DataFrame(rows)
return await MinioDataFramePayload.from_dataframe(
dataframe,
minio_repo=self.minio_repository,
workflow_metadata=metadata,
model_name=model_name,
operation='initial',
logger=self.logger,
)
@activity.defn(name='export_payload_to_postgres')
async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
"""
Export a payload to PostgreSQL.
"""
metadata = input_data.get('metadata')
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = await payload.retrieve(self.minio_repository, metadata)
return await self.export_data_to_postgres(
{
**input_data,
'data': data,
}
)
@activity.defn(name='cleanup_minio_objects_expired')
async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Delete objects under the given prefixes that are older than the retention window.
Args (input_data):
metadata (dict): Workflow metadata for logging and metrics.
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
Returns:
dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
metadata = input_data.get('metadata', {})
payload = MinioDataFramePayload.from_dict(input_data['data'])
prefix = payload.cleanup_prefix()
base = now()
cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
hours=self.retention_hours
)
report: dict[str, Any] = {
'failed': {},
'deleted': {},
'failed_count': 0,
'deleted_count': 0,
}
try:
keys = await self.minio_repository.list_objects(
prefix=prefix,
recursive=True,
metadata=metadata,
)
for key in keys:
try:
ts = MinioDataFramePayload.parse_object_timestamp(key)
if ts is None:
continue
if ts >= cutoff:
continue
await self.minio_repository.delete_file(
object_name=key,
metadata=metadata,
)
except Exception as e:
report['failed'][key] = {
'success': False,
'message': str(e),
}
report['failed_count'] += 1
continue
report['deleted'][key] = {
'success': True,
'message': 'Deleted',
}
report['deleted_count'] += 1
except Exception as e:
trace = traceback.format_exc()
await self.send_notification_async(
metadata=metadata,
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
message=f'Error cleaning up MinIO objects: {e}',
block='cleanup_minio_objects_expired',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
self.error(trace, metadata)
else:
# Cleanup success is expected in normal flow; avoid noisy INFO notifications
# that do not impact behavior and can flood observability in test runs.
self.info('MinIO objects cleaned up successfully', metadata)
return report
def close(self) -> None:
"""Close Storage resources (MinIO client and Postgres engine)."""
if hasattr(self, 'engine'):
Postgres.close(self)
MinioManager.close(self)
def __del__(self):
self.close()

239
laborious/metrics.py Normal file
View File

@@ -0,0 +1,239 @@
"""
Laborious Metrics Module
This module defines all Prometheus metrics used by the Sientia DataOps Laborious system
for monitoring and observability. The metrics provide insights into system performance,
prediction quality, and operational health.
The metrics are designed to be scraped by Prometheus and can be visualized in
Grafana or other monitoring dashboards to provide real-time visibility into
the system's operation.
Key Metric Categories:
- Application Health: Overall system status and availability
- Prediction Operations: Count and performance of prediction operations
- Data Quality: Confidence levels and validation results
- Export Operations: Database and OPC export performance
- Response Times: Performance monitoring for various operations
Metric Labels:
- pod_id: Kubernetes pod identifier for multi-instance deployments
- runtime: Runtime / environment identifier (matches ``RUNTIME`` env, see ``SientiaMonitoring``)
- model_name: Name of the ML model being used
- workflow_name: Name of the prediction pipeline
- opc_server_id: Identifier for OPC server operations
"""
from prometheus_client import Counter, Gauge, Histogram
from sientia_do.observability.metrics import CORE_LABELS
# Application health metric
APP_UP = Gauge(
'app_up',
'Indicates if the application is running (1) or shutting down (0)',
['pod_id'],
)
# Prediction operation metrics
PREDICTIONS_WRITTEN_COUNT = Counter(
'laborious_predictions_written_count',
'Number of predictions written to the database table predictions',
CORE_LABELS,
)
# Prediction quality metrics
PREDICTION_CONFIDENCE_MONITOR = Gauge(
'laborious_prediction_confidence_monitor',
'Current confidence of each prediction',
CORE_LABELS,
)
# Prediction total response time
PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
'laborious_prediction_response_time_monitor',
'Current response time of each prediction',
CORE_LABELS,
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)
# ================== OPC metrics ==================
PREDICTION_OPC_WRITING_COUNT = Counter(
'laborious_prediction_opc_writing_count',
'Number of predictions written to the OPC server',
[*CORE_LABELS, 'opc_server_id', 'tag'],
)
PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
'laborious_prediction_opc_writing_response_time_monitor',
'Current response time of each prediction written to the OPC server',
[*CORE_LABELS, 'opc_server_id', 'tag'],
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)
OPC_CONNECTIONS_TOTAL = Counter(
'opc_connections_initiated_total',
'Total connection attempts to OPC servers',
['pod_id', 'server_name'],
)
OPC_CONNECTIONS_FAILED = Counter(
'opc_connections_failed_total',
'Total failed connection attempts to OPC servers',
['pod_id', 'server_name'],
)
OPC_CONNECTION_STATUS = Gauge(
'opc_connection_status',
'Connection status with the OPC server (1=connected, 0=disconnected)',
['pod_id', 'server_name', 'server_url'],
)
_OPC_SESSION_DEBUG_LABELS = ['pod_id', 'server_name', 'runtime', 'opc_server_id', 'session_id']
OPC_SESSION_CREATED_TOTAL = Counter(
'opc_session_created_total',
'OPC UA sessions established (after successful connect)',
_OPC_SESSION_DEBUG_LABELS,
)
OPC_SESSION_CLOSED_TOTAL = Counter(
'opc_session_closed_total',
'OPC UA client disconnects completed (session tear-down initiated)',
_OPC_SESSION_DEBUG_LABELS,
)
OPC_SESSION_REVISED_TIMEOUT_MS = Gauge(
'opc_session_revised_timeout_milliseconds',
'Server-revised OPC UA session timeout (RevisedSessionTimeout) in ms after connect',
_OPC_SESSION_DEBUG_LABELS,
)
OPC_WRITE_ATTEMPT_LABELS = [*_OPC_SESSION_DEBUG_LABELS, 'model_id', 'model_name', 'result']
OPC_WRITE_ATTEMPTS_TOTAL = Counter(
'opc_write_attempts_total',
'OPC UA write attempts with session and outcome (result=OK or exception class name)',
OPC_WRITE_ATTEMPT_LABELS,
)
OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL = Counter(
'opc_write_inter_arrival_over_session_timeout_total',
'Successful writes where seconds since the previous successful write exceeded RevisedSessionTimeout (ms)',
_OPC_SESSION_DEBUG_LABELS,
)
# ================== Model metrics ==================
MODEL_READ_LAG = Histogram(
'laborious_model_read_lag',
'Lag between the start and read of read operations',
CORE_LABELS,
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)
MODEL_WRITE_LAG = Histogram(
'laborious_model_write_lag',
'Lag between the start and end of write operations',
CORE_LABELS,
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)
MODEL_READ_COUNT = Counter(
'laborious_model_read_count',
'Number of reads from the model',
CORE_LABELS,
)
MODEL_WRITE_COUNT = Counter(
'laborious_model_write_count',
'Number of writes to the model',
CORE_LABELS,
)
MODEL_READ_ERROR_COUNT = Counter(
'laborious_model_read_error_count',
'Number of errors reading from the model',
CORE_LABELS,
)
MODEL_WRITE_ERROR_COUNT = Counter(
'laborious_model_write_error_count',
'Number of errors writing to the model',
CORE_LABELS,
)
MODEL_ANALYZE_LAG = Histogram(
'laborious_model_analyze_lag',
'Lag between the start and end of analyze operations',
CORE_LABELS,
buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
)
MODEL_ANALYZE_COUNT = Counter(
'laborious_model_analyze_count',
'Number of analyze operations',
CORE_LABELS,
)
MODEL_ANALYZE_ERROR_COUNT = Counter(
'laborious_model_analyze_error_count',
'Number of errors during analyze operations',
CORE_LABELS,
)
# ================== Model import metrics ==================
#
# The failing step and the error code travel in `operation_type`, the one label the shared
# `CORE_LABELS` leaves free. Nothing else about a failure is a label: an error reason is a sentence
# and an exception is a log line, neither of which belongs in a metric's cardinality.
MODEL_IMPORT_STARTED_COUNT = Counter(
'laborious_model_import_started_count',
'Number of model imports that claimed their import log row',
CORE_LABELS,
)
MODEL_IMPORT_COMPLETED_COUNT = Counter(
'laborious_model_import_completed_count',
'Number of model imports that reached a registered version and a model document',
CORE_LABELS,
)
MODEL_IMPORT_ERROR_COUNT = Counter(
'laborious_model_import_error_count',
'Number of failed model imports, with the failing step or error code in operation_type',
CORE_LABELS,
)
MODEL_IMPORT_KDF_DURATION = Histogram(
'laborious_model_import_kdf_duration_seconds',
'Time spent deriving the bundle key with Argon2id',
CORE_LABELS,
buckets=[0.5, 1.0, 2.0, 3.0, 5.0, 8.0, 13.0, 21.0],
)
MODEL_IMPORT_STEP_DURATION = Histogram(
'laborious_model_import_step_duration_seconds',
'Duration of each model import step, named in operation_type',
CORE_LABELS,
buckets=[0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 300.0],
)
# ================== PI Web API metrics ==================
PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name']
PI_WEB_API_PREDICTION_WRITTEN_COUNT = Counter(
'laborious_pi_web_api_prediction_written_count',
'Number of predictions written to the PI Web API',
PI_WEB_API_LABELS,
)
PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT = Counter(
'laborious_pi_web_api_prediction_written_error_count',
'Number of errors writing predictions to the PI Web API',
PI_WEB_API_LABELS,
)

View File

View File

@@ -0,0 +1,27 @@
"""Reading `.sientia` bundles: the wire format, the seven gates and the step vocabulary.
There is no writer here. The producer is the Streamlit app in `sientia-projects-templates`
(`app/src/operations/model_export/`), and a second production writer for a format with one producer
would be a fork waiting to happen. What stands in for it in the tests is
`tests/helpers/bundle_factory.py` plus the golden fixtures under `tests/fixtures/bundle/`.
"""
from laborious.utils.bundle.format import BundleFormatError, BundleHeader
from laborious.utils.bundle.reader import (
BundleGateError,
BundleLimits,
BundleReader,
OpenedBundle,
)
from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep
__all__ = [
'GATE_TO_STEP',
'BundleFormatError',
'BundleGateError',
'BundleHeader',
'BundleLimits',
'BundleReader',
'ImportStep',
'OpenedBundle',
]

View File

@@ -0,0 +1,216 @@
"""The `.sientia` wire format, as the producer writes it.
Every constant here is copied from the producer that emits the bundles this reader opens — the
Streamlit app in `sientia-projects-templates`, modules
`app/src/operations/model_export/crypto.py` and `app/src/operations/model_export/bundle.py`. The
producer file each value came from is named next to it. The golden fixtures under
`tests/fixtures/bundle/` are the anti-drift device: if the producer's layout moves, the tests here
fail against the committed bytes instead of failing in production.
This module deliberately imports nothing from Laborious, Temporal, MLflow or libsodium: it is pure
`struct` and `hashlib` over bytes, so it stays liftable into a shared package if the format ever
gets one (the cancelled QTZPOC-15). `reader.py` is where the crypto lives.
"""
from __future__ import annotations
import hashlib
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import Final
# --- Producer: app/src/operations/model_export/crypto.py ---
MAGIC: Final[bytes] = b'SIENTIA1'
FORMAT_VERSION: Final[int] = 1
AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM: Final[int] = 1
KDF_ID_ARGON2ID: Final[int] = 1
SALT_SIZE: Final[int] = 16
DIGEST_SIZE: Final[int] = 32
RESERVED_SIGNATURE_SIZE: Final[int] = 64
# `KDF_MEMLIMIT_BYTES` / `KDF_OPSLIMIT` on the producer: what it writes is also the ceiling the
# reader refuses to exceed, so header bytes cannot buy an arbitrary allocation on the worker.
MAX_KDF_MEMLIMIT: Final[int] = 256 * 1024 * 1024 # 268435456
MAX_KDF_OPSLIMIT: Final[int] = 3
# `KDF_PARALLELISM_RESERVED`: libsodium's `crypto_pwhash` has no lane parameter, so the field is
# reserved and must be exactly 1.
PARALLELISM: Final[int] = 1
CHUNK_SIZE: Final[int] = 64 * 1024 # 65536, fixed by the format
KEY_SIZE: Final[int] = 32 # crypto_secretstream_xchacha20poly1305_KEYBYTES
# magic(8s) version(B) aead_id(B) kdf_id(B) salt(16s) memlimit(Q) opslimit(Q)
# parallelism(B) digest(32s) reserved_signature(64s) — producer's `_HEADER_STRUCT`.
HEADER_STRUCT: Final[struct.Struct] = struct.Struct('>8sBBB16sQQB32s64s')
HEADER_SIZE: Final[int] = HEADER_STRUCT.size # 140
# --- libsodium constants the framing depends on ---
# Not producer constants, but constants of the primitive the producer uses. Kept as literals so
# this module stays import-free; `tests/laborious/utils/bundle/test_format.py` asserts they equal
# the values `nacl.bindings` reports, which is what makes the literals safe.
STREAM_HEADER_SIZE: Final[int] = 24 # crypto_secretstream_xchacha20poly1305_HEADERBYTES
ABYTES: Final[int] = 17 # crypto_secretstream_xchacha20poly1305_ABYTES
# The pre-upload check prefix of `08-security-and-encryption.md` § 5.3: enough bytes to hold the
# header, the stream header and one full chunk. Derived from constants, never read from the file.
PREFIX_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + CHUNK_SIZE + ABYTES # 65717
# The digest covers the chunk stream only. The producer computes `sha256(ciphertext)` over the
# joined chunks and *then* writes `header + stream_header + ciphertext`, so on disk the covered
# range starts after both headers. Verified against a bundle produced by the producer's own code:
# `sha256(file[164:])` matches the header field and `sha256(file[140:])` does not.
DIGEST_OFFSET: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE # 164
MIN_FILE_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE
SUPPORTED_FORMAT_VERSIONS: Final[frozenset[int]] = frozenset({FORMAT_VERSION})
SUPPORTED_AEAD_IDS: Final[frozenset[int]] = frozenset({AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM})
SUPPORTED_KDF_IDS: Final[frozenset[int]] = frozenset({KDF_ID_ARGON2ID})
_HASH_BLOCK_SIZE: Final[int] = 1024 * 1024
class BundleFormatError(Exception):
"""The bytes are not a `.sientia` bundle this reader implements.
Raised for every header-level rejection: wrong magic, unknown version, unknown primitive ids,
cost parameters above the ceiling, a reserved field carrying a value, a file too short to hold
a header, or a chunk stream whose digest does not match the header's. `reader.py` turns this
into a gate-2 rejection; nothing here knows about gates or import steps.
"""
@dataclass(frozen=True)
class BundleHeader:
"""The 140-byte plaintext header, parsed and validated."""
format_version: int
aead_id: int
kdf_id: int
salt: bytes
kdf_memlimit_bytes: int
kdf_opslimit: int
kdf_parallelism: int
ciphertext_digest: bytes
reserved_signature: bytes
def parse_header(raw: bytes) -> BundleHeader:
"""Parse and validate the plaintext header.
Everything the reader trusts about the file's shape is decided here, before any offset past
the header is used and before a key is derived.
Args:
raw: at least `HEADER_SIZE` bytes read from the start of the object.
Returns:
BundleHeader: the validated header.
Raises:
BundleFormatError: on any invalid or unsupported field.
"""
if len(raw) < HEADER_SIZE:
raise BundleFormatError(
f'file is shorter than the {HEADER_SIZE}-byte header ({len(raw)} bytes available)'
)
(
magic,
format_version,
aead_id,
kdf_id,
salt,
kdf_memlimit_bytes,
kdf_opslimit,
kdf_parallelism,
ciphertext_digest,
reserved_signature,
) = HEADER_STRUCT.unpack(raw[:HEADER_SIZE])
if magic != MAGIC:
raise BundleFormatError('file does not start with the expected bundle marker')
if format_version not in SUPPORTED_FORMAT_VERSIONS:
raise BundleFormatError(f'unsupported bundle format version: {format_version}')
if aead_id not in SUPPORTED_AEAD_IDS:
raise BundleFormatError(f'unsupported encryption identifier: {aead_id}')
if kdf_id not in SUPPORTED_KDF_IDS:
raise BundleFormatError(f'unsupported key derivation identifier: {kdf_id}')
if kdf_memlimit_bytes > MAX_KDF_MEMLIMIT:
raise BundleFormatError(
f'declared key derivation memory {kdf_memlimit_bytes} exceeds the '
f'{MAX_KDF_MEMLIMIT} ceiling'
)
if kdf_opslimit > MAX_KDF_OPSLIMIT:
raise BundleFormatError(
f'declared key derivation passes {kdf_opslimit} exceeds the {MAX_KDF_OPSLIMIT} ceiling'
)
if kdf_memlimit_bytes <= 0 or kdf_opslimit <= 0:
raise BundleFormatError('key derivation parameters must be positive')
if kdf_parallelism != PARALLELISM:
raise BundleFormatError(
f'reserved parallelism field must be {PARALLELISM}, found {kdf_parallelism}'
)
if reserved_signature != bytes(RESERVED_SIGNATURE_SIZE):
raise BundleFormatError('reserved signature field must be empty in this format version')
return BundleHeader(
format_version=format_version,
aead_id=aead_id,
kdf_id=kdf_id,
salt=salt,
kdf_memlimit_bytes=kdf_memlimit_bytes,
kdf_opslimit=kdf_opslimit,
kdf_parallelism=kdf_parallelism,
ciphertext_digest=ciphertext_digest,
reserved_signature=reserved_signature,
)
def read_header(path: Path | str) -> BundleHeader:
"""Read and validate the header of the file at `path`.
Raises:
BundleFormatError: when the file cannot hold a header plus a stream header, or when the
header itself is invalid.
"""
file_path = Path(path)
size = file_path.stat().st_size
if size < MIN_FILE_SIZE:
raise BundleFormatError(
f'file is {size} bytes, smaller than the {MIN_FILE_SIZE} bytes a bundle needs for its '
'header and encrypted stream header'
)
with file_path.open('rb') as handle:
return parse_header(handle.read(HEADER_SIZE))
def compute_ciphertext_digest(path: Path | str) -> bytes:
"""SHA-256 of the chunk stream — the bytes from `DIGEST_OFFSET` to the end of the file."""
digest = hashlib.sha256()
with Path(path).open('rb') as handle:
handle.seek(DIGEST_OFFSET)
while block := handle.read(_HASH_BLOCK_SIZE):
digest.update(block)
return digest.digest()
def verify_ciphertext_digest(path: Path | str, header: BundleHeader) -> None:
"""Compare the file's chunk-stream digest against the header's.
Raises:
BundleFormatError: when they differ — a truncated, overwritten or corrupted transfer. This
runs before any key derivation, so a damaged file never costs an Argon2id pass.
"""
if compute_ciphertext_digest(path) != header.ciphertext_digest:
raise BundleFormatError(
'the encrypted content does not match the digest recorded in the bundle header'
)
def file_digest(path: Path | str) -> str:
"""Hex SHA-256 of the whole file, for the `expected_digest` comparison of gate 1."""
digest = hashlib.sha256()
with Path(path).open('rb') as handle:
while block := handle.read(_HASH_BLOCK_SIZE):
digest.update(block)
return digest.hexdigest()

View File

@@ -0,0 +1,473 @@
"""Reading a `.sientia` bundle: the seven gates, in order, in one function.
`BundleReader.open()` is the only entry point. Gate ordering is a property of that function rather
than of the caller, so it cannot be reordered by accident and can be unit-tested without Temporal:
1. object-level limits (`check_object_limits`, called by the download activity before the body is
fetched — the only gate that does not need the file)
2. header validation, including the chunk-stream digest, before any key derivation
3. AEAD decryption — the integrity check; a wrong password and a tampered file fail identically
4. archive inspection over the central directory, writing nothing to disk
5. extraction into a fresh, isolated directory
6. structure and schema of the extracted tree
7. content policy over the file set
Nothing here logs, notifies or emits a metric: rejections are raised as `BundleGateError` and the
activity layer owns the one failure channel.
"""
from __future__ import annotations
import json
import os
import shutil
import stat
import time
import unicodedata
import zipfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import nacl.bindings as sodium
import nacl.exceptions
from laborious.utils.bundle import format as fmt
from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep
# Gate 6 — the layout `bundle.build_zip` writes and `model_repository.py` reads back.
METADATA_FILE_NAME = 'metadata.json'
ARTIFACTS_DIR_NAME = 'artifacts'
REQUIRED_TOP_LEVEL_KEYS = frozenset({'parameters', 'metadata'})
REQUIRED_METADATA_FIELDS = (
'model_name',
'experiment_name',
'run_id',
'model_version',
'model_project',
'export_timestamp',
)
REQUIRED_PARAMETERS = ('target_variable',)
REQUIRED_MODEL_DIRS = ('prediction_model', 'data_model')
REQUIRED_MODEL_FILES = ('MLmodel', 'model.pkl')
# The import record's `experiment_name` column is `VARCHAR(50)` with a length-3 floor, and the name
# is written onto the record before provisioning starts — so the column's limits are enforced here,
# not three activities later as a constraint violation.
EXPERIMENT_NAME_MIN_LENGTH = 3
EXPERIMENT_NAME_MAX_LENGTH = 50
# Gate 7 — an MLflow 2.x artifact tree for this platform: the model directories' known files, the
# model card, the CSVs the training template logs and the bundle's own metadata.
ALLOWED_FILE_NAMES = frozenset(
{
'MLmodel',
'metadata.json',
'model.pkl',
'conda.yaml',
'python_env.yaml',
'requirements.txt',
'model_card.json',
'model_card.svg',
}
)
ALLOWED_FILE_SUFFIXES = frozenset({'.csv', '.json', '.yaml', '.yml', '.txt', '.pkl', '.svg', '.md'})
_DECRYPT_FAILURE_MESSAGE = (
'could not open this bundle: wrong password, or the file is damaged or was altered'
)
class BundleGateError(Exception):
"""A gate rejected the bundle.
Carries the gate number and the import step that gate maps onto, so the workflow never has to
guess which step to record.
"""
def __init__(self, message: str, *, gate: int, step: ImportStep | None = None) -> None:
super().__init__(message)
self.gate = gate
self.step = step if step is not None else GATE_TO_STEP[gate]
@dataclass(frozen=True)
class BundleLimits:
"""Ceilings enforced by gates 1 and 4. Values come from `build_import_config()`."""
max_object_bytes: int = 1024 * 1024 * 1024
max_entries: int = 5000
max_uncompressed_bytes: int = 4 * 1024 * 1024 * 1024
max_compression_ratio: float = 200.0
object_prefix: str = 'imported_models/'
object_suffix: str = '.sientia'
@dataclass(frozen=True)
class OpenedBundle:
"""What a bundle yields once every gate has passed. Nothing secret is carried."""
extracted_dir: Path
metadata: dict[str, Any] = field(default_factory=dict)
parameters: dict[str, Any] = field(default_factory=dict)
digest: str = ''
header: fmt.BundleHeader | None = None
# How long Argon2id took, so the worker's most expensive step is measurable without the
# activity layer having to reach inside the reader.
kdf_seconds: float = 0.0
class BundleReader:
"""Opens `.sientia` bundles produced by the exporter in `sientia-projects-templates`."""
def __init__(self, limits: BundleLimits | None = None) -> None:
self.limits = limits or BundleLimits()
# ------------------------------------------------------------------ gate 1
def check_object_limits(self, *, size: int, key: str) -> None:
"""Gate 1, object level: reject before the object's body is fetched.
Args:
size: the object's size as reported by `stat_object`.
key: the object key.
Raises:
BundleGateError: oversized object, key outside the import prefix, or wrong suffix.
"""
if size > self.limits.max_object_bytes:
raise BundleGateError(
f'uploaded object is {size} bytes, above the '
f'{self.limits.max_object_bytes} byte limit',
gate=1,
)
if size <= 0:
raise BundleGateError('uploaded object is empty', gate=1)
if not key.startswith(self.limits.object_prefix):
raise BundleGateError('uploaded object is not under the import prefix', gate=1)
if not key.endswith(self.limits.object_suffix):
raise BundleGateError(
f'uploaded object does not end in {self.limits.object_suffix}', gate=1
)
# --------------------------------------------------------------------- open
def open(
self,
encrypted_path: Path | str,
password: str,
extract_root: Path | str,
work_dir: Path | str | None = None,
) -> OpenedBundle:
"""Run gates 2 to 7 in order and return the opened bundle.
Args:
encrypted_path: the downloaded `.sientia` file.
password: the plaintext bundle password. Never stored, never logged, never returned.
extract_root: the directory to extract into. It MUST NOT exist yet (gate 5).
work_dir: where the decrypted zip is written. Defaults to `extract_root`'s parent.
Returns:
OpenedBundle: extracted tree, `metadata`, `parameters`, digest and header.
Raises:
BundleGateError: with the gate that rejected the bundle and its import step.
"""
encrypted = Path(encrypted_path)
root = Path(extract_root)
staging = Path(work_dir) if work_dir is not None else root.parent
header = self._gate_2_header(encrypted)
zip_path, kdf_seconds = self._gate_3_decrypt(encrypted, password, header, staging)
try:
self._gate_4_inspect(zip_path)
self._gate_5_extract(zip_path, root)
metadata, parameters = self._gate_6_structure(root)
self._gate_7_content_policy(root)
except BaseException:
# A rejection after decryption leaves nothing extracted behind (gate 5's contract) and
# never leaves a decrypted zip on the worker for the next import to trip over.
shutil.rmtree(root, ignore_errors=True)
zip_path.unlink(missing_ok=True)
raise
zip_path.unlink(missing_ok=True)
return OpenedBundle(
extracted_dir=root,
metadata=metadata,
parameters=parameters,
digest=header.ciphertext_digest.hex(),
header=header,
kdf_seconds=kdf_seconds,
)
# ------------------------------------------------------------------ gate 2
def _gate_2_header(self, encrypted: Path) -> fmt.BundleHeader:
"""Parse the header and check the chunk-stream digest, before deriving anything."""
try:
header = fmt.read_header(encrypted)
fmt.verify_ciphertext_digest(encrypted, header)
except fmt.BundleFormatError as error:
raise BundleGateError(str(error), gate=2) from error
except OSError as error:
raise BundleGateError('downloaded bundle could not be read', gate=2) from error
return header
# ------------------------------------------------------------------ gate 3
def _gate_3_decrypt(
self, encrypted: Path, password: str, header: fmt.BundleHeader, staging: Path
) -> tuple[Path, float]:
"""Derive the key with the header's parameters and decrypt the stream to a file.
The AEAD *is* the integrity check: a wrong password and a tampered byte fail here with the
same message, which never claims to distinguish them.
"""
staging.mkdir(parents=True, exist_ok=True)
zip_path = staging / 'bundle.zip'
try:
started = time.monotonic()
key = self._derive_key(password, header)
kdf_seconds = time.monotonic() - started
self._pull_stream(encrypted, key, zip_path)
except BundleGateError:
zip_path.unlink(missing_ok=True)
raise
except (nacl.exceptions.CryptoError, ValueError, RuntimeError) as error:
zip_path.unlink(missing_ok=True)
raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3) from error
except OSError as error:
zip_path.unlink(missing_ok=True)
raise BundleGateError('decrypted bundle could not be written', gate=3) from error
return zip_path, kdf_seconds
@staticmethod
def _derive_key(password: str, header: fmt.BundleHeader) -> bytes:
"""Argon2id over the NFC-normalised password, with the header's own cost parameters.
The password is normalised and **not** trimmed, matching the producer's `_derive_key`: the
same password typed on another OS must derive the same key, while a trailing space is part
of the secret.
"""
normalized = unicodedata.normalize('NFC', password)
return sodium.crypto_pwhash_alg(
fmt.KEY_SIZE,
normalized.encode('utf-8'),
header.salt,
header.kdf_opslimit,
header.kdf_memlimit_bytes,
sodium.crypto_pwhash_ALG_ARGON2ID13,
)
@staticmethod
def _pull_stream(encrypted: Path, key: bytes, zip_path: Path) -> None:
"""Decrypt the chunk stream into `zip_path`, requiring the FINAL tag."""
chunk_size = fmt.CHUNK_SIZE + fmt.ABYTES
state = sodium.crypto_secretstream_xchacha20poly1305_state()
with encrypted.open('rb') as source, zip_path.open('wb') as target:
source.seek(fmt.HEADER_SIZE)
stream_header = source.read(fmt.STREAM_HEADER_SIZE)
sodium.crypto_secretstream_xchacha20poly1305_init_pull(state, stream_header, key)
last_tag: int | None = None
while chunk := source.read(chunk_size):
plaintext, last_tag = sodium.crypto_secretstream_xchacha20poly1305_pull(
state, chunk
)
target.write(plaintext)
if last_tag != sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL:
# A stream that stops without its FINAL tag was cut short. The digest already catches
# truncation of a stored file; this catches a stream that was framed to look complete.
raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3)
# ------------------------------------------------------------------ gate 4
def _gate_4_inspect(self, zip_path: Path) -> None:
"""Inspect the central directory only. Nothing is written to disk by this gate."""
try:
with zipfile.ZipFile(zip_path) as archive:
entries = archive.infolist()
except zipfile.BadZipFile as error:
raise BundleGateError('bundle content is not a readable archive', gate=4) from error
if len(entries) > self.limits.max_entries:
raise BundleGateError(
f'archive declares {len(entries)} entries, above the '
f'{self.limits.max_entries} limit',
gate=4,
)
if not entries:
raise BundleGateError('archive is empty', gate=4)
uncompressed = sum(entry.file_size for entry in entries)
compressed = sum(entry.compress_size for entry in entries)
if uncompressed > self.limits.max_uncompressed_bytes:
raise BundleGateError(
f'archive declares {uncompressed} uncompressed bytes, above the '
f'{self.limits.max_uncompressed_bytes} limit',
gate=4,
)
ratio = uncompressed / max(compressed, 1)
if ratio > self.limits.max_compression_ratio:
raise BundleGateError(
f'archive compression ratio {ratio:.1f} is above the '
f'{self.limits.max_compression_ratio} limit',
gate=4,
)
for entry in entries:
self._check_entry_name(entry.filename)
self._check_entry_mode(entry)
@staticmethod
def _check_entry_name(name: str) -> None:
"""Reject absolute names and traversal, from the declared name alone."""
normalised = name.replace('\\', '/')
if normalised.startswith('/') or (len(normalised) > 1 and normalised[1] == ':'):
raise BundleGateError('archive declares an absolute entry path', gate=4)
if any(part == '..' for part in normalised.split('/')):
raise BundleGateError('archive declares an entry escaping its own tree', gate=4)
@staticmethod
def _check_entry_mode(entry: zipfile.ZipInfo) -> None:
"""Reject symlinks and anything that is neither a regular file nor a directory."""
mode = entry.external_attr >> 16
if stat.S_IFMT(mode) == 0:
# No file-type bits stored. This is the normal case for two kinds of entry: a zip
# written by a tool that records no Unix mode at all (`mode == 0`), and
# `ZipFile.writestr`, which stores permissions only — the producer's `metadata.json`
# arrives as `0o600 << 16`. Nothing to check: the member is extracted as a plain file
# with the worker's own permissions either way.
return
if stat.S_ISLNK(mode):
raise BundleGateError('archive declares a symbolic link', gate=4)
if not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)):
raise BundleGateError('archive declares an entry that is not a regular file', gate=4)
# ------------------------------------------------------------------ gate 5
def _gate_5_extract(self, zip_path: Path, root: Path) -> None:
"""Extract into a fresh directory, containing every member by resolved path.
A pre-existing extraction root is a failure rather than something to clear: it means
another import, or a previous attempt, owns that path.
"""
if root.exists():
raise BundleGateError('extraction directory already exists', gate=5)
root.mkdir(parents=True)
resolved_root = root.resolve()
try:
with zipfile.ZipFile(zip_path) as archive:
for entry in archive.infolist():
self._extract_member(archive, entry, root, resolved_root)
except BundleGateError:
shutil.rmtree(root, ignore_errors=True)
raise
except (OSError, zipfile.BadZipFile) as error:
shutil.rmtree(root, ignore_errors=True)
raise BundleGateError('bundle content could not be unpacked', gate=5) from error
@staticmethod
def _extract_member(
archive: zipfile.ZipFile, entry: zipfile.ZipInfo, root: Path, resolved_root: Path
) -> None:
"""Write one member, re-checking containment on the resolved path."""
name = entry.filename.replace('\\', '/')
target = root / name
resolved = Path(os.path.realpath(target))
if resolved != resolved_root and resolved_root not in resolved.parents:
raise BundleGateError('archive member resolves outside the extraction root', gate=5)
if entry.is_dir():
resolved.mkdir(parents=True, exist_ok=True)
return
if resolved.exists():
raise BundleGateError('archive member would overwrite an existing file', gate=5)
resolved.parent.mkdir(parents=True, exist_ok=True)
# Written as a plain file: the archive's stored permissions are never restored, so no
# member can arrive executable.
with archive.open(entry) as source, resolved.open('wb') as sink:
shutil.copyfileobj(source, sink)
# ------------------------------------------------------------------ gate 6
def _gate_6_structure(self, root: Path) -> tuple[dict[str, Any], dict[str, Any]]:
"""Validate the extracted layout and the metadata schema."""
metadata_path = root / METADATA_FILE_NAME
artifacts_dir = root / ARTIFACTS_DIR_NAME
if not metadata_path.is_file():
raise BundleGateError('bundle has no metadata document at its root', gate=6)
if not artifacts_dir.is_dir():
raise BundleGateError('bundle has no artifacts directory', gate=6)
document = self._load_metadata_document(metadata_path)
metadata = document['metadata']
parameters = document['parameters']
self._check_metadata_fields(metadata)
self._check_parameters(parameters)
self._check_model_directories(artifacts_dir)
return metadata, parameters
@staticmethod
def _load_metadata_document(metadata_path: Path) -> dict[str, Any]:
"""Read `metadata.json` and check its two top-level keys."""
try:
document = json.loads(metadata_path.read_text(encoding='utf-8'))
except (json.JSONDecodeError, UnicodeDecodeError, OSError) as error:
raise BundleGateError('bundle metadata is not readable JSON', gate=6) from error
if not isinstance(document, dict) or set(document) != REQUIRED_TOP_LEVEL_KEYS:
raise BundleGateError(
'bundle metadata does not carry exactly the parameters and metadata blocks', gate=6
)
if not isinstance(document['metadata'], dict) or not isinstance(
document['parameters'], dict
):
raise BundleGateError('bundle metadata blocks are not objects', gate=6)
return document
@staticmethod
def _check_metadata_fields(metadata: dict[str, Any]) -> None:
"""Every origin field the import records must be present and usable."""
for field_name in REQUIRED_METADATA_FIELDS:
value = metadata.get(field_name)
if not isinstance(value, str) or not value.strip():
raise BundleGateError(f'bundle metadata is missing the {field_name} field', gate=6)
experiment_name = metadata['experiment_name']
if not (EXPERIMENT_NAME_MIN_LENGTH <= len(experiment_name) <= EXPERIMENT_NAME_MAX_LENGTH):
raise BundleGateError(
"bundle experiment name does not fit the platform's limits", gate=6
)
@staticmethod
def _check_parameters(parameters: dict[str, Any]) -> None:
"""`target_variable` is what `model_config.target` is built from; nothing else supplies it."""
for name in REQUIRED_PARAMETERS:
value = parameters.get(name)
if value is None or not str(value).strip():
raise BundleGateError(f'bundle parameters are missing {name}', gate=6)
@staticmethod
def _check_model_directories(artifacts_dir: Path) -> None:
"""`prediction_model` and `data_model` must each be a loadable MLflow model directory."""
for directory in REQUIRED_MODEL_DIRS:
model_dir = artifacts_dir / directory
if not model_dir.is_dir():
raise BundleGateError(f'bundle has no {directory} artifacts', gate=6)
for required in REQUIRED_MODEL_FILES:
if not (model_dir / required).is_file():
raise BundleGateError(f'bundle {directory} artifacts are incomplete', gate=6)
# ------------------------------------------------------------------ gate 7
def _gate_7_content_policy(self, root: Path) -> None:
"""Only the file kinds an MLflow artifact tree calls for may be present."""
for path in sorted(root.rglob('*')):
if path.is_dir():
continue
if not path.is_file():
raise BundleGateError('bundle contains an entry that is not a file', gate=7)
if path.name in ALLOWED_FILE_NAMES:
continue
if path.suffix.lower() in ALLOWED_FILE_SUFFIXES:
continue
raise BundleGateError('bundle contains a file a model export should not carry', gate=7)

View File

@@ -0,0 +1,55 @@
"""The import pipeline's step vocabulary, and the gate-to-step mapping.
`ImportStep` is the join between three things: the machine-readable step written into the import
record's JSON detail, the key of the error catalog (one code and one sentence per step) and the
value the frontend switches on. It lives in the bundle package rather than in the activities module
because the reader raises with it — `laborious/activities/model_import.py` imports it from here so
there is exactly one definition.
It holds **pipeline** steps only. The two failures that are not steps of the pipeline — failing to
claim the record and failing to write it — are reported by the failure channel under their own
markers (`not_claimable`, `status_write`), deliberately outside this enum so the
catalog-completeness test keeps meaning "every pipeline step has a sentence".
"""
from __future__ import annotations
from enum import StrEnum
from typing import Final
class ImportStep(StrEnum):
"""Where an import stopped. Values are what gets written to the record."""
RECEIVED = 'received'
DOWNLOAD = 'download'
DECRYPTION = 'decryption'
ARCHIVE_INSPECTION = 'archive_inspection'
EXTRACTION = 'extraction'
STRUCTURE_VALIDATION = 'structure_validation'
CONTENT_POLICY = 'content_policy'
EXPERIMENT_CREATION = 'experiment_creation'
ARTIFACT_UPLOAD = 'artifact_upload'
REGISTRATION = 'registration'
MODEL_DOCUMENT = 'model_document'
CLEANUP = 'cleanup'
# The seven gates map onto the first seven steps, with gates 1 and 2 sharing one: both say "these
# are not the bytes that were sent" — an object over the size ceiling, a key with the wrong shape, a
# digest that disagrees with what the uploader computed, a wrong magic, an unsupported version, a
# chunk stream that does not match the header's digest. The catalog's `download` sentence ("The
# uploaded file could not be read, or it is not the file that was sent. Upload it again.") is true
# for all of them, and none of them may borrow the `decryption` sentence, which implicates the
# password.
GATE_TO_STEP: Final[dict[int, ImportStep]] = {
1: ImportStep.DOWNLOAD,
2: ImportStep.DOWNLOAD,
3: ImportStep.DECRYPTION,
4: ImportStep.ARCHIVE_INSPECTION,
5: ImportStep.EXTRACTION,
6: ImportStep.STRUCTURE_VALIDATION,
7: ImportStep.CONTENT_POLICY,
}
GATE_COUNT: Final[int] = 7

View File

@@ -0,0 +1,171 @@
import json
import tempfile
from os import getenv, path
from typing import Any
def build_mlflow_config() -> dict[str, Any]:
"""
Build MLFlow server configuration from environment variables.
This function constructs an MLFlow configuration dictionary from
environment variables with sensible defaults for local development.
It handles server connection and authentication parameters.
Environment Variables:
MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
MLFLOW_PORT: MLFlow server port (default: 5080)
MLFLOW_USERNAME: MLFlow username (default: aignosi)
MLFLOW_PASSWORD: MLFlow password (default: aignosi)
Returns:
dict: MLFlow configuration dictionary with all required parameters
"""
return {
'host': getenv('MLFLOW_HOST', 'http://localhost'),
'port': int(getenv('MLFLOW_PORT', '5080')),
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
}
def build_opc_config() -> dict[str, Any]:
"""
Build OPC server configuration from environment variables.
This function constructs an OPC server configuration dictionary from
environment variables. It supports both single server and multi-server
configurations with flexible parameter handling.
Environment Variables:
OPC_CONFIG: JSON string containing multiple OPC server configurations
OPC_ID: OPC server ID (fallback, default: 1)
OPC_URL: Single OPC server URL (fallback, default: opc.tcp://localhost:4840)
OPC_SERVER_URI: Single OPC server URI (fallback, default: opc.tcp://localhost:4840)
OPC_CERT_PATH: Client certificate path (fallback, default: None)
OPC_PRIVATE_KEY_PATH: Client private key path (fallback, default: None)
OPC_SERVER_CERT_PATH: Server certificate path (fallback, default: None)
OPC_RECONNECTION_INTERVAL: Reconnection interval in seconds (fallback, default: 120)
Returns:
dict: OPC server configuration dictionary
"""
opc_raw = getenv('OPC_CONFIG', None)
if opc_raw:
return json.loads(opc_raw)
return {
getenv('OPC_ID', '1'): {
'id': getenv('OPC_ID', '1'),
'server_name': getenv('OPC_SERVER_NAME', 'default_server'),
'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'),
'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'),
'cert_path': getenv('OPC_CERT_PATH', None),
'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None),
'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None),
'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')),
}
}
def build_minio_config() -> dict[str, Any]:
"""
Build MinIO (S3-compatible) configuration from environment variables.
Environment Variables:
MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000)
MINIO_ACCESS_KEY: Access key (default: minioadmin)
MINIO_SECRET_KEY: Secret key (default: minioadmin)
MINIO_REGION: Region name for S3 client (default: us-east-1)
MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious)
MINIO_SECURE: Whether to use HTTPS (default: false)
Returns:
dict: MinIO configuration dictionary
"""
return {
'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'),
'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'),
'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'),
'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')),
'secure': getenv('MINIO_SECURE', 'false') == 'true',
}
def build_import_config() -> dict[str, Any]:
"""
Build the `.sientia` model-import configuration from environment variables.
The password key is read as the raw base64 string and **not** validated here: the worker must
start whether or not `IMPORT_MODEL_ENABLED` is set, and a missing or malformed key is an import
failure at the decryption step (`laborious.utils.import_password.load_envelope_key`), not a
boot failure.
Environment Variables:
IMPORT_PASSWORD_KEY: base64 of the 32-byte AES-256-GCM key that wraps bundle passwords
IMPORT_BUNDLE_BUCKET: bucket holding the uploaded bundles (default: sientia)
IMPORT_BUNDLE_PREFIX: key prefix the uploads must sit under (default: imported_models/)
IMPORT_WORK_DIR: worker scratch directory (default: <tempdir>/sientia-import)
IMPORT_MAX_OBJECT_BYTES: gate 1 object size ceiling (default: 1 GiB)
IMPORT_MAX_ARCHIVE_ENTRIES: gate 4 entry count ceiling (default: 5000)
IMPORT_MAX_UNCOMPRESSED_BYTES: gate 4 uncompressed size ceiling (default: 4 GiB)
IMPORT_MAX_COMPRESSION_RATIO: gate 4 compression ratio ceiling (default: 200)
IMPORT_MODELS_COLLECTION: MongoDB collection holding the model listing (default: models)
IMPORT_BUNDLE_RETENTION_DAYS: lifecycle expiry for the uploaded bundle (default: 7)
Returns:
dict: import configuration dictionary
"""
prefix = getenv('IMPORT_BUNDLE_PREFIX', 'imported_models/')
return {
'password_key': getenv('IMPORT_PASSWORD_KEY'),
'bucket': getenv('IMPORT_BUNDLE_BUCKET', 'sientia'),
'prefix': prefix if prefix.endswith('/') else f'{prefix}/',
'work_dir': getenv('IMPORT_WORK_DIR', path.join(tempfile.gettempdir(), 'sientia-import')),
'max_object_bytes': int(getenv('IMPORT_MAX_OBJECT_BYTES', str(1024 * 1024 * 1024))),
'max_entries': int(getenv('IMPORT_MAX_ARCHIVE_ENTRIES', '5000')),
'max_uncompressed_bytes': int(
getenv('IMPORT_MAX_UNCOMPRESSED_BYTES', str(4 * 1024 * 1024 * 1024))
),
'max_compression_ratio': float(getenv('IMPORT_MAX_COMPRESSION_RATIO', '200')),
'models_collection': getenv('IMPORT_MODELS_COLLECTION', 'models'),
'retention_days': int(getenv('IMPORT_BUNDLE_RETENTION_DAYS', '7')),
}
def build_import_status_config() -> dict[str, Any]:
"""
Build the connection to the import log database — the POSTGRES_* server, another database.
The import log is `public.experiment_run` in the database owned by the Spring Boot BFF
(`sientia-core-mlops-bff`), which sits on the same Postgres server as the `sientia` database
Laborious itself uses. Only the database name differs, so only the database name has a variable
of its own; host, port, user and password come from POSTGRES_*. The table and schema are
constants of the code, never inputs.
The four IMPORT_STATUS_DB_* connection variables remain as overrides, for a deployment where the
log genuinely lives on another server — the e2e suite is exactly that, two containers.
Environment Variables:
IMPORT_STATUS_DB_NAME: database (default: sientia-core-mlops-bff)
POSTGRES_HOST: server host, unless IMPORT_STATUS_DB_HOST overrides it (default: localhost)
POSTGRES_PORT: server port, unless IMPORT_STATUS_DB_PORT overrides it (default: 5432)
POSTGRES_USER: user, unless IMPORT_STATUS_DB_USER overrides it (default: sientia)
POSTGRES_PASSWORD: password, unless IMPORT_STATUS_DB_PASSWORD overrides it (default: empty)
Returns:
dict: import status database configuration dictionary
"""
# Written as `override or fallback or default` rather than
# `getenv(override) or getenv(fallback, default)`: inside an `or`, the two-argument `getenv`
# widens to `str | None` for mypy, so the `int(...)` below fails to type-check. Putting the
# default last as its own term keeps every value a plain `str` and reads in precedence order.
return {
'host': getenv('IMPORT_STATUS_DB_HOST') or getenv('POSTGRES_HOST') or 'localhost',
'port': int(getenv('IMPORT_STATUS_DB_PORT') or getenv('POSTGRES_PORT') or '5432'),
'dbname': getenv('IMPORT_STATUS_DB_NAME') or 'sientia-core-mlops-bff',
'user': getenv('IMPORT_STATUS_DB_USER') or getenv('POSTGRES_USER') or 'sientia',
'password': getenv('IMPORT_STATUS_DB_PASSWORD') or getenv('POSTGRES_PASSWORD') or '',
'schema': 'public',
}

View File

@@ -0,0 +1,34 @@
from typing import Any
from pandas import DataFrame
DEFAULT_MAX_DEBUG_DATAFRAME_ROWS = 100
def build_dataframe_debug_message(
message: str,
data: Any,
max_rows: int = DEFAULT_MAX_DEBUG_DATAFRAME_ROWS,
) -> str:
"""
Build a safe debug message for dataframe payloads
Args:
- message (str): Base message to identify the logged payload
- data (Any): Payload to evaluate for dataframe-aware logging
- max_rows (int): Maximum dataframe row count allowed for full payload logging
Return:
Formatted debug message with full dataframe content or compact summary
"""
if not isinstance(data, DataFrame):
return f'{message} {data}'
rows = data.shape[0]
if rows <= max_rows:
return f'{message}\n{data.to_csv()}'
return (
f'{message} skipped because dataframe has {rows} rows '
f'(max: {max_rows}). Shape: {data.shape}'
)

View File

View File

@@ -0,0 +1,48 @@
from pandas import DataFrame
def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool:
"""
Filter to check if specific variables contain null values.
This function examines a DataFrame to determine if any of the specified variables
contain null (NaN) values. It returns True if null values are found for any of
the specified variables, False otherwise.
Args:
data (DataFrame): The pandas DataFrame to be examined. Must contain columns
named 'variable' and 'value'.
config (dict): Configuration dictionary containing the following key:
- variables (list): List of variable names to check for null values
Returns:
bool: True if any of the specified variables contain null values,
False if none of the specified variables contain null values.
"""
if data.empty:
return False
return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty
def filter_empty_data(data: DataFrame, _config: dict) -> bool:
"""
Filter to check if the DataFrame is empty.
This function determines whether the provided DataFrame contains any data.
It's a simple utility function that can be used in conditional logic to
handle cases where no data is available.
Args:
data (DataFrame): The pandas DataFrame to be checked for emptiness.
_config (dict): Configuration dictionary (unused in this function).
The underscore prefix indicates this parameter is required for
interface consistency but not used in the implementation.
Returns:
bool: True if the DataFrame is empty (has no rows), False if it contains data.
"""
return data.empty

View File

@@ -0,0 +1,64 @@
import numpy as np
from pandas import DataFrame
def api_error_filter(response: dict, _config: dict) -> bool:
"""
Filter MLFlow API responses for error conditions.
This function analyzes MLFlow API responses to detect error conditions
and determine if the response should be filtered out due to quality
or reliability issues.
Args:
response: MLFlow API response data (dict)
_config: Filter configuration dictionary
Required keys:
- error_codes (list, optional): List of error codes to detect
- error_keywords (list, optional): List of error keywords to detect
- check_structure (bool, optional): Whether to validate response structure
Returns:
bool: True if data should be filtered (contains errors), False otherwise
"""
if not response:
return True
if not response['success']:
return True
return False
def nan_values_filter(predictions: DataFrame, _config: dict) -> bool:
"""
Filter data for NaN (Not a Number) values.
This function detects NaN values in MLFlow prediction results and
determines if the data quality is sufficient for further processing
or export operations.
Args:
predictions: DataFrame containing prediction data to check for NaN values
_config: Filter configuration dictionary
Required keys:
- max_nan_ratio (float, optional): Maximum allowed NaN value ratio (0.0 to 1.0)
- max_nan_count (int, optional): Maximum allowed NaN value count
- check_nested (bool, optional): Whether to check nested data structures
Returns:
bool: True if data should be filtered (too many NaN values), False otherwise
"""
data = (
predictions.replace({None: np.nan})
.drop(columns=['timestamp'], errors='ignore')
.infer_objects()
)
if data.isna().all().all():
return True
return False

View File

@@ -0,0 +1,98 @@
"""The bundle password envelope: AES-256-GCM, opened inside the activity that uses the password.
The workflow only ever carries `password_envelope`, so no plaintext password reaches the Temporal
event history. This module is imported by the bundle-opening activity and by nothing else — in
particular not by `laborious/workflows/import_model.py`, which a test asserts.
Envelope layout, chosen because WebCrypto gives the Angular frontend AES-GCM natively:
base64( nonce(12 bytes) || ciphertext || GCM tag(16 bytes) )
The key is a 32-byte value held in a Kubernetes Secret and exposed as `IMPORT_PASSWORD_KEY`,
base64-encoded. Neither the key, the envelope nor the recovered password is ever logged, notified,
used as a metric label or written to the import record.
"""
from __future__ import annotations
import base64
import binascii
from typing import Final
from cryptography.exceptions import InvalidTag
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
NONCE_SIZE: Final[int] = 12
TAG_SIZE: Final[int] = 16
KEY_SIZE: Final[int] = 32
MIN_ENVELOPE_SIZE: Final[int] = NONCE_SIZE + TAG_SIZE
class EnvelopeError(Exception):
"""The envelope or its key is unusable.
Deliberately one type for every cause — missing key, wrong key, malformed base64, altered
bytes — because the failure is reported at the decryption step either way. What it must never
be confused with is a wrong *bundle* password: that one comes from the reader's gate 3, and the
two carry different sentences.
"""
def load_envelope_key(raw_key: str | None) -> bytes:
"""Decode the configured envelope key.
Args:
raw_key: the value of `IMPORT_PASSWORD_KEY` — base64 of exactly 32 bytes.
Returns:
bytes: the 32-byte key.
Raises:
EnvelopeError: when the value is absent, not base64, or not 32 bytes long. There is no
fallback interpretation: an unset key never means "the envelope is plaintext".
"""
if not raw_key:
raise EnvelopeError('the import password key is not configured')
try:
key = base64.b64decode(raw_key, validate=True)
except (binascii.Error, ValueError) as error:
raise EnvelopeError('the import password key is not valid base64') from error
if len(key) != KEY_SIZE:
raise EnvelopeError(f'the import password key must be {KEY_SIZE} bytes')
return key
def decrypt_password_envelope(envelope: str, key: bytes) -> str:
"""Recover the bundle password from its envelope.
Args:
envelope: base64 of `nonce || ciphertext || tag`.
key: the 32-byte key from `load_envelope_key`.
Returns:
str: the plaintext password, exactly as it was typed — not trimmed, not case-folded. NFC
normalisation happens at key derivation, where the producer does it.
Raises:
EnvelopeError: malformed envelope, wrong key or altered bytes (AES-GCM authentication).
"""
if not envelope:
raise EnvelopeError('the import request carries no password envelope')
try:
payload = base64.b64decode(envelope, validate=True)
except (binascii.Error, ValueError) as error:
raise EnvelopeError('the password envelope is not valid base64') from error
if len(payload) < MIN_ENVELOPE_SIZE:
raise EnvelopeError('the password envelope is too short to hold a nonce and a tag')
nonce, sealed = payload[:NONCE_SIZE], payload[NONCE_SIZE:]
try:
plaintext = AESGCM(key).decrypt(nonce, sealed, None)
except InvalidTag as error:
raise EnvelopeError('the password envelope could not be authenticated') from error
except ValueError as error:
raise EnvelopeError('the password envelope could not be opened') from error
try:
return plaintext.decode('utf-8')
except UnicodeDecodeError as error:
raise EnvelopeError('the recovered password is not valid UTF-8') from error

View File

View File

@@ -0,0 +1,348 @@
"""
MinIO-backed DataFrame payload for Temporal workflows.
Data is never stored as a pandas ``DataFrame`` field on the dataclass.
Instead, the DataFrame is only provided as an input to:
`from_dataframe` / `from_dataframe_to_dict`.
At build time, the DataFrame is evaluated for its serialized size; if it exceeds
the configured threshold, it is serialized to parquet bytes and uploaded to MinIO.
Otherwise, it is inlined as a Temporal-friendly ``dict``.
"""
import pickle
import re
from collections.abc import Hashable
from dataclasses import dataclass
from datetime import datetime
from io import BytesIO
from os import getenv
from typing import Any, Literal
from pandas import DataFrame, read_parquet
from sientia_do.observability.logger import Logger
from sientia_do.repository.minio_repository import MinioRepository
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now
# Keys that are part of the serialized wire format (not arbitrary metadata).
_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'})
_OBJECT_TIMESTAMP_PATTERN = re.compile(
r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$'
)
OFFLOAD_THRESHOLD_BYTES = int(
float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024
)
# Relative prefix used for storing offloaded prediction datasets in MinIO.
# It is also the root directory for retention cleanup listing.
PREDICTION_DATASETS_PREFIX = 'prediction_datasets'
OperationKind = Literal['initial', 'transform', 'predict']
def _build_object_key(
model_name: str, operation: OperationKind, timestamp: str
) -> tuple[str, str | None]:
"""
Build the MinIO object key and the directory prefix used for retention listing.
Args:
model_name: Registered model name used in the pipeline.
operation: Either initial (pre-transform load) or transform (post-MLFlow transform).
timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME.
Return:
tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root).
"""
# Naming convention:
# - Directory is always `prediction_datasets/<model_name>`
# - Filename follows the retention-parsing pattern
basename = f'{model_name}-{operation}-{timestamp}.parquet'
model_dir = model_name.strip().strip('/')
prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}'
return f'{prefix}/{basename}', prefix
@dataclass
class MinioDataFramePayload:
"""
Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys.
Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`.
Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class.
"""
last_timestamp: str
status: dict[str, Any] | None = None
data: dict[Hashable, Any] | None = None
bucket: str | None = None
object_key: str | None = None
object_prefix: str | None = None
uri: str | None = None
@staticmethod
def _debug(
logger: Logger | None,
message: str,
metadata: dict[str, Any] | None = None,
) -> None:
"""
Emit debug logs only when logger is provided
Args:
- logger (Logger | None): Logger instance used for debug messages
- message (str): Message to be logged
- metadata (dict[str, Any] | None): Optional workflow metadata context
"""
if logger is None:
return
logger.custom_debug(message, metadata)
@classmethod
def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload':
"""
Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization.
Temporal converts dataclass return values into plain dicts when crossing
workflow/activity boundaries. This method rebuilds the typed instance so
that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are
available on the receiving side.
If the argument is already a MinioDataFramePayload, it is returned as-is.
Args:
raw: Dict with keys matching the dataclass fields
(last_timestamp, status, data, bucket, object_key, object_prefix, uri),
or an existing MinioDataFramePayload instance.
Return:
MinioDataFramePayload: Reconstructed (or original) instance.
"""
if isinstance(raw, MinioDataFramePayload):
return raw
return cls(
last_timestamp=raw['last_timestamp'],
status=raw.get('status'),
data=raw.get('data'),
bucket=raw.get('bucket'),
object_key=raw.get('object_key'),
object_prefix=raw.get('object_prefix'),
uri=raw.get('uri'),
)
@staticmethod
def estimate_size_bytes(
df: DataFrame,
metadata: dict[str, Any] | None = None,
logger: Logger | None = None,
) -> int:
"""
Approximate serialized size of the DataFrame as the default-orient dict.
Args:
df: DataFrame whose tabular content size is estimated.
Return:
int: Estimated size in bytes (pickle of dict representation).
"""
try:
size = len(pickle.dumps(df.to_dict()))
except Exception:
size = len(pickle.dumps(df))
MinioDataFramePayload._debug(
logger,
f'DataFrame size: {size} bytes',
metadata,
)
return size
@staticmethod
def parse_object_timestamp(object_key: str) -> datetime | None:
"""
Parse the timestamp embedded in the object key basename (before .parquet).
Args:
object_key: S3/MinIO object key whose basename follows
``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``.
Return:
datetime | None: Parsed UTC-naive datetime from the key, or None if not matched.
"""
basename = object_key.rsplit('/', 1)[-1]
match = _OBJECT_TIMESTAMP_PATTERN.search(basename)
if not match:
return None
try:
return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME)
except ValueError:
return None
def cleanup_prefix(self) -> str | None:
"""
Return True if cleanup is enabled for this payload.
"""
if self.object_key is not None and self.data is None:
return self.object_prefix
return None
def has_data(self) -> bool:
"""
Return True if the payload has some data internally or in MinIO.
"""
return (self.data is not None and self.data != {}) or self.object_key is not None
@classmethod
async def from_dataframe(
cls,
dataframe: DataFrame | None,
minio_repo: MinioRepository,
model_name: str,
operation: OperationKind,
status: dict[str, Any] | None = None,
workflow_metadata: dict | None = None,
last_timestamp: str | None = None,
logger: Logger | None = None,
) -> 'MinioDataFramePayload':
"""
Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO.
The DataFrame is not stored on the returned instance.
Args:
dataframe: Tabular data to evaluate and persist (inline or MinIO).
metadata: Small metadata dict merged into the payload (e.g. success, message).
minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`.
workflow_metadata: Metadata passed to MinIO store for logging/metrics.
model_name: Registered model name used in the object basename.
operation: Either ``initial`` (query load) or ``transform`` (post-transform).
key_prefix: Backward-compatible parameter (currently ignored for object naming).
size_threshold_bytes: Byte limit before offload. When None, the module-level
environment-derived default is used.
Return:
MinioDataFramePayload: Instance with data and/or MinIO fields set.
"""
if dataframe is None or dataframe.empty:
cls._debug(
logger,
'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload',
workflow_metadata,
)
return cls(
data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
)
if last_timestamp is None:
last_timestamp = max(dataframe['timestamp'].values.tolist())
dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger)
cls._debug(
logger,
(
f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes '
f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)'
),
workflow_metadata,
)
if dataframe_size <= OFFLOAD_THRESHOLD_BYTES:
cls._debug(
logger,
'MinioDataFramePayload.from_dataframe using inline payload',
workflow_metadata,
)
return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status)
timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
object_key, object_prefix = _build_object_key(model_name, operation, timestamp)
cls._debug(
logger,
(
'MinioDataFramePayload.from_dataframe offloading payload to MinIO '
f'with key {object_key}'
),
workflow_metadata,
)
# Upload using the relative object key. The upstream repository will
# prefix it internally under its MinIO namespace.
parquet_buffer = BytesIO()
dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
file_bytes = parquet_buffer.getvalue()
upload_result = await minio_repo.upload_file(
file_bytes=file_bytes,
relative_key=object_key,
metadata=workflow_metadata,
)
bucket = minio_repo.bucket
object_key_full = upload_result.get('minio_object_name', object_key)
uri = f's3://{bucket}/{object_key_full}' if bucket else None
cls._debug(
logger,
f'MinioDataFramePayload.from_dataframe upload completed: {uri}',
workflow_metadata,
)
return cls(
data=None,
bucket=bucket,
object_key=object_key_full,
object_prefix=object_prefix,
uri=uri,
last_timestamp=last_timestamp,
status=status,
)
async def retrieve(
self,
minio_repo: MinioRepository,
workflow_metadata: dict[str, Any] | None = None,
logger: Logger | None = None,
) -> DataFrame:
"""
Load parquet from MinIO when object_key is set and populate inline data.
Args:
minio_repo: sientia_do MinioRepository (or compatible) with download_file().
workflow_metadata: Metadata passed to MinIO read for logging/metrics.
Return:
dict[str, Any]: Flat dict with data filled (same keys as to_dict after load).
"""
if self.data is not None:
self._debug(
logger,
'MinioDataFramePayload.retrieve using inline payload data',
workflow_metadata,
)
return DataFrame(self.data)
if not self.has_data():
self._debug(
logger,
'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe',
workflow_metadata,
)
return DataFrame()
self._debug(
logger,
f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}',
workflow_metadata,
)
file_bytes = await minio_repo.download_file(
object_name=self.object_key, metadata=workflow_metadata
)
df = read_parquet(BytesIO(file_bytes))
self._debug(
logger,
f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}',
workflow_metadata,
)
return df

View File

@@ -0,0 +1,32 @@
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.minio_repository import MinioRepository
class MinioManager(SientiaMonitoring):
minio_repository: MinioRepository | None = None
def __init__(
self,
minio_repository: MinioRepository | None = None,
logger: Logger | None = None,
notification_handler: NotificationHandler | None = None,
metrics_controller: MetricsController | None = None,
):
if self.minio_repository is None:
self.minio_repository = minio_repository
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
def close(self) -> None:
"""
Close the MinioManager and clean up resources.
"""
if self.minio_repository is not None:
try:
self.minio_repository.close()
finally:
self.minio_repository = None
SientiaMonitoring.shutdown(self)

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,866 @@
import asyncio
import json
import time
import traceback
from datetime import datetime
from pathlib import Path
from typing import Any
from asyncua import Client
from asyncua.crypto.security_policies import SecurityPolicyBasic256
from asyncua.ua import DataValue, Variant, VariantType
from asyncua.ua.uaerrors import UaStatusCodeError
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from laborious import metrics
# Requested session and secure channel lifetime (ms) before server revision; 10 minutes.
OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS = 10 * 60 * 1000
class OpcClientAlreadyExistsError(RuntimeError):
"""Raised when _create_client is called while self.client is already set."""
class OpcSessionAlreadyConnectedError(RuntimeError):
"""Raised when _open_session is called while a UA session is already open."""
class OpcClientNotInitializedError(RuntimeError):
"""Raised when _open_session is called before _create_client."""
RECONNECTABLE_OPC_BAD_NAMES: frozenset[str] = frozenset(
{
'BadSessionIdInvalid',
'BadSessionClosed',
'BadSessionNotActivated',
'BadSecureChannelIdInvalid',
'BadSecureChannelClosed',
'BadSecureChannelTokenUnknown',
'BadTcpSecureChannelUnknown',
'BadServerNotConnected',
'BadConnectionClosed',
'BadDisconnect',
'BadConnectionRejected',
'BadCommunicationError',
'BadRequestInterrupted',
'BadUnknownResponse',
'BadTimeout',
'BadRequestTimeout',
'BadSequenceNumberInvalid',
'BadSequenceNumberUnknown',
'BadSecurityModeInsufficient',
'BadRequestHeaderInvalid',
'BadInvalidState',
}
)
def _opc_authentication_token_str(client: Client | None) -> str:
"""
Serialize the current OPC UA authentication token (session handle) for logging and metrics.
Return:
str: Token string, or "unknown" if unavailable.
"""
if client is None:
return 'unknown'
try:
proto = client.uaclient.protocol
if proto is None:
return 'unknown'
tok = getattr(proto, 'authentication_token', None)
if tok is None:
return 'unknown'
return str(tok)
except Exception:
return 'unknown'
def _opc_status_from_exception(exc: BaseException) -> str:
"""
Resolve OPC UA status name from an exception, including chained UaStatusCodeError causes.
Args:
exc (BaseException): Raised error from asyncua.
Return:
str: Status class name or generic Python exception name.
"""
current: BaseException | None = exc
while current is not None:
if isinstance(current, UaStatusCodeError):
return type(current).__name__
current = current.__cause__
return type(exc).__name__
def is_reconnectable_opcua_bad(exc: BaseException) -> bool:
"""
Return whether the exception is a Tier-1 OPC UA Bad* that should trigger reconnect.
Args:
exc (BaseException): Raised error from get_node or write_value.
Return:
bool: True if reconnect should be scheduled.
"""
return _opc_status_from_exception(exc) in RECONNECTABLE_OPC_BAD_NAMES
def _model_labels_from_write_metadata(metadata: dict[str, Any] | None) -> dict[str, str]:
"""
Extract model_id and model_name from write metadata for Prometheus labels.
Args:
metadata (dict[str, Any] | None): Context passed into write_data; may omit keys.
Return:
dict[str, str]: Labels model_id and model_name, defaulting to "unknown".
"""
if not metadata:
return {'model_id': 'unknown', 'model_name': 'unknown'}
return {
'model_id': str(metadata.get('model_id', 'unknown')),
'model_name': str(metadata.get('model_name', 'unknown')),
}
data_type_map = {
'float': {
'converter': float,
'opc_type': VariantType.Float,
},
'double': {
'converter': float,
'opc_type': VariantType.Double,
},
'int': {
'converter': int,
'opc_type': VariantType.Int32,
},
'bool': {
'converter': bool,
'opc_type': VariantType.Boolean,
},
'str': {
'converter': str,
'opc_type': VariantType.String,
},
}
class OpcRepository(SientiaMonitoring):
def __init__(
self,
opc_id: str,
url: str,
server_name: str,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
reconnection_interval: int = 60,
server_uri: str | None = None,
cert_path: str | None = None,
private_key_path: str | None = None,
server_cert_path: str | None = None,
):
self.url = url
self.id = opc_id
self.server_name = server_name
self.server_uri = server_uri
self.cert_path = cert_path
self.private_key_path = private_key_path
self.server_cert_path = server_cert_path
self.reconnection_interval = reconnection_interval
self.last_reconnection_time: None | datetime = None
self.disconnection_interval = 10.0
self.notification_handler = notification_handler
self.client: None | Client = None
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.metadata = {
'model_name': '-',
'model_id': '-',
'workflow_name': 'opc_repository',
'schedule_name': '-',
}
self._last_write_mono: float | None = None
self._connection_lock = asyncio.Lock()
self._session_ready = asyncio.Event()
self._reconnect_task: asyncio.Task[None] | None = None
self._allow_reconnect = True
def _opc_debug_tags(self, session_id: str) -> dict[str, str]:
"""
Build Prometheus/log label tags for OPC session-scoped metrics.
Args:
session_id (str): OPC UA session token string.
Return:
dict[str, str]: Labels pod_id, server_name, runtime, opc_server_id, session_id.
"""
return {
'pod_id': str(getattr(self, 'pod_id', 'unknown')),
'server_name': self.server_name,
'runtime': str(getattr(self, 'runtime', 'unknown')),
'opc_server_id': self.id,
'session_id': session_id,
}
def _is_session_open(self) -> bool:
"""
Return whether the asyncua client has an open transport session.
Return:
bool: True when protocol exists and is not closed.
"""
if self.client is None:
return False
try:
proto = self.client.uaclient.protocol
return proto is not None and proto.state != 'closed'
except Exception:
return False
def _reconnection_window_elapsed(self) -> bool:
"""
Return whether enough time has passed since the last reconnect attempt.
Return:
bool: True if a new reconnect is allowed.
"""
if self.last_reconnection_time is None:
return True
return (
datetime.now() - self.last_reconnection_time
).total_seconds() > self.reconnection_interval
def _not_connected_error(self) -> dict[str, Any]:
"""
Build the standard error payload when validate_connection finds no open protocol.
Return:
dict[str, Any]: Notification fields for OPC_CONNECTION_NOT_READY.
"""
return {
'notification_id': f'OPC_CONNECTION_NOT_READY_{self.id}',
'message': f'OPC server {self.id} is not connected',
'block': 'opc_repository',
'level': NotificationLevel.WARNING,
}
async def set_security(self) -> None:
"""
Configure certificates and timeouts on the asyncua client.
Raises:
ValueError: If cert paths or client are missing.
"""
if self.cert_path is None or self.private_key_path is None:
raise ValueError(
'Certificate and private key paths must be provided for secure connection.'
)
cert = Path(self.cert_path)
private_key = Path(self.private_key_path)
server_cert = Path(self.server_cert_path) if self.server_cert_path else None
if self.client is None:
raise ValueError('Client must be initialized before setting security')
self.client.application_uri = self.server_uri
self.info('Setting security...', self.metadata)
await self.client.set_security(
SecurityPolicyBasic256,
certificate=str(cert),
private_key=str(private_key),
server_certificate=str(server_cert) if server_cert else None,
)
self.client.secure_channel_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
self.client.session_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
async def _create_client(self) -> None:
"""
Instantiate the asyncua Client and apply security when configured.
Caller must hold _connection_lock. Does not open a UA session.
Raises:
OpcClientAlreadyExistsError: If self.client is already set.
"""
if self.client is not None:
raise OpcClientAlreadyExistsError(
f'OPC client already exists for server {self.id}; '
'call disconnect() before creating a new client'
)
self.client = Client(self.url, timeout=10, watchdog_intervall=50) # type: ignore[attr-defined]
self.client.name = self.pod_id
self.client.application_name = self.pod_id
pod_uri = self.pod_id.replace('-', ':')
self.client.application_uri = pod_uri
self.client.product_uri = pod_uri
if self.cert_path:
await self.set_security()
async def _open_session(self) -> tuple[bool, dict[str, Any]]:
"""
Open the OPC UA session on the existing client.
Caller must hold _connection_lock.
Raises:
OpcClientNotInitializedError: If self.client is None.
OpcSessionAlreadyConnectedError: If a session is already open.
Return:
tuple[bool, dict[str, Any]]: Success flag and error payload on connect failure.
"""
if self.client is None:
raise OpcClientNotInitializedError(
f'OPC client is not initialized for server {self.id}; '
'call _create_client() before opening a session'
)
if self._is_session_open():
raise OpcSessionAlreadyConnectedError(
f'OPC session already connected for server {self.id}; '
'call disconnect() before connecting again'
)
tags = {
'pod_id': self.pod_id,
'server_name': self.server_name,
}
await self.emit_metric(metrics.OPC_CONNECTIONS_TOTAL, tags)
try:
await self.client.connect()
session_id = _opc_authentication_token_str(self.client)
revised_session_timeout_ms = int(self.client.session_timeout)
revised_secure_channel_timeout_ms = int(self.client.secure_channel_timeout)
self.info(
f'OPC new session connected opc_server_id={self.id} session_id={session_id} '
f'revised_session_timeout_ms={revised_session_timeout_ms} '
f'revised_secure_channel_timeout_ms={revised_secure_channel_timeout_ms}',
self.metadata,
)
await self.emit_metric(
metrics.OPC_SESSION_CREATED_TOTAL, self._opc_debug_tags(session_id)
)
await self.emit_metric(
metric_object=metrics.OPC_SESSION_REVISED_TIMEOUT_MS,
method='set',
tags=self._opc_debug_tags(session_id),
value=revised_session_timeout_ms,
)
await self.emit_metric(
metric_object=metrics.OPC_CONNECTION_STATUS,
method='set',
tags={**tags, 'server_url': self.url},
value=1,
)
self._last_write_mono = None
self._session_ready.set()
return True, {}
except Exception as e:
await self._disconnect_locked()
trace = traceback.format_exc()
self.error(trace, self.metadata)
await self.emit_metric(metrics.OPC_CONNECTIONS_FAILED, tags)
return False, {
'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
'message': f'Failed to connect to OPC server: {e}',
'block': 'opc_repository',
'level': NotificationLevel.ERROR,
'attachment_content': trace,
}
async def _connect_locked(self) -> tuple[bool, dict[str, Any]]:
"""
Create the client when absent, then open a UA session.
Caller must hold _connection_lock.
Raises:
OpcSessionAlreadyConnectedError: If a session is already open.
Return:
tuple[bool, dict[str, Any]]: Result from _open_session on connect failure.
"""
if self._is_session_open():
raise OpcSessionAlreadyConnectedError(
f'OPC session already connected for server {self.id}; '
'call disconnect() before connecting again'
)
if self.client is None:
await self._create_client()
return await self._open_session()
async def _disconnection_fallback(self) -> list[dict[str, Any]]:
"""
Try up to five times to disconnect from the OPC UA server.
"""
assert self.client is not None
error_stack: list[dict[str, Any]] = []
for i in range(5):
try:
self.info(
f'Disconnecting from OPC UA server, attempt {i + 1} of 5',
self.metadata,
)
await self.client.disconnect()
return []
except Exception as e:
self.error(
f'Failed to disconnect from OPC UA server in attempt {i + 1} of 5: {e}',
self.metadata,
)
error_stack.append(
{
'attempt': i + 1,
'error': str(e),
'traceback': traceback.format_exc(),
}
)
await asyncio.sleep(self.disconnection_interval * i)
return error_stack
async def _disconnect_locked(self) -> None:
"""
Tear down the current session and client.
Caller must hold _connection_lock.
"""
self._last_write_mono = None
self._session_ready.clear()
if self.client is None:
return
session_id = _opc_authentication_token_str(self.client)
self.info(
f'OPC disconnecting opc_server_id={self.id} session_id={session_id}',
self.metadata,
)
await self.emit_metric(metrics.OPC_SESSION_CLOSED_TOTAL, self._opc_debug_tags(session_id))
errors = await self._disconnection_fallback()
if errors:
await self.send_notification_async(
metadata=self.metadata,
notification_id=f'OPC_DISCONNECTION_ERROR_{self.id}',
message='Failed to disconnect from OPC server in 5 attempts.',
block='opc_repository',
level=NotificationLevel.ERROR,
attachment_content=json.dumps(errors, indent=4),
)
else:
self.warning(f'Disconnected from OPC server {self.id} successfully', self.metadata)
await self.emit_metric(
metric_object=metrics.OPC_CONNECTION_STATUS,
method='set',
tags={
'pod_id': self.pod_id,
'server_name': self.server_name,
'server_url': self.url,
},
value=0,
)
self.client = None
async def _reconnect_locked(self) -> tuple[bool, dict[str, Any]]:
"""
Close the current session and open a new one.
Caller must hold _connection_lock. Records last_reconnection_time for interval gating.
Return:
tuple[bool, dict[str, Any]]: Result from _connect_locked after teardown.
"""
self.last_reconnection_time = datetime.now()
await self._disconnect_locked()
return await self._connect_locked()
async def connect(self) -> tuple[bool, dict[str, Any]]:
"""
Open an OPC UA session under the connection lock (worker initialization).
"""
async with self._connection_lock:
self.info(
f'Starting connection to OPC server {self.id}:{self.server_name}...',
self.metadata,
)
return await self._connect_locked()
async def disconnect(self) -> None:
"""
Gracefully disconnect from the OPC server under the connection lock.
Disables background reconnect so late writes during worker shutdown do not
respawn sessions.
"""
async with self._connection_lock:
self._allow_reconnect = False
await self._disconnect_locked()
async def validate_connection(self) -> tuple[bool, dict[str, Any]]:
"""
Read-only check that the asyncua protocol is open.
Caller must ensure _session_ready before writing. Does not connect or reconnect.
Return:
tuple[bool, dict[str, Any]]: (True, {}) when open, otherwise (False, error).
"""
if self._is_session_open():
return True, {}
self.error(f'OPC server {self.id} is not connected', self.metadata)
return False, self._not_connected_error()
def _reconnect_task_in_progress(self) -> bool:
"""
Return whether a background reconnect task is currently running.
Return:
bool: True when a reconnect task exists and has not finished.
"""
return self._reconnect_task is not None and not self._reconnect_task.done()
async def _start_reconnect(self, reason: str, session_id: str) -> None:
"""
Schedule a background reconnect when allowed by interval and task state.
Clears _session_ready before starting the task. No-op when _allow_reconnect is
False, the reconnection window has not elapsed, or a reconnect is already running.
Args:
reason (str): Trigger for reconnect (OPC status name or synthetic reason).
session_id (str): Session token before failure.
"""
if not self._allow_reconnect:
return
if not self._reconnection_window_elapsed():
self.warning(
f'OPC reconnect skipped reason=reconnection_window opc_server_id={self.id} '
f'reconnect_reason={reason}',
self.metadata,
)
return
if self._reconnect_task_in_progress():
self.warning(
f'OPC reconnect skipped reason=in_progress opc_server_id={self.id} '
f'reconnect_reason={reason}',
self.metadata,
)
return
self._session_ready.clear()
self.info(
f'OPC reconnect scheduled reconnect_reason={reason} opc_server_id={self.id} '
f'old_session_id={session_id}',
self.metadata,
)
self._reconnect_task = asyncio.create_task(self._run_reconnect(reason, session_id))
async def _run_reconnect(self, reason: str, session_id: str) -> None:
"""
Background task that tears down and re-establishes the OPC UA session.
Args:
reason (str): Trigger for reconnect (OPC status or ProtocolClosed).
session_id (str): Previous session token string for logging.
"""
try:
async with self._connection_lock:
self.info(
f'OPC reconnect started reconnect_reason={reason} opc_server_id={self.id} '
f'old_session_id={session_id}',
self.metadata,
)
success, error = await self._reconnect_locked()
if not success:
self.error(
f'OPC reconnect failed reconnect_reason={reason} opc_server_id={self.id}',
self.metadata,
)
if error:
self.error(error.get('message', ''), self.metadata)
except Exception:
self.error(
f'OPC reconnect task failed opc_server_id={self.id} reconnect_reason={reason}',
self.metadata,
)
self.error(traceback.format_exc(), self.metadata)
async def _log_write_inter_arrival(self, session_id: str, node: str) -> None:
"""
Log elapsed wall time since the previous successful OPC write on this repository.
Args:
session_id (str): Current OPC UA session token string.
node (str): Node id written in this operation.
"""
now = time.monotonic()
if self._last_write_mono is not None:
delta_s = now - self._last_write_mono
self.info(
f'OPC write inter-arrival_s={delta_s:.6f} opc_server_id={self.id} '
f'session_id={session_id} node={node}',
self.metadata,
)
if self.client is not None:
session_timeout_ms = float(self.client.session_timeout)
if session_timeout_ms > 0 and delta_s > (session_timeout_ms / 1000.0):
await self.emit_metric(
metrics.OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL,
self._opc_debug_tags(session_id),
)
self._last_write_mono = now
async def _emit_opc_write_metric(
self, session_id: str, result: str, metadata: dict[str, Any] | None
) -> None:
"""
Emit opc_write_attempts_total for a single write attempt outcome.
Args:
session_id (str): OPC UA session token string, or "unknown".
result (str): Outcome label (OK, OPC status name, ProtocolClosed, etc.).
metadata (dict[str, Any] | None): Write context for model_id/model_name labels.
"""
await self.emit_metric(
metrics.OPC_WRITE_ATTEMPTS_TOTAL,
{
**self._opc_debug_tags(session_id),
**_model_labels_from_write_metadata(metadata),
'result': result,
},
)
def _write_failure_payload(
self,
notification_id: str,
message: str,
level: NotificationLevel = NotificationLevel.ERROR,
attachment_content: str | None = None,
opc_error_kind: str | None = None,
opc_status: str | None = None,
) -> dict[str, Any]:
"""
Build a structured error dict returned from failed write_data paths.
Args:
notification_id (str): Stable notification identifier.
message (str): Human-readable failure message.
level (NotificationLevel): Severity for downstream notifications.
attachment_content (str | None): Optional traceback or diagnostic text.
opc_error_kind (str | None): Classifier (session_bad, connection_lost, etc.).
opc_status (str | None): OPC UA status name or synthetic reason.
Return:
dict[str, Any]: Error payload consumed by the OPC activity layer.
"""
payload: dict[str, Any] = {
'notification_id': notification_id,
'message': message,
'block': 'opc_repository',
'level': level,
}
if attachment_content is not None:
payload['attachment_content'] = attachment_content
if opc_error_kind is not None:
payload['opc_error_kind'] = opc_error_kind
if opc_status is not None:
payload['opc_status'] = opc_status
return payload
async def _handle_tier1_bad(
self,
exc: BaseException,
session_id: str,
node: str,
metadata: dict[str, Any],
phase: str,
) -> tuple[bool, dict[str, Any]]:
"""
Record metrics/logs and schedule reconnect after a Tier-1 Bad* error.
Args:
exc (BaseException): Tier-1 OPC UA error.
session_id (str): Session token at failure time.
node (str): Node id being written.
metadata (dict[str, Any]): Write context.
phase (str): get_node or write_value.
Return:
tuple[bool, dict[str, Any]]: Always (False, error payload).
"""
opc_status = _opc_status_from_exception(exc)
trace = traceback.format_exc()
self.error(trace, metadata)
await self._emit_opc_write_metric(session_id, opc_status, metadata)
self.error(
f'OPC write failed opc_status={opc_status} opc_server_id={self.id} '
f'session_id={session_id} model_id={metadata.get("model_id", "unknown")} '
f'model_name={metadata.get("model_name", "unknown")} node={node} phase={phase}',
metadata,
)
await self._start_reconnect(opc_status, session_id)
return False, self._write_failure_payload(
notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}',
message=f'Failed to {phase} on OPC server: {exc} | metadata: {metadata}',
attachment_content=trace,
opc_error_kind='session_bad',
opc_status=opc_status,
)
async def _write_reconnect_in_progress(
self, metadata: dict[str, Any]
) -> tuple[bool, dict[str, Any]]:
"""
Fail a write because a background reconnect task is already running.
Args:
metadata (dict[str, Any]): Write context passed through to the activity.
Return:
tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind reconnect_in_progress).
"""
await self._emit_opc_write_metric('unknown', 'ReconnectInProgress', metadata)
self.warning(
f'OPC write rejected reconnect_in_progress opc_server_id={self.id} '
f'model_id={metadata.get("model_id", "unknown")} '
f'model_name={metadata.get("model_name", "unknown")}',
metadata,
)
return False, {
'notification_id': f'OPC_WRITE_RECONNECT_IN_PROGRESS_{self.id}',
'message': f'OPC write skipped: reconnect in progress | metadata: {metadata}',
'block': 'opc_repository',
'level': NotificationLevel.WARNING,
'opc_error_kind': 'reconnect_in_progress',
}
async def _write_connection_lost(
self, metadata: dict[str, Any], opc_status: str
) -> tuple[bool, dict[str, Any]]:
"""
Fail a write after scheduling reconnect for a closed or stale session.
Args:
metadata (dict[str, Any]): Write context passed through to the activity.
opc_status (str): Synthetic reason (ProtocolClosed, SessionNotReady).
Return:
tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind connection_lost).
"""
await self._emit_opc_write_metric('unknown', opc_status, metadata)
return False, self._write_failure_payload(
notification_id=f'OPC_WRITE_CONNECTION_LOST_{self.id}',
message=f'OPC write skipped: connection lost ({opc_status}) | metadata: {metadata}',
level=NotificationLevel.WARNING,
opc_error_kind='connection_lost',
opc_status=opc_status,
)
async def write_data(
self, node: str, value: Any, data_type: str, metadata: dict[str, Any]
) -> tuple[bool, dict[str, Any]]:
"""
Write data to OPC server with a single attempt and background reconnect scheduling.
Reconnect is scheduled on Tier-1 Bad*, closed protocol, or stale session readiness.
There is no retry within the same call.
Args:
node (str): OPC UA node id to write.
value (Any): Value to convert and send.
data_type (str): Logical type key (float, int, bool, str, double).
metadata (dict[str, Any]): Activity context (model_id, model_name, etc.).
Return:
tuple[bool, dict[str, Any]]: (True, {response_time}) on success, or
(False, structured error info) on failure.
"""
if self._reconnect_task_in_progress():
return await self._write_reconnect_in_progress(metadata)
if not self._session_ready.is_set():
session_id = _opc_authentication_token_str(self.client)
await self._start_reconnect('SessionNotReady', session_id)
if self._reconnect_task_in_progress():
return await self._write_reconnect_in_progress(metadata)
return await self._write_connection_lost(metadata, 'SessionNotReady')
is_connected, _error = await self.validate_connection()
if not is_connected:
session_id = _opc_authentication_token_str(self.client)
await self._start_reconnect('ProtocolClosed', session_id)
return await self._write_connection_lost(metadata, 'ProtocolClosed')
start_time = time.time()
session_id = _opc_authentication_token_str(self.client)
try:
node_obj = self.client.get_node(node) # type: ignore[union-attr]
except Exception as e:
if is_reconnectable_opcua_bad(e):
return await self._handle_tier1_bad(e, session_id, node, metadata, 'get_node')
trace = traceback.format_exc()
self.error(trace, metadata)
await self._emit_opc_write_metric(
session_id, f'GetNodeError:{type(e).__name__}', metadata
)
return False, self._write_failure_payload(
notification_id=f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
message=f'Failed to get node from OPC server: {e} | metadata: {metadata}',
attachment_content=trace,
)
if data_type not in data_type_map:
await self._emit_opc_write_metric(session_id, 'UnsupportedDataType', metadata)
return False, self._write_failure_payload(
notification_id=f'OPC_WRITE_DATA_TYPE_ERROR_{self.id}',
message=f'Unsupported data type: {data_type} | metadata: {metadata}',
)
data = data_type_map[data_type]['converter'](value)
self.info(f'Writing {data} - {type(data)} to {node}', metadata)
ua_data = DataValue(
Variant(data, data_type_map[data_type]['opc_type']),
)
try:
await node_obj.write_value(ua_data)
end_time = time.time()
response_time = end_time - start_time
except Exception as e:
if is_reconnectable_opcua_bad(e):
return await self._handle_tier1_bad(e, session_id, node, metadata, 'write_value')
trace = traceback.format_exc()
self.error(trace, metadata)
await self._emit_opc_write_metric(session_id, type(e).__name__, metadata)
return False, self._write_failure_payload(
notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}',
message=f'Failed to write data to OPC server: {e} | metadata: {metadata}',
attachment_content=trace,
)
await self._emit_opc_write_metric(session_id, 'OK', metadata)
await self._log_write_inter_arrival(session_id, node)
return True, {
'response_time': response_time,
}

View File

349
laborious/worker/worker.py Normal file
View File

@@ -0,0 +1,349 @@
"""
Laborious Worker Module
This module provides the main worker implementation for the Sientia DataOps Laborious system.
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
prediction and retraining workflows.
The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``):
- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload)
- minimal_retrain-{runtime}-queue: Model retraining workflows
- drift-{runtime}-queue: Drift detection workflows
- simple_metrics-{runtime}-queue: Simple metrics workflows
- import_model-{runtime}-queue: `.sientia` model import — registered **only** when
``IMPORT_MODEL_ENABLED`` is truthy, so the containment boundary stays a deployment decision
``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the
same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``).
Key Features:
- Resource-based scaling with WorkerTuner (CPU and memory aware)
- Automatic polling scaling with PollerBehaviorAutoscaling
- Prometheus metrics integration
- Comprehensive error handling and logging
- Graceful shutdown with cleanup
- Multiple worker instances for different workflow types
Environment Variables:
- RUNTIME: Required non-empty string; suffix for all task queue names
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious)
- POD_ID: Kubernetes pod identifier for metrics
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
- PROJECT_NAME: Project name for notifications (default: laborious)
- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default
- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_NAME: import configuration, see
``laborious.utils.connectors_config`` (the import log's server is POSTGRES_*)
"""
from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
with workflow.unsafe.imports_passed_through():
import asyncio
import os
import sys
from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import get_logger
from sientia_do.temporal.worker.prepare_worker import prepare_worker
from sientia_do.utils.connectors_config import (
build_api_config,
build_mongodb_config,
build_postgres_config,
)
from laborious import metrics
from laborious.activities.activities import Activities
from laborious.utils.connectors_config import (
build_import_config,
build_import_status_config,
build_minio_config,
build_mlflow_config,
build_opc_config,
)
from laborious.workflows.drift import Drift
from laborious.workflows.import_model import ImportModel
from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.simple_metrics import SimpleMetrics
from laborious.workflows.sub_workflows.format_and_export_prediction import (
FormatAndExportPrediction,
)
from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
POD_ID = os.getenv('HOSTNAME')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
# The values that turn the import worker on. Anything else — unset, empty, `0`, `false`, a typo —
# leaves the four default workers alone and polls no import queue at all.
IMPORT_ENABLED_VALUES = ('1', 'true', 'yes', 'on')
def import_worker_enabled() -> bool:
"""Whether this runtime registers the fifth (import) worker.
Read at worker start, not per workflow: containment (QTZPOC-19) wants the importer in its own
runtime with its own credentials, and this flag is what makes that a deployment decision.
"""
return os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in IMPORT_ENABLED_VALUES
async def main():
"""
Main entry point for the Laborious worker application.
This function initializes and starts all components of the worker:
1. Sets up logging and metadata
2. Starts Prometheus metrics server
3. Initializes notification handler
4. Creates and configures activities
5. Initializes OPC connections
6. Starts Temporal client and workers
7. Manages worker lifecycle and graceful shutdown
The function runs indefinitely until interrupted or an error occurs.
On error, it performs cleanup and exits with a non-zero status code.
Raises:
Exception: Any unhandled exception during worker execution
SystemExit: On graceful shutdown or error conditions
"""
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
logger = get_logger(__name__)
metadata = {
'pod_id': POD_ID,
'model_name': '-',
'model_id': '-',
'workflow_name': '-',
'schedule_name': '-',
}
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
runtime = os.getenv('RUNTIME', '').strip()
if not runtime:
logger.custom_critical(
'RUNTIME environment variable is required and must be non-empty',
metadata,
)
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
sys.exit(1)
metadata_runtime = {**metadata, 'runtime': runtime}
logger.custom_info('Starting prometheus client...', metadata_runtime)
start_prometheus_server()
logger.custom_info('Starting Notification Handler...', metadata_runtime)
mongo_config = build_mongodb_config()
notification_handler = NotificationHandler(
connection_string=mongo_config['connection_string'],
database=mongo_config['database_name'],
logger=logger,
project_name=os.getenv('PROJECT_NAME', 'laborious'),
)
logger.custom_info('Starting Activities...', metadata_runtime)
activities = Activities(
postgres_config=build_postgres_config(),
mlflow_config=build_mlflow_config(),
minio_config=build_minio_config(),
opc_config=build_opc_config(),
pi_web_api_config=build_api_config(),
logger=logger,
notification_handler=notification_handler,
mongo_config=mongo_config,
import_config=build_import_config(),
import_status_config=build_import_status_config(),
)
logger.custom_info('Initializing OPC...', metadata_runtime)
await activities.init_opc()
logger.custom_info(
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...',
metadata_runtime,
)
new_runtime = Runtime(
telemetry=TelemetryConfig(
metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
)
)
logger.custom_info(f'Starting Temporal Client at {host}...', metadata_runtime)
temporal_client = await client.Client.connect(
target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
runtime=new_runtime,
)
logger.custom_info(f'Starting Workers (runtime={runtime})...', metadata_runtime)
workers = [
prepare_worker(
temporal_client=temporal_client,
main_workflow=MinimalRetrain,
other_workflows=[],
activities=[
activities.load_query_with_minio_offload,
activities.retrain_model,
activities.update_production_model,
activities.format_retrain_report,
activities.export_data_to_postgres,
],
logger=logger,
runtime=runtime,
),
prepare_worker(
temporal_client=temporal_client,
main_workflow=SimpleMetrics,
other_workflows=[],
activities=[
activities.load_custom_query,
activities.calculate_simple_metrics,
activities.export_data_to_postgres,
],
logger=logger,
runtime=runtime,
),
prepare_worker(
temporal_client=temporal_client,
main_workflow=Drift,
other_workflows=[],
activities=[
activities.load_custom_query,
activities.get_reference_data,
activities.calculate_drift,
activities.export_data_to_postgres,
],
logger=logger,
runtime=runtime,
),
prepare_worker(
temporal_client=temporal_client,
main_workflow=PredictionsBatch,
other_workflows=[PredictionProcess, FormatAndExportPrediction],
activities=[
# MLFlow
activities.request_predict,
activities.request_transform,
# Gates
activities.input_gate,
activities.mlflow_response_gate,
activities.mlflow_content_gate,
activities.format_transformed_data,
activities.format_prediction,
activities.format_default_prediction,
# OPC
activities.write_opc_data,
# Postgres / MinIO offload
activities.load_query_with_minio_offload,
activities.cleanup_minio_objects_expired,
activities.repeat_last_prediction,
activities.export_data_to_postgres,
activities.export_payload_to_postgres,
activities.write_metrics,
# Pi Web API
activities.write_pi_web_api_data,
],
logger=logger,
runtime=runtime,
),
]
# The fifth worker is opt-in: containment (QTZPOC-19) wants the importer in its own runtime with
# its own credentials, and a flag lets that be a deployment decision instead of forcing every
# Laborious runtime to poll the import queue today. Unset, the four workers above are unaffected
# and no import queue is polled at all.
if import_worker_enabled():
logger.custom_info('Import model worker enabled', metadata_runtime)
workers.append(
prepare_worker(
temporal_client=temporal_client,
main_workflow=ImportModel,
other_workflows=[],
activities=[
# Import log row (the BFF database, update-only)
activities.claim_import_status,
activities.record_import_status,
activities.record_import_names,
activities.record_import_terminal_status,
activities.report_import_status_write_failure,
# Bundle
activities.download_import_bundle,
activities.open_import_bundle,
# MLflow provisioning
activities.create_import_experiment,
activities.upload_import_artifacts,
activities.register_import_model_version,
# The model listing
activities.write_import_model_document,
# Cleanup and retention
activities.cleanup_import_files,
activities.ensure_import_bundle_retention,
],
logger=logger,
runtime=runtime,
)
)
handlers = []
for w in workers:
handlers.append(w.run())
logger.custom_info('Workers started successfully', metadata_runtime)
exit_code = 0
try:
# This will run the workers and wait for them to complete.
# If an exception occurs in any of the worker handlers, it will be propagated here.
await asyncio.gather(*handlers)
except BaseException as e: # NOSONAR
logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
exit_code = 1
finally:
if notification_handler:
notification_handler.shutdown()
if activities:
await activities.shutdown()
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
sys.exit(exit_code)
def start_prometheus_server():
"""
Starts the Prometheus metrics server for monitoring and observability.
This function initializes the Prometheus HTTP server on the configured port
and sets the application health metric to indicate the service is running.
The server exposes metrics that can be scraped by Prometheus for monitoring
the health and performance of the Laborious worker.
Environment Variables:
HTTP_METRICS_PORT: Port for the metrics server (default: 9090)
POD_ID: Pod identifier for metrics labeling
Raises:
SystemExit: If the metrics server fails to start
"""
try:
port = int(os.getenv('HTTP_METRICS_PORT', 9090))
start_http_server(port)
print(f'Prometheus server started on port {port}.')
metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP
except Exception as e:
print(f'Failed to start Prometheus server: {e}')
os._exit(1)
if __name__ == '__main__':
asyncio.run(main())

View File

View File

@@ -0,0 +1,107 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name='drift')
class Drift:
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the drift workflow.
This method orchestrates the complete drift process by:
1. Loading data using the provided custom SQL query
2. Preparing prediction configuration and filters
3. Delegating to the PredictionProcess workflow for ML operations
"""
metadata = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'drift',
}
}
print(f'Input data: {input_data}', metadata)
model_config = input_data['model_config']
target_name = model_config['target']
gathering_query = f"""
SELECT *
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
WHERE
model_id = '{input_data['model_id']}' AND
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
ORDER BY timestamp ASC
""" # nosec B608 - values come from internal Temporal workflow config, not user input
target_data_handler = workflow.start_activity_method(
Activities.load_custom_query,
{
**metadata,
'query': gathering_query,
'datetime_columns': ['timestamp', 'created_at'],
'orient': 'records',
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
reference_data_handler = workflow.start_activity_method(
Activities.get_reference_data,
{**metadata, 'model_name': input_data['model_name']},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
target_data = await target_data_handler
reference_data = await reference_data_handler
if not target_data:
return
drift_data = await workflow.execute_local_activity_method(
Activities.calculate_drift,
{
**metadata,
'target_data': target_data,
'reference_data': reference_data,
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'target_name': target_name,
'drift_metrics': input_data.get(
'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
),
'chunk_period': input_data.get('chunk_period', 'min'),
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
if drift_data:
await workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
**metadata,
'data': drift_data,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)

View File

@@ -0,0 +1,424 @@
"""The `import_model` workflow: a `.sientia` object becomes a registered model, or a durable ERROR.
Shape of this workflow, and the reasons it has that shape:
- **The claim comes first.** The record already exists — the frontend inserted it through the BFF —
so the workflow's first act is to take ownership of it. Everything after that failure is recorded
*on* the record; a failed claim is recorded nowhere, because the workflow owns no row.
- **Fail-fast, and nothing else.** No rollback, no compensation, no resume, no reconciliation, no
`continue_as_new`. A failed import stays failed; trying again is a new import with a new upload.
- **Only the status writes retry**, and only because `UPDATE ... SET <fixed values> WHERE id` is
idempotent. Every provisioning call keeps `maximum_attempts=1`, because a repeated registry call or
Mongo insert can write twice.
- **One `finally`, two actions**: cleanup, then the single terminal write that classifies the
outcome once. No `except` block writes a verdict of its own.
- **The workflow holds ids, never contents.** The decrypted bundle never crosses an activity
boundary, and the password reaches this module only as an envelope it cannot open.
"""
from datetime import timedelta
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
from laborious.activities.model_import import (
IMPORT_STATUS_HINT_RETRY_POLICY,
IMPORT_STATUS_RETRY_POLICY,
derive_run_name,
)
from laborious.activities.model_import_errors import ImportInputError, failure_location
from laborious.utils.bundle.steps import ImportStep
# `start_to_close_timeout` is per attempt. The terminal write additionally gets a
# `schedule_to_close_timeout` larger than its retry envelope (10 attempts, ≈151 s of backoff, ≈251 s
# if every attempt also burns its timeout) — set it smaller and Temporal cuts the retries short,
# which is the one configuration mistake that would silently undo design D20.
STATUS_WRITE_TIMEOUT = timedelta(seconds=10)
STATUS_WRITE_ENVELOPE = timedelta(seconds=300)
DOWNLOAD_TIMEOUT = timedelta(seconds=600)
OPEN_TIMEOUT = timedelta(seconds=900)
EXPERIMENT_TIMEOUT = timedelta(seconds=60)
ARTIFACT_UPLOAD_TIMEOUT = timedelta(seconds=900)
REGISTRATION_TIMEOUT = timedelta(seconds=120)
DOCUMENT_TIMEOUT = timedelta(seconds=60)
CLEANUP_TIMEOUT = timedelta(seconds=120)
RETENTION_TIMEOUT = timedelta(seconds=60)
# Fields no import may carry. The project is a frontend concern: the model listing is the MongoDB
# `models` document, no migrated database holds a project or model table, and nothing in this
# repository reads a project-to-model link (design D22). An input naming one is rejected rather than
# accepted and ignored, so no caller believes a project link was recorded. The same goes for a
# relational target: there is none, and for a plaintext password, which does not belong in an event
# history whatever else is true.
REJECTED_INPUT_FIELDS = (
'project_id',
'project_name',
'provisioning_target',
'schema',
'table_name',
'status_table',
'password',
)
DIGEST_LENGTH = 64
_HEX_DIGITS = frozenset('0123456789abcdef')
@workflow.defn(name='import_model')
class ImportModel:
"""Turn an encrypted `.sientia` object in MinIO into a registered model, or into an ERROR."""
@workflow.run
async def run(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""Execute the import.
Args:
input_data: the import request.
Required keys:
- import_run_id (int): primary key of the import log row the frontend inserted
- object_key (str): the uploaded object's key
- expected_digest (str): the SHA-256 the uploader computed for those bytes
- password_envelope (str): the encrypted bundle password, never plaintext
Optional keys:
- bucket (str): defaults to the configured import bucket
Returns:
dict: `import_run_id`, `run_id`, `model_name`, `version` and the model document's id.
Raises:
ImportLogRowNotClaimableError: the record is not this workflow's — the run ends FAILED
with that error type and nothing is written to any row.
Exception: whatever failed. The record carries the failing step, its code and its
sentence; this exception carries the technical truth.
"""
workflow_id = workflow.info().workflow_id
import_run_id = self._require_import_run_id(input_data.get('import_run_id'))
metadata: dict[str, Any] = {
'model_name': '-',
'model_id': str(import_run_id),
'workflow_name': 'import_model',
'schedule_name': '-',
}
base = {
'metadata': metadata,
'import_run_id': import_run_id,
'workflow_id': workflow_id,
}
# Class A. Outside the try/finally on purpose: a claim that matches nothing leaves this
# workflow owning no row, so there is nothing to clean up and nothing to record.
claim = await workflow.execute_activity_method(
Activities.claim_import_status,
base,
retry_policy=IMPORT_STATUS_RETRY_POLICY,
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
)
current_step = ImportStep.RECEIVED
failed_step: str | None = None
gate: int | None = None
succeeded = False
source: dict[str, Any] = {}
result: dict[str, Any] = {'import_run_id': import_run_id}
try:
self._validate_input(input_data, claim)
current_step = ImportStep.DOWNLOAD
await self._hint(base, current_step)
await self._ensure_retention(base, input_data)
download = await workflow.execute_activity_method(
Activities.download_import_bundle,
{
**base,
'bucket': input_data.get('bucket'),
'object_key': input_data['object_key'],
'expected_digest': input_data['expected_digest'],
},
retry_policy=retry_policy,
start_to_close_timeout=DOWNLOAD_TIMEOUT,
)
current_step = ImportStep.DECRYPTION
await self._hint(base, current_step)
opened = await workflow.execute_activity_method(
Activities.open_import_bundle,
{
**base,
'path': download['path'],
'password_envelope': input_data['password_envelope'],
},
retry_policy=retry_policy,
start_to_close_timeout=OPEN_TIMEOUT,
)
source = opened['metadata']
parameters = opened['parameters']
model_name = source['model_name']
metadata['model_name'] = model_name
run_name = derive_run_name(model_name, source['model_version'])
# The names go on the record before anything exists in MLflow, so a person watching the
# import list sees *which* model is being imported while it is still running.
await self._record_names(base, source, run_name)
current_step = ImportStep.EXPERIMENT_CREATION
await self._hint(base, current_step, source)
experiment = await workflow.execute_activity_method(
Activities.create_import_experiment,
{**base, 'experiment_name': source['experiment_name']},
retry_policy=retry_policy,
start_to_close_timeout=EXPERIMENT_TIMEOUT,
)
current_step = ImportStep.ARTIFACT_UPLOAD
await self._hint(base, current_step, source)
run = await workflow.execute_activity_method(
Activities.upload_import_artifacts,
{
**base,
'experiment_id': experiment['experiment_id'],
'run_name': run_name,
'extracted_dir': opened['extracted_dir'],
'parameters': parameters,
'source': source,
},
retry_policy=retry_policy,
start_to_close_timeout=ARTIFACT_UPLOAD_TIMEOUT,
)
current_step = ImportStep.REGISTRATION
await self._hint(base, current_step, source)
version = await workflow.execute_activity_method(
Activities.register_import_model_version,
{**base, 'model_name': model_name, 'run_id': run['run_id']},
retry_policy=retry_policy,
start_to_close_timeout=REGISTRATION_TIMEOUT,
)
current_step = ImportStep.MODEL_DOCUMENT
await self._hint(base, current_step, source)
document = await workflow.execute_activity_method(
Activities.write_import_model_document,
{
**base,
'model_name': model_name,
'target': parameters['target_variable'],
},
retry_policy=retry_policy,
start_to_close_timeout=DOCUMENT_TIMEOUT,
)
succeeded = True
result = {
'import_run_id': import_run_id,
'run_id': run['run_id'],
'model_name': model_name,
'version': version['version'],
'document_id': document.get('id'),
}
except BaseException as error: # NOSONAR - classified once below, then re-raised untouched
located_step, located_gate = failure_location(error)
failed_step = located_step or current_step.value
gate = located_gate
raise
finally:
cleanup = await self._cleanup(base)
await self._record_terminal(
base,
succeeded=succeeded,
step=failed_step or current_step.value,
gate=gate,
source=source,
cleanup_failed=not cleanup.get('cleaned', True),
)
return result
# ------------------------------------------------------------------ validation
@staticmethod
def _require_import_run_id(value: Any) -> int:
"""The one check that runs before the claim, because the claim needs its result.
A missing or unusable id means there is no record to claim and none to write: the run ends
FAILED with this error type, and the failure is visible in the workflow's terminal state, its
error type, the error metric and a notification — everywhere except a row, which is the
point.
"""
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
raise ImportInputError('import_run_id must be a positive integer')
return value
@staticmethod
def _validate_input(input_data: dict[str, Any], claim: dict[str, Any]) -> None:
"""Validate the rest of the request, and check it against the record it just claimed.
This runs **after** the claim so its failure lands on the record rather than only in the
logs: by now the row is this workflow's, so a bad request is reported to the person who made
it. Everything here fails at step `received`.
"""
unknown = [name for name in REJECTED_INPUT_FIELDS if name in input_data]
if unknown:
raise ImportInputError(
f'the import request carries fields this workflow does not accept: '
f'{", ".join(sorted(unknown))}'
)
for required in ('object_key', 'expected_digest', 'password_envelope'):
if not input_data.get(required):
raise ImportInputError(f'the import request is missing {required}')
digest = str(input_data['expected_digest']).lower()
if len(digest) != DIGEST_LENGTH or not set(digest) <= _HEX_DIGITS:
raise ImportInputError('expected_digest must be a 64-character hexadecimal SHA-256')
object_key = input_data['object_key']
file_name = claim.get('file_name')
if file_name and file_name != object_key:
raise ImportInputError('the record names a different object than the request')
request_digest = claim.get('request_digest')
if request_digest and request_digest != digest:
raise ImportInputError('the record names a different digest than the request')
# --------------------------------------------------------------- status writes
async def _hint(
self,
base: dict[str, Any],
step: ImportStep,
source: dict[str, Any] | None = None,
) -> None:
"""Record the step about to be attempted. A hint: its failure never fails the import."""
try:
await workflow.execute_activity_method(
Activities.record_import_status,
{**base, 'step': step.value, 'source': source or None},
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
)
except BaseException as error: # NOSONAR - a lost hint must not fail a working import
await self._report_status_write_failure(base, error)
async def _record_names(
self, base: dict[str, Any], source: dict[str, Any], run_name: str
) -> None:
"""Write the experiment and run names. Also a hint: the run is created either way."""
try:
await workflow.execute_activity_method(
Activities.record_import_names,
{
**base,
'experiment_name': source['experiment_name'],
'run_name': run_name,
'source': source,
'step': ImportStep.EXPERIMENT_CREATION.value,
},
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
)
except BaseException as error: # NOSONAR - the names are for a human reading the list
await self._report_status_write_failure(base, error)
async def _record_terminal(
self,
base: dict[str, Any],
*,
succeeded: bool,
step: str | None,
gate: int | None,
source: dict[str, Any],
cleanup_failed: bool,
) -> None:
"""The single authoritative write, and the last thing this workflow does.
Its own failure is reported through the channel with the verdict it could not write, and it
never turns a completed import into a failed one — nor masks the cause of a real failure.
"""
try:
await workflow.execute_activity_method(
Activities.record_import_terminal_status,
{
**base,
'succeeded': succeeded,
'step': step,
'gate': gate,
'source': source or None,
'cleanup_failed': cleanup_failed,
},
retry_policy=IMPORT_STATUS_RETRY_POLICY,
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
)
except (
BaseException
) as error: # NOSONAR - the outcome stands even when it cannot be written
await self._report_status_write_failure(
base,
error,
verdict=f'succeeded={succeeded} step={step} gate={gate}',
)
async def _report_status_write_failure(
self,
base: dict[str, Any],
error: BaseException,
verdict: str | None = None,
) -> None:
"""Route a status write that could not be made through the one failure channel."""
try:
await workflow.execute_activity_method(
Activities.report_import_status_write_failure,
{**base, 'detail': str(error), 'verdict': verdict},
retry_policy=retry_policy,
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
)
except BaseException: # noqa: S110 # NOSONAR - the report of a report is where this stops
# Deliberately silent: the channel this was trying to reach is the thing that logs, so
# there is nowhere left to say it, and an import that worked must not fail over it.
pass
# ----------------------------------------------------------- side-effect steps
@staticmethod
async def _ensure_retention(base: dict[str, Any], input_data: dict[str, Any]) -> None:
"""Ensure the bundle bucket's expiry rule, without making it a reason to fail an import.
The rule is applied by the importer rather than assumed from a Helm value, so the acceptance
criterion is verifiable from the code that depends on it. But a bucket policy the importer
cannot set is an operational problem, not a bad import: failing a valid model import because
a lifecycle API refused would be the wrong trade, so the failure is dropped here and left
where Temporal already records it — the activity's own failure in the event history.
"""
try:
await workflow.execute_activity_method(
Activities.ensure_import_bundle_retention,
{**base, 'bucket': input_data.get('bucket')},
retry_policy=retry_policy,
start_to_close_timeout=RETENTION_TIMEOUT,
)
except BaseException: # noqa: S110 # NOSONAR - retention is not a precondition of an import
# Dropped, not reported: the activity's own failure is already in the event history,
# and a bucket policy the importer cannot set is an operational problem rather than a
# bad import. Failing a valid model import over a lifecycle API would be the wrong
# trade.
pass
@staticmethod
async def _cleanup(base: dict[str, Any]) -> dict[str, Any]:
"""Remove this import's temporary files. Runs on every path, before the terminal write."""
try:
return await workflow.execute_activity_method(
Activities.cleanup_import_files,
base,
retry_policy=retry_policy,
start_to_close_timeout=CLEANUP_TIMEOUT,
)
except BaseException: # NOSONAR - cleanup never replaces the failure that brought us here
return {'cleaned': False}

View File

@@ -0,0 +1,137 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
@workflow.defn(name='minimal_retrain')
class MinimalRetrain:
"""
Automated model retraining workflow for the Laborious system.
This workflow implements a complete model retraining pipeline that loads
training data, executes model retraining, updates production models,
and maintains comprehensive audit trails. It's designed for automated
model lifecycle management with minimal manual intervention.
The workflow provides a robust retraining process with:
- Automated data loading from configured data sources
- MLFlow model retraining with quality validation
- Production model updates with version control
- Comprehensive reporting and audit trail maintenance
- Error handling and notification integration
"""
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the automated model retraining workflow.
This method orchestrates the complete model retraining process by:
1. Loading training data using the provided custom SQL query
2. Executing MLFlow model retraining with the loaded data
3. Updating production models with newly trained versions
4. Persisting comprehensive retraining reports to database
The method implements comprehensive error handling and ensures all
required parameters are properly configured before proceeding.
Args:
input_data: Complete configuration for the retraining workflow
Required keys:
- schedule_name (str): Schedule identifier for the retraining
- model_name (str): Name of the ML model to retrain
- model_id (int): Unique identifier for the model version
- query (str): SQL query for training data loading
- schema (str, optional): Database schema for report storage
- table_name (str, optional): Target table for retraining reports
- datetime_columns (list[str], optional): Columns to treat as datetime
Returns:
None: The workflow completes successfully when all steps finish
Raises:
Exception: If any required parameters are missing or if the workflow fails
during data loading, retraining, or model update operations
"""
metadata = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'minimal_retrain',
}
}
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
storage_result = await workflow.execute_activity_method(
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': model_name,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=600),
)
storage_payload = MinioDataFramePayload.from_dict(storage_result)
if not storage_payload.has_data():
raise ValueError('No data returned from query')
experiment_response = await workflow.execute_activity_method(
Activities.retrain_model,
{
**metadata,
'data': storage_result,
'model_name': model_name,
'model_config': model_config,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(hours=1),
)
if experiment_response['success']:
update_report = await workflow.execute_activity_method(
Activities.update_production_model,
{**metadata, 'model_name': model_name, **experiment_response},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
else:
update_report = {}
report = await workflow.execute_local_activity_method(
Activities.format_retrain_report,
{
**metadata,
'experiment_response': experiment_response,
'model_name': model_name,
'model_id': input_data['model_id'],
'update_report': update_report,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
await workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
**metadata,
'data': report,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=600),
)

View File

@@ -0,0 +1,127 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name='predictions_batch')
class PredictionsBatch:
"""
Main batch prediction workflow for the Laborious system.
This workflow orchestrates the complete batch prediction process, handling
data loading, configuration management, and workflow delegation. It serves
as the primary entry point for batch prediction operations and ensures
proper data preparation before ML model inference.
The workflow implements a robust data processing pipeline with:
- Custom SQL query execution for data loading
- Comprehensive configuration management
- Data quality filter application
- MLFlow model integration
- Workflow delegation to specialized sub-workflows
Workflow Execution:
1. Data Loading: Executes custom SQL query to load prediction data
2. Configuration Preparation: Sets up prediction parameters and filters
3. Workflow Delegation: Spawns PredictionProcess child workflow
4. Error Handling: Implements comprehensive error handling and retry policies
"""
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the batch prediction workflow.
This method orchestrates the complete batch prediction process by:
1. Loading data using the provided custom SQL query
2. Preparing prediction configuration and filters
3. Delegating to the PredictionProcess workflow for ML operations
The method implements comprehensive error handling and ensures all
required parameters are properly configured before proceeding.
Args:
input_data: Complete configuration for the batch prediction
Required keys:
- schedule_name (str): Schedule identifier for the prediction
- model_name (str): Name of the ML model to use
- model_id (int): Unique identifier for the model
- query (str): SQL query for data loading
- schema (dict, optional): Data schema definition
- table_name (str, optional): Target table for predictions
- input_filters (dict, optional): Data quality filters
- mlflow_transform_filters (dict, optional): MLFlow transform filters
- mlflow_predict_filters (dict, optional): MLFlow prediction filters
- model_retention (int, optional): Model retention period in minutes
- path_priority (list[str]): Decision path priority configuration
- opc_output_config (dict, optional): OPC server export configuration
- pi_web_api_output_config (dict, optional): PI Web API export configuration
- datetime_columns (list[str], optional): Columns to treat as datetime
- save_transform (bool, optional): Whether to save transformed data (default: True)
- prediction_store_policy (str, optional): Data retention policy (default: 'lts:1')
Returns:
None: The workflow completes successfully when the child workflow finishes
Raises:
Exception: If any required parameters are missing or if the workflow fails
during data loading or workflow delegation
"""
metadata = {
'metadata': {
'schedule_name': input_data['schedule_name'],
'model_name': input_data['model_name'],
'model_id': input_data['model_id'],
'workflow_name': 'predictions_batch',
}
}
# Load data using custom query with optional MinIO offload for large frames
data = await workflow.execute_activity_method(
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
# Prepare input for prediction_process workflow
prediction_input = {
'metadata': metadata,
'data': data,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'input_filters': input_data.get(
'input_filters', {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
),
'mlflow_transform_filters': input_data.get(
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
),
'mlflow_predict_filters': input_data.get(
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
),
'model_config': input_data.get('model_config', {}),
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
'opc_output_config': input_data.get('opc_output_config', {}),
'on_conflict': input_data.get('on_conflict', 'error'),
'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
'save_transform': input_data.get('save_transform', True),
}
# Execute prediction process workflow
await workflow.execute_child_workflow('subworkflow.prediction_process', prediction_input)

View File

@@ -0,0 +1,95 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name='simple_metrics')
class SimpleMetrics:
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the simple metrics workflow.
"""
metadata = {
'metadata': {
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
'workflow_name': 'simple_metrics',
'schedule_name': input_data['schedule_name'],
}
}
model_id = input_data['model_id']
interval_minutes = input_data['interval_minutes']
model_config = input_data['model_config']
target_name = model_config['target']
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 = '{model_id}' and
p.prediction is not null and
ld.variable = '{target_name}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
order by
p."timestamp" desc;
""" # nosec B608 - values come from internal Temporal workflow config, not user input
target_data = await workflow.execute_activity_method(
Activities.load_custom_query,
{
**metadata,
'query': query,
'datetime_columns': ['timestamp'],
'orient': 'records',
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
if not target_data:
return
simple_metrics = await workflow.execute_local_activity_method(
Activities.calculate_simple_metrics,
{
**metadata,
'model_id': model_id,
'target_data': target_data,
'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']),
'interval_minutes': interval_minutes,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)
if not simple_metrics:
return
await workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
**metadata,
'data': simple_metrics,
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=300),
)

View File

@@ -0,0 +1,220 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name='subworkflow.format_and_export_prediction')
class FormatAndExportPrediction:
"""
Data formatting and export workflow for prediction results.
This workflow handles the final stages of the prediction pipeline, including
data formatting, database persistence, OPC server export, and metrics recording.
It implements flexible formatting based on prediction quality and provides
comprehensive export capabilities to multiple destinations.
The workflow supports two main prediction paths:
1. Normal Prediction: Formats and exports successful prediction results
2. Default Prediction: Creates fallback predictions for error conditions
Export Destinations:
- PostgreSQL Database: Persistent storage with timestamp conversion
- PI Web API: Real-time industrial system integration for prediction and confidence values
- OPC Servers: Real-time industrial system integration
- Prometheus Metrics: Performance monitoring and operational visibility
"""
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the prediction formatting and export workflow.
This method orchestrates the complete data export process by:
1. Determining the appropriate formatting strategy based on path_flag
2. Formatting prediction data according to quality and requirements
3. Exporting data to PI Web API for real-time industrial access (if configured)
4. Exporting data to OPC servers for real-time industrial access (if configured)
5. Persisting data to PostgreSQL database with comprehensive metadata
6. Recording performance metrics for operational monitoring
The method implements flexible formatting strategies:
- Normal predictions: Full data formatting with confidence scores
- Error predictions: Default formatting with error indicators
- Comprehensive export: Multi-destination data distribution
Args:
input_data: Complete configuration for the export workflow
Required keys:
- metadata (dict): Workflow execution metadata
- path_flag (str | None): Decision path flag for formatting strategy
- None: Normal prediction path with full formatting
- Any other value: Default prediction path for error conditions
- data (dict[str, Any]): Prediction data to format and export
- prediction_confidence (float): Confidence score for the prediction
- timestamp (str): ISO-formatted timestamp for the prediction
- model_id (int): Unique identifier for the ML model
- model_name (str): Name of the ML model
- schema (str): Database schema for data storage
- table_name (str): Target table for data persistence
Optional keys:
- opc_output_config (dict[str, Any]): OPC server export configuration
- pi_web_api_output_config (dict[str, Any]): PI Web API export configuration
Contains endpoint, prediction_tags, and confidence_tags mappings
- transformed_data (dict[str, Any]): Transformed data to export separately
Only processed when path_flag is None
- transform_table_name (str): Target table for transformed data export
Required if transformed_data is provided
- prediction_store_policy (str): Data retention policy (e.g., 'lts:1', 'erl:2')
Required when path_flag is None
- comment (str): Operational comment or error description
Required when path_flag is not None
Returns:
None: The workflow completes successfully when all export operations finish
Note:
When transformed_data is provided and path_flag is None, the workflow will:
1. Format the transformed data using format_transformed_data
2. Export it to a separate table (transform_table_name) asynchronously
3. Wait for both prediction and transformed data exports to complete
"""
metadata = input_data['metadata']
path_flag = input_data['path_flag']
data = input_data['data']
transformed_data = input_data.get('transformed_data', None)
prediction_confidence = input_data['prediction_confidence']
opc_output_config = input_data.get('opc_output_config', None)
pi_web_api_output_config = input_data.get('pi_web_api_output_config', None)
if path_flag is None:
# Normal prediction path: format prediction data with full metadata
prediction = await workflow.execute_local_activity_method(
Activities.format_prediction,
{
**metadata,
'data': data,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'prediction_store_policy': input_data['prediction_store_policy'],
'model_name': input_data['model_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
# Optionally format and export transformed data to separate table
if transformed_data is not None:
transformed = await workflow.execute_local_activity_method(
Activities.format_transformed_data,
{
**metadata,
'data': transformed_data,
'model_id': input_data['model_id'],
'model_name': input_data['model_name'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
write_transformed_handler = workflow.start_activity_method(
Activities.export_payload_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['transform_table_name'],
'data': transformed,
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
},
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
else:
write_transformed_handler = None
else:
# Error path: create default prediction with error indicators
prediction = await workflow.execute_local_activity_method(
Activities.format_default_prediction,
{
**metadata,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'comment': input_data['comment'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
write_transformed_handler = None
opc_metrics = {}
# write to pi web api
if pi_web_api_output_config:
prediction = await workflow.execute_activity_method(
Activities.write_pi_web_api_data,
{
'pi_web_api_output_config': pi_web_api_output_config,
'data': prediction,
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
# write to opc
if opc_output_config:
prediction, opc_metrics = await workflow.execute_activity_method(
Activities.write_opc_data,
{
'opc_output_config': opc_output_config,
'data': prediction,
**metadata,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)
# write to postgres
await workflow.execute_activity_method(
Activities.export_data_to_postgres,
{
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction,
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
'on_conflict': input_data.get('on_conflict', 'error'),
'unique_columns': ['model_id', 'timestamp'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=180),
)
if write_transformed_handler is not None:
await write_transformed_handler
await workflow.execute_activity_method(
Activities.write_metrics,
{
**metadata,
'prediction': prediction,
'opc_metrics': opc_metrics,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60),
)

View File

@@ -0,0 +1,346 @@
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from datetime import timedelta
from typing import Any
from sientia_do.temporal.policies import retry_policy
from laborious.activities.activities import Activities
@workflow.defn(name='subworkflow.prediction_process')
class PredictionProcess:
"""
Core prediction processing workflow for the Laborious system.
This workflow implements the complete ML model inference pipeline, handling
data quality validation, MLFlow model interactions, and prediction processing.
It serves as the central orchestrator for all prediction operations and ensures
data quality throughout the entire process.
The workflow implements a robust data processing pipeline with:
- Data quality validation using configurable filters
- MLFlow model transformation and prediction
- Response validation and quality assurance
- Flexible decision path handling
- Comprehensive error handling and retry policies
Workflow Execution:
1. Timestamp Retrieval: Gets last processed timestamp for incremental processing
2. Input Data Gate: Applies data quality filters
3. Path Decision: Determines processing path based on filter results
4. MLFlow Transform: Requests data transformation using MLFlow models
5. Response Validation: Filters transform responses for quality assurance
6. MLFlow Prediction: Executes prediction using transformed data
7. Content Validation: Filters prediction responses for final quality check
8. Export Delegation: Delegates to FormatAndExportPrediction workflow
"""
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
Execute the prediction process workflow.
This method orchestrates the complete prediction processing pipeline by:
1. Retrieving the last processed timestamp for incremental processing
2. Applying data quality filters to validate input data
3. Executing MLFlow model transformation and prediction
4. Validating all responses for quality assurance
5. Delegating to export workflow for data persistence
The method implements comprehensive error handling and ensures all
data quality requirements are met before proceeding with ML operations.
Args:
input_data: Complete configuration for the prediction process
Required keys:
- metadata (dict): Workflow execution metadata
- data (dict): Input data for prediction processing
- schema (dict): Data schema definition
- table_name (str): Target table for predictions
- model_id (str): ML model identifier
- model_name (str): ML model name
- input_filters (dict): Data quality filters
- mlflow_transform_filters (dict): MLFlow transform filters
- mlflow_predict_filters (dict): MLFlow prediction filters
- model_retention (int): Model retention period in minutes
- path_priority (list[str]): Decision path priority configuration
- opc_output_config (dict, optional): OPC server export configuration
- pi_web_api_output_config (dict, optional): PI Web API export configuration
- save_transform (bool, optional): Whether to save transformed data (default: True)
- prediction_store_policy (str, optional): Data retention policy (default: 'lts:1')
Returns:
None: The workflow completes successfully when export workflow finishes
Raises:
Exception: If any required parameters are missing or if the workflow fails
during data processing, MLFlow operations, or workflow delegation
"""
metadata = input_data['metadata']
data = input_data['data']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
save_transform = input_data.get('save_transform', True)
try:
await self._run_prediction_pipeline(
input_data,
metadata,
data,
model_id,
model_name,
model_config,
save_transform,
)
await workflow.execute_activity_method(
Activities.cleanup_minio_objects_expired,
{**metadata, 'data': data},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
except Exception as e:
await workflow.execute_activity_method(
Activities.cleanup_minio_objects_expired,
{**metadata, 'data': data},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
raise e
async def _run_prediction_pipeline(
self,
input_data: dict[str, Any],
metadata: dict[str, Any],
data: dict[str, Any],
model_id: str,
model_name: str,
model_config: dict[str, Any],
save_transform: bool,
) -> None:
last_timestamp = data['last_timestamp']
# Apply input data quality gates
gate_input = {
**metadata,
'filters': input_data['input_filters'],
'data': data,
'path_priority': input_data['path_priority'],
}
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.input_gate,
gate_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
# Handle path decision based on filter results
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
# Request MLFlow model transformation
transformed_data = await workflow.execute_activity_method(
Activities.request_transform,
{**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
# Validate MLFlow transform response
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate,
{
**metadata,
'filters': input_data['mlflow_transform_filters'],
'data': transformed_data,
'type': 'transform',
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
# Handle path decision based on transform validation
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_content_gate,
{
**metadata,
'filters': input_data['mlflow_transform_filters'],
'data': transformed_data,
'type': 'transform',
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
predicted_data = await workflow.execute_activity_method(
Activities.request_predict,
{
**metadata,
'data': transformed_data,
'model_name': model_name,
'model_config': model_config,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=5),
)
# Validate MLFlow prediction response
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate,
{
**metadata,
'filters': input_data['mlflow_predict_filters'],
'data': predicted_data,
'type': 'predict',
'path_priority': input_data['path_priority'],
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
# Handle path decision based on prediction validation
if await self.path_flag_handler(
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
# Delegate to export workflow for data persistence
await workflow.execute_child_workflow(
'subworkflow.format_and_export_prediction',
{
'metadata': metadata,
'on_conflict': input_data.get('on_conflict', 'error'),
'path_flag': path_flag,
'data': predicted_data,
'transformed_data': transformed_data if save_transform else None,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model_id,
'model_name': model_name,
'model_config': model_config,
'opc_output_config': input_data['opc_output_config'],
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],
'comment': comment,
'prediction_store_policy': input_data['prediction_store_policy'],
},
)
async def path_flag_handler(
self,
data: dict[str, Any],
path_flag: str,
input_data: dict,
confidence: int,
last_timestamp: str,
comment: str,
) -> bool:
"""
Handle path decisions based on filter results and confidence levels.
This method determines the appropriate action based on the path flag
returned by data quality filters. It can stop processing, continue,
or repeat operations based on the configured path priority.
Args:
data: Input data for processing
path_flag: Path decision from filter (STOP, CONTINUE, REPEAT)
input_data: Complete workflow input configuration including:
- metadata (dict): Workflow execution metadata
- schema (str): Database schema
- table_name (str): Target table for predictions
- transform_table_name (str): Target table for transformed data
- model_id (str): ML model identifier
- model_name (str): ML model name
- model_config (dict, optional): Model configuration
- opc_output_config (dict, optional): OPC server export configuration
- pi_web_api_output_config (dict, optional): PI Web API export configuration
- prediction_store_policy (str, optional): Data retention policy
confidence: Confidence level from filter validation
last_timestamp: Last processed timestamp
comment: Additional information about the filter result
Returns:
bool: True if processing should stop, False to continue
Path Handling:
- STOP: Terminates workflow execution
- CONTINUE: Delegates to FormatAndExportPrediction workflow with current data
- REPEAT: Repeats last prediction if available
"""
metadata = input_data['metadata']
schema = input_data['schema']
table_name = input_data['table_name']
transform_table_name = input_data['transform_table_name']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
path_flag = path_flag.upper() if path_flag else ''
if path_flag == 'STOP':
# Stop processing and exit workflow
return True
elif path_flag == 'REPEAT':
# Repeat last prediction if available
await workflow.execute_activity_method(
Activities.repeat_last_prediction,
{
**metadata,
'schema': schema,
'table_name': table_name,
'model': model_id,
'last_timestamp': last_timestamp,
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
return True
elif path_flag == 'CONTINUE':
# call write workflow
await workflow.execute_child_workflow(
'subworkflow.format_and_export_prediction',
{
'metadata': metadata,
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model_id,
'model_name': model_name,
'model_config': model_config,
'schema': schema,
'table_name': table_name,
'transform_table_name': transform_table_name,
'comment': comment,
'opc_output_config': input_data['opc_output_config'],
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
'prediction_store_policy': input_data['prediction_store_policy'],
'on_conflict': input_data.get('on_conflict', 'error'),
},
)
return True
return False

189
pyproject.toml Normal file
View File

@@ -0,0 +1,189 @@
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "laborious"
version = "0.0.0"
description = "Sientia DataOps Laborious - ML Model Orchestration System"
readme = "README.md"
requires-python = ">=3.11"
authors = [
{name = "Aignosi", email = "dev@aignosi.com"}
]
[tool.ruff]
line-length = 100
target-version = "py311"
exclude = [
".git",
".venv",
"venv",
"__pycache__",
"*.pyc",
".pytest_cache",
"htmlcov",
"tests/laborious/workflows/subworkflows/test_prediction_process.py",
]
[tool.ruff.lint]
select = [
"E", # pycodestyle errors
"W", # pycodestyle warnings
"F", # pyflakes
"I", # isort
"B", # flake8-bugbear
"C4", # flake8-comprehensions
"UP", # pyupgrade
"N", # pep8-naming
"YTT", # flake8-2020
"S", # flake8-bandit
"BLE", # flake8-blind-except
"A", # flake8-builtins
"C90", # mccabe complexity
]
ignore = [
"BLE001", # ignore blind except, we need to send notifications with any error
"E501", # line too long (handled by formatter)
"S101", # use of assert (needed for tests)
"S105", # possible hardcoded password (false positives)
"S106", # possible hardcoded password (false positives)
"S608", # potential sql injection (false positives)
"N802", # function name should be lowercase (temporal decorators)
"N806", # variable in function should be lowercase
]
[tool.ruff.lint.per-file-ignores]
"tests/**/*.py" = [
"S101", # assert allowed in tests
"S105", # hardcoded passwords ok in tests
"S106", # hardcoded passwords ok in tests
]
"e2e/conftest.py" = [
# Imports below the sientia/ModelAnalysis stub are load-order critical: the stub
# must be in sys.modules before Activities is imported, so they cannot move up.
"E402",
]
"scripts/**/*.py" = [
# Notebook-style scripts split into `# %%` cells: each cell imports what it needs
# where it needs it, so imports legitimately appear below other statements.
"E402",
]
[tool.ruff.lint.mccabe]
max-complexity = 15
[tool.ruff.format]
quote-style = "single"
indent-style = "space"
line-ending = "auto"
[tool.mypy]
python_version = "3.11"
warn_return_any = false
warn_unused_configs = true
disallow_untyped_defs = false
disallow_incomplete_defs = false
check_untyped_defs = true
no_implicit_optional = true
warn_redundant_casts = true
warn_unused_ignores = false
warn_no_return = true
strict_equality = true
ignore_missing_imports = true
# Ignore missing imports for external packages
[[tool.mypy.overrides]]
module = "temporalio.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia_do.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "mlflow.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "prometheus_client.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia.*"
ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "pandas.*"
ignore_missing_imports = true
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = [
"-v",
"--strict-markers",
# pytest>=9.1 has a known bug where its unraisableexception plugin crashes
# (tracemalloc partially-initialized AttributeError) when 2+ unraisable
# exceptions land close together — e.g. "coroutine was never awaited" from
# AsyncMock-mocked sync methods (metrics_controller, minio_repository) being
# GC'd. Harmless mock artifacts turned into a hard ERROR by the plugin itself.
"-p", "no:unraisableexception",
]
filterwarnings = [
# `kafka-python`, pulled in transitively by `sientia_do`, reads its protocol schemas through the
# legacy `importlib.resources` API that Python 3.11 deprecated. It emits ~109 warnings per run,
# twice each (once at the call site, once inside `importlib`), which is enough noise to bury a
# warning worth reading. Nothing here can fix it — the change belongs in `kafka-python`.
#
# Matched on the message rather than by module so the filter cannot quietly swallow an unrelated
# deprecation: any *other* DeprecationWarning, including one from our own code, still shows up.
"ignore:read_text is deprecated:DeprecationWarning",
"ignore:open_text is deprecated:DeprecationWarning",
]
markers = [
"asyncio: marks tests as async",
"integration: marks tests as integration tests",
"unit: marks tests as unit tests",
"opc: marks tests that use the in-process OPC UA server (OpcRepository E2E)",
# The .sientia import e2e suite. Outside `testpaths`, so `validate --only-tests` and the
# coverage gate never collect it; run it on request with `pytest e2e -m import_e2e`.
"import_e2e: marks the model-import end-to-end tests (needs Docker and an MLflow server)",
"slow: marks tests that run on a real clock and take tens of seconds",
]
[tool.coverage.run]
source = ["laborious"]
omit = [
"*/tests/*",
"*/venv/*",
"*/__pycache__/*",
"*/site-packages/*",
]
branch = true
[tool.coverage.report]
precision = 2
show_missing = true
skip_covered = false
exclude_lines = [
"pragma: no cover",
"def __repr__",
"def __str__",
"raise AssertionError",
"raise NotImplementedError",
"if __name__ == .__main__.:",
"if TYPE_CHECKING:",
"class .*\\bProtocol\\):",
"@(abc\\.)?abstractmethod",
]
[tool.coverage.html]
directory = "htmlcov"
[tool.bandit]
exclude_dirs = ["tests", "venv", ".venv"]
skips = ["B101", "B601", "B608"] # Skip assert, shell injection, and SQL injection (false positives)

25
requirements-dev.txt Normal file
View File

@@ -0,0 +1,25 @@
# Development and Testing Dependencies
# These packages are only needed for development, testing, and code quality checks
# Install with: pip install -r requirements-dev.txt
# Code Quality & Linting
ruff>=0.1.0 # Fast Python linter and formatter (replaces flake8, black, isort)
mypy>=1.7.0 # Static type checker
bandit>=1.7.5 # Security vulnerability scanner
pandas-stubs>=2.0.0 # Type stubs for pandas
types-requests>=2.31.0 # Type stubs for requests
# Testing
pytest>=7.4.0 # Testing framework
pytest-cov>=4.1.0 # Coverage plugin for pytest
pytest-asyncio>=0.21.0 # Async test support (already in main requirements)
testcontainers[postgres,minio,mongodb]>=4.15.0 # PostgreSQL, MinIO and MongoDB containers for E2E tests
# Floor is the `testcontainers.community.*` namespace: the flat
# `testcontainers.minio` / `.postgres` / `.mongodb` paths the e2e suite used
# to import are deprecated shims that warn on import.
# Development Tools
ipython>=8.12.0 # Enhanced Python shell
ipdb>=0.13.13 # IPython debugger
ipykernel==6.30.1 # IPython kernel for Jupyter notebooks
python-dotenv>=1.0.0 # .env loading for scripts/import_model_manual_run.py

20
requirements-light.txt Normal file
View File

@@ -0,0 +1,20 @@
temporalio
psycopg2-binary
sqlalchemy
asyncua==1.0.6
redis
sientia_do>=1.12.1
mlflow
prometheus-client
botocore
boto3
s3fs
pyarrow
kaleido
hyperopt
shap
pycurl
scipy<1.14.0
scikit-learn==1.4.2
pycaret==3.3.2
joblib

26
requirements-local.txt Normal file
View File

@@ -0,0 +1,26 @@
temporalio
psycopg2-binary
sqlalchemy
asyncua==1.0.6
redis
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.1
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.41.1
prometheus-client
# fsspec/s3fs/gcsfs pinned as a set: nannyml (via the mlops library) pulls both s3fs and
# gcsfs, and each pins fsspec exactly, so a free `s3fs` makes the resolver fail against
# evidently's fsspec>=2024.2.0. Versions match the library-distribution-controller freeze.
fsspec==2026.7.0
s3fs==2026.7.0
gcsfs==2026.8.0
# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3.
botocore==1.43.56
boto3==1.43.56
pyarrow
kaleido
hyperopt
shap
pycurl
scipy<1.14.0
scikit-learn==1.4.2
pycaret==3.3.2
joblib

33
requirements.txt Normal file
View File

@@ -0,0 +1,33 @@
temporalio
psycopg2-binary
sqlalchemy
asyncua==1.0.6
redis
sientia_do>=1.12.1
sientia>=0.41.1
prometheus-client
# fsspec/s3fs/gcsfs pinned as a set: nannyml (via sientia) pulls both s3fs and gcsfs, and
# each pins fsspec exactly, so a free `s3fs` makes the resolver fail against evidently's
# fsspec>=2024.2.0. Versions match the library-distribution-controller freeze.
fsspec==2026.7.0
s3fs==2026.7.0
gcsfs==2026.8.0
# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3.
botocore==1.43.56
boto3==1.43.56
pyarrow
kaleido
hyperopt
shap
pycurl
scipy<1.14.0
scikit-learn==1.4.2
pycaret==3.3.2
joblib
# .sientia bundle import (Argon2id + XChaCha20-Poly1305 secretstream): same pin as the
# producer in sientia-projects-templates, so both sides open the same bytes.
pynacl==1.5.0
# AES-256-GCM password envelope; already present transitively, declared explicitly.
cryptography
# MongoDB models collection (the model listing); transitive via sientia_do, declared explicitly.
pymongo

View File

@@ -0,0 +1,502 @@
"""Run one `.sientia` model import by hand, standing in for the frontend that does not exist yet.
**What this is.** The Angular screen plus the Spring Boot BFF will, between them, pick a `.sientia`
file, hash it, wrap the bundle password for the worker, upload the object and insert one `PENDING`
row into `public.experiment_run` before starting `import_model`. None of that exists yet. This
script does exactly those four things, against whatever environment `.env` points at, so the import
path can be exercised for real and so whoever writes the frontend has a working reference for the
wire shapes (the digest, the envelope, the object key, the row).
**What it writes.** Two things, and it can undo both: one object under
`{IMPORT_BUNDLE_BUCKET}/{IMPORT_BUNDLE_PREFIX}` in MinIO, and one row in `public.experiment_run` in
the BFF Postgres database. Everything else it does is a read. If the workflow runs, *it* creates the
MLflow experiment, run and registered version and the MongoDB model document — cell 10 can remove
those too, behind its own flag.
**The safety rule.** Every target comes from the environment, and this file contains no host, user,
password or key literal. A `.env` inherited from a cluster deployment can perfectly well name a
production database and a production bucket, so cell 5 and cell 6 — the only cells that write — are
gated on cell 2, where the operator has to restate the Postgres host that cell 1 printed. This is a
development tool. It is not for production, and the guard is a mitigation, not a guarantee.
**Cells.** Run them one at a time in the VS Code interactive window, or run the module top to
bottom. Cell order is the contract:
1 resolve every target and print it 6 insert the PENDING row -> import_run_id
2 confirmation guard (first write is after) 7 print the workflow input
3 pick the bundle, compute expected_digest 8 optionally start the workflow and wait
4 build the password envelope 9 read the outcome back
5 upload the object, verify the stored bytes 10 undo
**Boundaries.** This script imports nothing from `laborious.activities` or `laborious.workflows` —
it stands in for the frontend, so it must not borrow the consumer's code. The one exception is
`laborious.utils.bundle.format`, which is pure `struct`/`hashlib` over bytes and lets cell 3 print
the header of the file it is about to upload. The envelope is built here in four visible lines
rather than imported, because those four lines are what the frontend has to transliterate to
WebCrypto; `laborious/utils/import_password.py` is the decrypt counterpart.
"""
# %%
# Cell 1 — resolve every target from the environment and print what we are pointed at.
# Reads only. No secret is printed here or anywhere else in this script.
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import Any
from dotenv import load_dotenv
REPO_ROOT = Path(__file__).resolve().parent.parent
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
load_dotenv(REPO_ROOT / '.env')
def env(name: str, default: str = '') -> str:
"""The environment variable, stripped, or `default`."""
return os.getenv(name, default).strip()
def env_required(name: str) -> str:
"""The environment variable, or a `RuntimeError` naming it."""
value = env(name)
if not value:
raise RuntimeError(f'{name} is not set — see .env.example')
return value
def _line(label: str, value: object) -> str:
"""One aligned `label value` line, for output that stays readable in a terminal."""
return f' {label:<24} {value}'
# Every connection target below is *required* rather than defaulted. The worker's own config falls
# back to localhost, which is right for a service that must boot; it is wrong for a script that
# writes, because a silent fallback is a write to somewhere nobody read out loud.
#
# The import log database: the POSTGRES_* server from .env, with only the database name of its own,
# because the log lives in the Spring Boot BFF's database rather than in `sientia`. The schema and
# table are constants of the worker's code, never configuration.
#
# `build_import_status_config` also honours IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} as overrides,
# for a deployment whose log sits on another server. This script does not: it is a development tool
# pointed at a real environment, and the override exists for the e2e suite's two containers.
STATUS_DB_HOST = env_required('POSTGRES_HOST')
STATUS_DB_PORT = int(env_required('POSTGRES_PORT'))
STATUS_DB_NAME = env_required('IMPORT_STATUS_DB_NAME')
STATUS_DB_USER = env_required('POSTGRES_USER')
STATUS_DB_PASSWORD = os.getenv('POSTGRES_PASSWORD', '')
STATUS_TABLE = 'public.experiment_run'
# Where the frontend uploads the bundle. The prefix keeps the worker's own default, which is a
# format decision rather than a target.
MINIO_ENDPOINT_URL = env_required('MINIO_ENDPOINT_URL')
MINIO_ACCESS_KEY = env_required('MINIO_ACCESS_KEY')
MINIO_SECRET_KEY = env_required('MINIO_SECRET_KEY')
BUNDLE_BUCKET = env_required('IMPORT_BUNDLE_BUCKET')
_raw_prefix = env('IMPORT_BUNDLE_PREFIX', 'imported_models/')
BUNDLE_PREFIX = _raw_prefix if _raw_prefix.endswith('/') else f'{_raw_prefix}/'
# Read-back targets. This script never writes to these; the workflow does.
MONGODB_URL = env_required('MONGODB_URL')
MONGODB_USERNAME = env('MONGODB_USERNAME')
MONGODB_PASSWORD = os.getenv('MONGODB_PASSWORD', '')
MONGODB_DATABASE = env_required('MONGODB_DATABASE')
MODELS_COLLECTION = env('IMPORT_MODELS_COLLECTION', 'models')
MLFLOW_HOST = env_required('MLFLOW_HOST')
MLFLOW_PORT = env_required('MLFLOW_PORT')
MLFLOW_TRACKING_URI = f'{MLFLOW_HOST}:{MLFLOW_PORT}'
TEMPORAL_HOST = env_required('TEMPORAL_HOST')
TEMPORAL_NAMESPACE = env_required('TEMPORAL_NAMESPACE')
RUNTIME = env_required('RUNTIME')
IMPORT_TASK_QUEUE = f'import_model-{RUNTIME}-queue'
# This script's own inputs.
BUNDLE_PATH = Path(
env('IMPORT_SCRIPT_BUNDLE_PATH', 'tests/fixtures/bundle/golden_model_v3.sientia')
)
if not BUNDLE_PATH.is_absolute():
BUNDLE_PATH = REPO_ROOT / BUNDLE_PATH
IMPORT_USERNAME = env_required('IMPORT_SCRIPT_USERNAME')
print('resolved targets — nothing has been written yet')
print(_line('status postgres', f'{STATUS_DB_HOST}:{STATUS_DB_PORT}/{STATUS_DB_NAME}'))
print(_line('status user', STATUS_DB_USER))
print(_line('status table', STATUS_TABLE))
print(_line('minio endpoint', MINIO_ENDPOINT_URL))
print(_line('minio bucket/prefix', f'{BUNDLE_BUCKET}/{BUNDLE_PREFIX}'))
print(_line('mongo', f'{MONGODB_URL}/{MONGODB_DATABASE}.{MODELS_COLLECTION}'))
print(_line('mlflow', MLFLOW_TRACKING_URI))
print(_line('temporal', f'{TEMPORAL_HOST} ns={TEMPORAL_NAMESPACE}'))
print(_line('task queue', IMPORT_TASK_QUEUE))
print(_line('bundle', BUNDLE_PATH))
print(_line('username', IMPORT_USERNAME))
print()
print(f'to continue, set the cell 2 confirmation to: {STATUS_DB_HOST}')
# %%
# Cell 2 — the confirmation guard. Cells 5 and 6 write; this is the last cell before them.
#
# Restate the Postgres host cell 1 printed. Either set IMPORT_SCRIPT_CONFIRM_TARGET in .env, or
# replace the right-hand side below with the value cell 1 told you. It is a read-then-restate on
# purpose: a boolean flag can be flipped without ever looking at what it is pointed at.
CONFIRM_TARGET = env('IMPORT_SCRIPT_CONFIRM_TARGET')
if CONFIRM_TARGET != STATUS_DB_HOST:
raise RuntimeError(
'refusing to write: the confirmation does not match the resolved status database host. '
f'confirmation={CONFIRM_TARGET!r}, resolved POSTGRES_HOST={STATUS_DB_HOST!r}. '
'Read the targets cell 1 printed, then restate the host here.'
)
print(f'target confirmed: {STATUS_DB_HOST} — the write cells may run')
# %%
# Cell 3 — pick the bundle, compute the digest the frontend must send, print its header.
#
# `expected_digest` is the SHA-256 of the exact bytes that get uploaded, computed the way the
# frontend computes it: over the file, not over anything derived from it. The header fields printed
# here are plaintext at the front of the file — reading them decrypts nothing.
import hashlib
from laborious.utils.bundle.format import DIGEST_OFFSET, HEADER_SIZE, parse_header
bundle_bytes = BUNDLE_PATH.read_bytes()
expected_digest = hashlib.sha256(bundle_bytes).hexdigest()
object_name = BUNDLE_PATH.name
object_key = f'{BUNDLE_PREFIX}{object_name}'
header = parse_header(bundle_bytes[:HEADER_SIZE])
print(_line('file', BUNDLE_PATH))
print(_line('size', f'{len(bundle_bytes)} bytes'))
print(_line('expected_digest', expected_digest))
print(_line('object key', f'{BUNDLE_BUCKET}/{object_key}'))
print(_line('format version', header.format_version))
print(_line('aead / kdf id', f'{header.aead_id} / {header.kdf_id}'))
print(_line('kdf memlimit', f'{header.kdf_memlimit_bytes} bytes'))
print(_line('kdf opslimit', header.kdf_opslimit))
print(_line('kdf parallelism', header.kdf_parallelism))
print(_line('salt', header.salt.hex()))
print(_line('ciphertext digest', header.ciphertext_digest.hex()))
print(_line('chunk stream from', f'offset {DIGEST_OFFSET}'))
# %%
# Cell 4 — build the password envelope, exactly as the frontend must build it.
#
# Wire shape: base64( nonce(12) || ciphertext || GCM tag(16) ) under AES-256-GCM with the 32-byte
# key that `IMPORT_PASSWORD_KEY` carries base64-encoded. The four lines below are the whole of it,
# written out rather than imported so the Angular side can transliterate them to WebCrypto:
#
# const k = await crypto.subtle.importKey('raw', keyBytes, 'AES-GCM', false, ['encrypt']);
# const nonce = crypto.getRandomValues(new Uint8Array(12));
# const sealed = await crypto.subtle.encrypt({name: 'AES-GCM', iv: nonce}, k, passwordBytes);
# const envelope = base64(concat(nonce, new Uint8Array(sealed)));
#
# WebCrypto appends the GCM tag to the ciphertext, and so does `cryptography`'s AESGCM, so the two
# produce byte-identical payloads. The decrypt counterpart the worker runs is
# `laborious/utils/import_password.py` (`decrypt_password_envelope`).
#
# Neither the key nor the password nor the plaintext is printed — only the envelope's length.
import base64
import secrets
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
envelope_key = base64.b64decode(env_required('IMPORT_PASSWORD_KEY'), validate=True)
nonce = secrets.token_bytes(12)
sealed = AESGCM(envelope_key).encrypt(
nonce, env_required('IMPORT_SCRIPT_BUNDLE_PASSWORD').encode('utf-8'), None
)
password_envelope = base64.b64encode(nonce + sealed).decode('ascii')
print(_line('envelope key', f'{len(envelope_key)} bytes (not printed)'))
print(_line('envelope', f'{len(password_envelope)} base64 characters (not printed)'))
# %%
# Cell 5 — FIRST WRITE. Upload the object, then prove the stored bytes are the ones we hashed.
#
# The re-read is the point: `expected_digest` is a promise about what is in the bucket, and the
# worker's gate 1 holds the upload to it. If the digest of the stored object differs, the frontend
# contract is already broken and there is no reason to insert a row.
import io
from urllib.parse import urlparse
from minio import Minio
_endpoint = urlparse(MINIO_ENDPOINT_URL)
minio_client = Minio(
_endpoint.netloc or MINIO_ENDPOINT_URL,
access_key=MINIO_ACCESS_KEY,
secret_key=MINIO_SECRET_KEY,
secure=_endpoint.scheme == 'https',
)
minio_client.put_object(
BUNDLE_BUCKET,
object_key,
io.BytesIO(bundle_bytes),
length=len(bundle_bytes),
content_type='application/octet-stream',
)
stat = minio_client.stat_object(BUNDLE_BUCKET, object_key)
_response = minio_client.get_object(BUNDLE_BUCKET, object_key)
try:
stored_bytes = _response.read()
finally:
_response.close()
_response.release_conn()
stored_digest = hashlib.sha256(stored_bytes).hexdigest()
if stored_digest != expected_digest:
raise RuntimeError(
f'the stored object does not hash to the digest we computed: stored {stored_digest}, '
f'expected {expected_digest} — do not insert a row for this upload'
)
print(_line('uploaded', f'{BUNDLE_BUCKET}/{object_key}'))
print(_line('reported size', f'{stat.size} bytes'))
print(_line('etag', stat.etag))
print(_line('re-read digest', f'{stored_digest} (matches expected_digest)'))
# %%
# Cell 6 — SECOND WRITE. Insert the import log row: the BFF's half of the contract.
#
# `URL.create` rather than a URL string, so a password containing `@ : /` survives. `experiment_name`
# is a placeholder: the column is NOT NULL and the frontend cannot see inside an encrypted bundle,
# so the worker overwrites it with the bundle's own experiment name as soon as the bundle is open
# (`record_import_names`). `file_name` is the object key, because the worker compares the two and
# refuses the import at the `received` step when they disagree.
import json
from datetime import UTC, datetime
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
status_engine = create_engine(
URL.create(
'postgresql',
username=STATUS_DB_USER,
password=STATUS_DB_PASSWORD,
host=STATUS_DB_HOST,
port=STATUS_DB_PORT,
database=STATUS_DB_NAME,
),
pool_pre_ping=True,
)
request_data = {
'expected_digest': expected_digest,
'bucket': BUNDLE_BUCKET,
'object_key': object_key,
'file_name': object_name,
'source': 'scripts/import_model_manual_run.py',
}
_now = datetime.now(UTC).replace(tzinfo=None)
_placeholder_experiment = f'import-{Path(object_name).stem}'[:50]
INSERT_ROW_SQL = text(
'INSERT INTO public.experiment_run '
'(experiment_name, username, status, created_at, updated_at, '
' bucket_name, file_name, request_data, run_type) '
'VALUES (:experiment_name, :username, :status, :created_at, :updated_at, '
' :bucket_name, :file_name, CAST(:request_data AS JSONB), :run_type) '
'RETURNING id'
)
with status_engine.begin() as connection:
import_run_id = connection.execute(
INSERT_ROW_SQL,
{
'experiment_name': _placeholder_experiment,
'username': IMPORT_USERNAME,
'status': 'PENDING',
'created_at': _now,
'updated_at': _now,
'bucket_name': BUNDLE_BUCKET,
'file_name': object_key,
'request_data': json.dumps(request_data),
'run_type': 'IMPORT',
},
).scalar_one()
print(_line('import_run_id', import_run_id))
print(_line('status / run_type', 'PENDING / IMPORT'))
print(_line('experiment_name', f'{_placeholder_experiment} (the worker overwrites this)'))
print(_line('file_name', object_key))
# %%
# Cell 7 — the workflow input, ready to paste into the Temporal UI.
#
# These five keys are the whole contract. `bucket` is optional — the worker falls back to
# IMPORT_BUNDLE_BUCKET — and is sent explicitly so the printed input is self-describing. No model
# name is sent: the worker takes it from the bundle's own metadata, verbatim.
workflow_input = {
'import_run_id': import_run_id,
'bucket': BUNDLE_BUCKET,
'object_key': object_key,
'expected_digest': expected_digest,
'password_envelope': password_envelope,
}
print(f'workflow: import_model task queue: {IMPORT_TASK_QUEUE}')
print(json.dumps(workflow_input, indent=2))
# %%
# Cell 8 — optionally start the workflow and wait for it.
#
# Skipping this cell is normal: the operator may prefer to start the run from the Temporal UI with
# the JSON cell 7 printed. Cell 9 works either way, because it looks the outcome up by row id.
#
# The coroutine runs on a loop of its own in a worker thread rather than through `asyncio.run`, so
# this cell behaves the same in the interactive window (which already has a running loop) as it
# does when the module is executed top to bottom (which does not).
START_THE_WORKFLOW = False
if START_THE_WORKFLOW:
import asyncio
from concurrent.futures import ThreadPoolExecutor
from temporalio.client import Client
async def start_and_wait() -> Any:
"""Start `import_model` on the import queue and return its result."""
client = await Client.connect(TEMPORAL_HOST, namespace=TEMPORAL_NAMESPACE)
handle = await client.start_workflow(
'import_model',
workflow_input,
id=f'manual-import-{import_run_id}',
task_queue=IMPORT_TASK_QUEUE,
)
print(_line('workflow id', handle.id))
print(_line('run id', handle.result_run_id))
return await handle.result()
with ThreadPoolExecutor(max_workers=1) as _pool:
workflow_result = _pool.submit(asyncio.run, start_and_wait()).result()
print(json.dumps(workflow_result, indent=2, default=str))
else:
print('not started — paste the input from cell 7 into the Temporal UI, or set')
print(f'START_THE_WORKFLOW = True and re-run this cell. Task queue: {IMPORT_TASK_QUEUE}')
# %%
# Cell 9 — read the outcome back from all three places it lands.
#
# The row is authoritative for the verdict; MLflow and MongoDB are what the verdict is about. Each
# block is labelled, so the output can be pasted into a ticket as it is.
print('--- import log row ---')
READ_ROW_SQL = text(
'SELECT id, experiment_name, run_name, status, error_message, run_type, '
' bucket_name, file_name, created_at, updated_at, '
' request_data, orchestrator_response_data '
'FROM public.experiment_run WHERE id = :id'
)
with status_engine.begin() as connection:
_result = connection.execute(READ_ROW_SQL, {'id': import_run_id})
row = _result.mappings().one_or_none()
if row is None:
print(f' no row with id {import_run_id}')
else:
for _column, _value in row.items():
print(_line(_column, _value))
print()
print('--- mlflow ---')
import mlflow
os.environ['MLFLOW_TRACKING_USERNAME'] = env('MLFLOW_USERNAME')
os.environ['MLFLOW_TRACKING_PASSWORD'] = os.getenv('MLFLOW_PASSWORD', '')
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
mlflow_client = mlflow.tracking.MlflowClient()
_experiment_name = row['experiment_name'] if row is not None else ''
_experiment = mlflow_client.get_experiment_by_name(_experiment_name) if _experiment_name else None
if _experiment is None:
print(f' no experiment named {_experiment_name!r}')
else:
print(_line('experiment', f'{_experiment.name} ({_experiment.experiment_id})'))
for _run in mlflow_client.search_runs([_experiment.experiment_id]):
print(_line('run', f'{_run.info.run_name} ({_run.info.run_id}) {_run.info.status}'))
for _artifact in mlflow_client.list_artifacts(_run.info.run_id):
print(_line(' artifact', _artifact.path))
# The run name is `import-{model_name}-v{version}` (`derive_run_name`), so the registered model
# name is what sits between the prefix and the version suffix.
_run_name = (row['run_name'] if row is not None else '') or ''
_registered = _run_name.removeprefix('import-').rsplit('-v', 1)[0] if _run_name else ''
if _registered:
try:
for _version in mlflow_client.search_model_versions(f"name='{_registered}'"):
print(
_line(
'version',
f'{_version.name} v{_version.version} stage={_version.current_stage} '
f'run={_version.run_id}',
)
)
print(_line(' source', _version.source))
except Exception as error:
# A read-back cell reports what it could not read; it does not fail the inspection.
print(f' registry read failed: {type(error).__name__}: {error}')
else:
print(' no run_name on the row yet, so no registered model to look up')
print()
print('--- mongodb model document ---')
from pymongo import MongoClient
_credentials = f'{MONGODB_USERNAME}:{MONGODB_PASSWORD}@' if MONGODB_USERNAME else ''
mongo_client: MongoClient = MongoClient(f'mongodb://{_credentials}{MONGODB_URL}')
_collection = mongo_client[MONGODB_DATABASE][MODELS_COLLECTION]
for _document in _collection.find({}, {'_id': False}).sort('id', 1):
print(_line('document', _document))
# %%
# Cell 10 — undo.
#
# The object and the row always: they are what this script created, and leaving them behind leaves a
# PENDING import nobody will ever run. What MLflow and MongoDB hold is left alone unless
# DELETE_PROVISIONING is set, because after a successful run that is the thing you wanted to look at.
DELETE_PROVISIONING = False
try:
minio_client.remove_object(BUNDLE_BUCKET, object_key)
print(_line('deleted object', f'{BUNDLE_BUCKET}/{object_key}'))
except Exception as error:
# Undo reports what it could not remove and carries on to the row.
print(f' could not delete the object: {type(error).__name__}: {error}')
with status_engine.begin() as connection:
_deleted = connection.execute(
text('DELETE FROM public.experiment_run WHERE id = :id'),
{'id': import_run_id},
).rowcount
print(_line('deleted rows', _deleted))
if DELETE_PROVISIONING:
if _experiment is not None:
mlflow_client.delete_experiment(_experiment.experiment_id)
print(_line('deleted experiment', _experiment.experiment_id))
if _registered:
try:
mlflow_client.delete_registered_model(_registered)
print(_line('deleted model', _registered))
except Exception as error:
# Same rule as above: report and carry on to the document.
print(f' could not delete the registered model: {type(error).__name__}: {error}')
_removed = _collection.delete_many({'name': _registered}).deleted_count
print(_line('deleted documents', _removed))
else:
print(' provisioning kept — set DELETE_PROVISIONING = True to remove it too')
status_engine.dispose()
mongo_client.close()

11
scripts/run_coverage.sh Executable file
View File

@@ -0,0 +1,11 @@
#!/bin/bash
# Exit on any error
set -e
echo "Activating virtual environment..."
source ./venv/bin/activate
pytest --cov=laborious --cov-report=html
xdg-open htmlcov/index.html

18
scripts/run_local.sh Executable file
View File

@@ -0,0 +1,18 @@
#!/bin/bash
# Exit on any error
set -e
echo "Activating virtual environment..."
source ./venv/bin/activate
echo "Loading environment variables from .env..."
if [ -f .env ]; then
export $(cat .env | grep -v '^#' | xargs)
echo "Environment variables loaded from .env"
else
echo "Warning: .env file not found. Continuing without environment variables."
fi
echo "Starting ingestor application..."
python -m laborious.worker.worker

11
sonar-project.properties Normal file
View File

@@ -0,0 +1,11 @@
sonar.projectKey=Aignosi_sientia-dataops-laborious_temporal_ca1a7039-6db9-49e5-be78-54d29bc93e4f
sonar.projectName=sientia-dataops-laborious_temporal
sonar.sources=laborious
sonar.tests=tests
sonar.projectVersion=1.0.0
sonar.coverage.exclusions=laborious/worker/*
sonar.qualitygate.wait=true
sonar.qualitygate.timeout=300
sonar.python.coverage.reportPaths=coverage.xml
sonar.python.xunit.reportPath=pytest.xml
sonar.python.version=3.11