Implementing a Data Capture System for Machine Learning Platforms on AWS: A Practical Approach
In the previous article on Data Capture (Introduction to Data Capture in Machine Learning Platforms on AWS: A Theoretical Approach) in Machine Learning (ML) platforms, the key concepts of the system and which AWS services can be the basis for building it in a modular and efficient way were explored.
Get Jose Cerezo Egea’s stories in your inbox
Join Medium for free to get updates from this writer.
This article shows the creation, from a practical point of view, of a simple system that captures input and output data from a model, stores it in S3, validates it using AWS Lambda, and ingests it into distributed transactional tables.
System Design
The data capture system has been designed to ensure efficient collection, validation, and storage of the information generated in each phase of the ML model life cycle. The architecture is composed of the following key elements:
- Python decorators. A decorator is made available to the user that allows data capture with total flexibility and transparency. It requires some parameters such as the the database and the table in which to record the data. It is responsible for storing them in a file in AWS S3, which allows their persistence and further processing.
- Event Trigger S3. To automate the processing flow, an event is configured in the AWS S3 bucket in which the files containing the data captured by the decorator are deposited and which acts as a trigger. Each time a new file is uploaded to the bucket, the event triggers an AWS Lambda function in charge of data validation and transformation.
- AWS Lambda. When invoked by the AWS S3 event, this function is responsible for validating that the data to be persisted is correct and corresponds to the schema, defined in AWS Glue, of the target table. It also applies the necessary transformations before storing them in the corresponding distributed transactional table. This function must have an alarm method to warn users if the ingest cannot be performed correctly.
- AWS Glue. To facilitate data analysis and querying, AWS Glue catalogs and organizes the data into a structured schema. There is also the option for the user to have the ability to create tables, which allows more control when defining them, for example, for a specific type of column and its name.
- AWS Athena. Allows the execution of SQL queries on this data, without the need to provision additional infrastructure, optimizing the access and analysis of the information efficiently.
This modular architecture allows the automation of the data capture process, guaranteeing its integrity and facilitating its subsequent exploitation for the generation of reports and advanced analysis.
Practical Implementation
Python Decorator for Data Capture
As discussed in the previous article, a Python decorator is used as a tool to capture the input and output data of an ML model in S3, since its use does not modify the main logic of the decorated functions (data preparation, functions to predict, etc). An example implementation of the capture decorator is presented below:
import pandas as pd
from datetime import datetime
def capture_data(bucket: str, database: str, table: str):
def decorator(func):
def wrapper(*args, **kwargs):
# Execute the function and capture the results
result = func(*args, **kwargs)
timestamp = datetime.utcnow().strftime("%Y/%m/%d/%H")
# Organized route for S3
full_path = f"s3://{bucket}/{database}/{table}/{timestamp}/output.parquet"
# Saving data in S3
if isinstance(result, pd.DataFrame):
result.to_parquet(full_path, engine='pyarrow', index=False)
print(f"Data captured in: {full_path}")
else:
raise ValueError("The output must be a DataFrame.")
return result
return wrapper
return decorator
# Example of use
@capture_data(bucket="mybucket", database="mydb", table="mytable")
def preprocess(data) -> pd.DataFrame:
# <your data preprocessing>
return preprocessed_dataThe key points of the decorator’s code are:
- Parquet format. The file with the data to be captured is uploaded to S3 in parquet format, as it is efficient in storage and compatible with Athena.
- S3 path. It is important to define within the organization a structure to store the data in S3, as it will later affect the way it is consumed. In this case the path is <database>/<table>/<timestamp>, but it could be extended for further segregation. An improvement on how to manage this timestamp field will be presented later in the best practices section.
- Arguments of the decorator. In this example there are 3 (bucket, database and table), but you could have more or less, depending on the need. For example, the bucket in which to store the data to capture could always be the same and not an argument to the user, but a constant.
S3 and Lambda Event Configuration
The S3 bucket where the parquet file with the data to be captured is deposited must have previously configured an Event notification with:
- Event types → Put, Multipart upload completed (to be launched only in case of file uploads).
- Filter on the suffix for “.parquet”.
- Destination type → Lambda function.
This notification invokes the lambda that is in charge of validating the schema of the data to be captured, performing the necessary transformations, and finally persisting the data in the corresponding table. The code of the lambda could be similar to:
import boto3
import pandas as pd
from io import BytesIO
s3 = boto3.client('s3')
def lambda_handler(event, context):
# Get uploaded file information
bucket = event['Records'][0]['s3']['bucket']['name']
key = event['Records'][0]['s3']['object']['key']
database, table = extract_db_table_from_key(key)
# Download the S3 file
response = s3.get_object(Bucket=bucket, Key=key)
file_content = response['Body'].read()
# Read file as DataFrame
df = pd.read_parquet(BytesIO(file_content))
# Validate the scheme
expected_columns = ['feature1', 'feature2']
if not all(col in df.columns for col in expected_columns):
raise ValueError("The file does not comply with the expected schema.")
# <your additional processing (optional)>
# Persist data in distributed transactional tables
insert_data_in_table(df, database, table)
return {"statusCode": 200, "body": "Successful processing"}The insert_data_in_table function is in charge of persisting the data in the corresponding distributed transitional table, a simple example of this function could be:
import pandas as pd
def insert_data_in_table(df, database, table):
# Generate a string representation of column names and values from a DataFrame
columns_str, values_str = generate_df_str(df)
# Create the SQL query
insert_query = f'INSERT INTO {database}.{table} ({columns_str}) VALUES {values_str}'
# Execute Athena query
start = time.time()
print(f'Start query: {query}')
response = pd.read_sql(insert_query)
print('Query done after {}'.format(time.time() - start))
return responseIntegration with Glue and Athena
Once the data has been persisted in the tables, Glue provides the schema for these tables so that they are easily exploitable, among others, by Athena, which through SQL queries offers the possibility to visualize the captured data.
This is an example of features that we may have captured during the training of a model, although it would also be interesting to capture columns such as the prediction id, the date on which the capture was performed, the predicted class, etc.
Best Practices
During the development of the component, a series of good practices have been learned that we recommend to take into account:
- Organization in S3. It is key to design clear and consistent paths to facilitate data access, analysis, permissions management, query efficiency, etc. For this case we propose a segregation by database, table and date of data capture, although more levels could be added as needed.
- Error Management. Consider moving the files that do not pass the validations to a graveyard for later analysis, so that data that are not correct will not be lost, as they could be useful later. Continuing with the previous example, they could be saved in the same path, but changing “s3://<bucket>/datacapture/…” to “s3://<bucket>/graveyard-datacapture/…”.
- Monitoring and Alerts. Configure AWS CloudWatch to receive alerts when errors occur in both the decorator and Lambda execution. These alarms could detect critical errors in the platform that need to be solved urgently, in addition to having greater control and monitoring of the component.
- Cost Optimization. Use efficient formats, such as Parquet, and define and configure an appropriate lifecycle for your data. For example, in the case of the graveyard you could set up a Transition action that defines at what point those files should be moved to another storage class so that they are cheaper to keep in storage. Another example would be, after the correct persistence of the data in the corresponding table, the Parquet file could be moved to S3 Glacier.
Conclusion and Next Steps
The modular and scalable system discussed demonstrates how it is possible to implement a basic data capture flow with AWS, through Python decorators, which provide complete flexibility and transparency to the user using them, S3 triggers to make the capture as immediate as possible, and validations/transformations of the data within the Lambda. This ensures that the data in an ML pipeline is captured correctly for future uses, such as triggering alarms in case of data bias, audits, etc.
In the next and final article, we will explore how to scale this approach for multi-entity platforms, addressing challenges such as privacy, cost separation, and protection of sensitive data.
Who we are
José Cerezo is an AI/ML Engineer with the role of Consultant in the Data Science practice unit at SDG Group España. He has experience in sectors such as insurance, defining and implementing different ML platforms in the AWS ecosystem. www.linkedin.com/in/jose-cerezo-egea/
Rubén Rodríguez is an ML architect and specialist participating in the architecture and methodology area of the Data Science practice unit at SDG Group España. He has experience in different sectors such as pharmaceuticals, insurance, telecommunications, and other sectors, managing different technologies in AWS ecosystems. www.linkedin.com/in/ruben-rodriguez-cardos-97a9708a/
