Using Amazon Aurora machine learning with Aurora PostgreSQL

By using Amazon Aurora machine learning with your Aurora PostgreSQL DB cluster, you can use Amazon Comprehend or Amazon SageMaker or Amazon Bedrock, depending on your needs. These services each support specific machine learning use cases.

Aurora machine learning is supported in certain AWS Regions and for specific versions of Aurora PostgreSQL only. Before trying to set up Aurora machine learning, check availability for your Aurora PostgreSQL version and your Region. For details, see Aurora machine learning with Aurora PostgreSQL.

Requirements for using Aurora machine learning with Aurora PostgreSQL

AWS machine learning services are managed services that are set up and run in their own production environments. Aurora machine learning supports integration with Amazon Comprehend, SageMaker, and Amazon Bedrock. Before trying to set up your Aurora PostgreSQL DB cluster to use Aurora machine learning, be sure you understand the following requirements and prerequisites.

Supported features and limitations of Aurora machine learning with Aurora PostgreSQL

Aurora machine learning supports any SageMaker endpoint that can read and write the comma-separated value (CSV) format through a ContentType value of text/csv. The built-in SageMaker algorithms that currently accept this format are the following.

When using Amazon Bedrock with Aurora machine learning, the following limitations apply:

Setting up your Aurora PostgreSQL DB cluster to use Aurora machine learning

For Aurora machine learning to work with your Aurora PostgreSQL DB cluster, you need to create an AWS Identity and Access Management (IAM) role for each of the services that you want to use. The IAM role allows your Aurora PostgreSQL DB cluster to use the Aurora machine learning service on the cluster's behalf. You also need to install the Aurora machine learning extension. In the following topics, you can find setup procedures for each of these Aurora machine learning services.

Setting up Aurora PostgreSQL to use Amazon Bedrock

In the procedure following, you first create the IAM role and policy that gives your Aurora PostgreSQL permission to use Amazon Bedrock on the cluster's behalf. You then attach the policy to an IAM role that your Aurora PostgreSQL DB cluster uses to work with Amazon Bedrock. For simplicity's sake, this procedure uses the AWS Management Console to complete all tasks.

The role (with its policy) are associated with the Aurora PostgreSQL DB cluster. When the process completes, the role is listed in the Current IAM roles for this cluster listing, as shown following.

The IAM setup for Amazon Bedrock is complete. Continue setting up your Aurora PostgreSQL to work with Aurora machine learning by installing the extension as detailed in Installing the Aurora machine learning extension

Setting up Aurora PostgreSQL to use Amazon Comprehend

In the procedure following, you first create the IAM role and policy that gives your Aurora PostgreSQL permission to use Amazon Comprehend on the cluster's behalf. You then attach the policy to an IAM role that your Aurora PostgreSQL DB cluster uses to work with Amazon Comprehend For simplicity's sake, this procedure uses the AWS Management Console to complete all tasks.

The role (with its policy) are associated with the Aurora PostgreSQL DB cluster. When the process completes, the role is listed in the Current IAM roles for this cluster listing, as shown following.

The IAM setup for Amazon Comprehend is complete. Continue setting up your Aurora PostgreSQL to work with Aurora machine learning by installing the extension as detailed in Installing the Aurora machine learning extension

Setting up Aurora PostgreSQL to use Amazon SageMaker

Before you can create the IAM policy and role for your Aurora PostgreSQL DB cluster, you need to have your SageMaker model setup and your endpoint available.

The role (with its policy) are associated with the Aurora PostgreSQL DB cluster. When the process completes, the role is listed in the Current IAM roles for this cluster listing.

The IAM setup for SageMaker is complete. Continue setting up your Aurora PostgreSQL to work with Aurora machine learning by installing the extension as detailed in Installing the Aurora machine learning extension.

Setting up Aurora PostgreSQL to use Amazon S3 for SageMaker (Advanced)

To use SageMaker with your own models rather than using the pre-built components provided by SageMaker, you need to set up an Amazon Simple Storage Service (Amazon S3) bucket for Aurora PostgreSQL DB cluster to use. This is an advanced topic, and isn't fully documented in this Amazon Aurora User Guide. The general process is the same as for integrating support for SageMaker, as follows.

For basic usage information, see Exporting data to Amazon S3 for SageMaker model training (Advanced).

Installing the Aurora machine learning extension

The Aurora machine learning extensions aws_ml 1.0 provides two functions that you can use to invoke Amazon Comprehend, SageMaker services and aws_ml 2.0 provides two additional functions that you can use to invoke Amazon Bedrock services. Installing these extensions on your Aurora PostgreSQL DB cluster also creates an administrative role for the feature.

labdb=> CREATE EXTENSION IF NOT EXISTS aws_ml CASCADE;
NOTICE:  installing required extension "aws_commons"
CREATE EXTENSION
labdb=> 

Installing the aws_ml extensions also creates the aws_ml administrative role and three new schemas, as follows.

The rds_superuser role is granted the aws_ml administrative role and is made the OWNER of these three Aurora machine learning schemas. To allow other database users to access the Aurora machine learning functions, the rds_superuser needs to grant EXECUTE privileges on the Aurora machine learning functions. By default, EXECUTE privileges are revoked from PUBLIC on the functions in the two Aurora machine learning schemas.

In a multi-tenant database configuration, you can prevent tenants from accessing Aurora machine learning functions by using REVOKE USAGE on the specific Aurora machine learning schema that you want to protect.

Using Amazon Bedrock with your Aurora PostgreSQL DB cluster

For Aurora PostgreSQL, Aurora machine learning provides the following Amazon Bedrock function for working with your text data. This function is available only after you install the aws_ml 2.0 extension and complete all setup procedures. For more information, see Setting up your Aurora PostgreSQL DB cluster to use Aurora machine learning.

The inputs and outputs of this function are as follows.

The following example shows how to invoke a Anthropic Claude 2 model for Bedrock using invoke_model.

SELECT aws_bedrock.invoke_model (
    model_id    := 'anthropic.claude-v2',
    content_type:= 'application/json',
    accept_type := 'application/json',
    model_input := '{"prompt": "\n\nHuman: You are a helpful assistant that answers questions directly and only using the information provided in the context below.\nDescribe the answer 
    in detail.\n\nContext: %s \n\nQuestion: %s \n\nAssistant:","max_tokens_to_sample":4096,"temperature":0.5,"top_k":250,"top_p":0.5,"stop_sequences":[]}'
);

The inputs and outputs of this function are as follows.

The following example shows how to generate an embedding using the Titan Embeddings G1 – Text embedding model for the phrase PostgreSQL I/O monitoring views.

SELECT aws_bedrock.invoke_model_get_embeddings(
   model_id      := 'amazon.titan-embed-text-v1',
   content_type  := 'application/json',
   json_key      := 'embedding',
   model_input   := '{ "inputText": "PostgreSQL I/O monitoring views"}') AS embedding;

Using Amazon Comprehend with your Aurora PostgreSQL DB cluster

For Aurora PostgreSQL, Aurora machine learning provides the following Amazon Comprehend function for working with your text data. This function is available only after you install the aws_ml extension and complete all setup procedures. For more information, see Setting up your Aurora PostgreSQL DB cluster to use Aurora machine learning.

The inputs and outputs of this function are as follows.

In the following, you can find examples of how to use this function.

SELECT feedback, s.sentiment,s.confidence 
    FROM support,aws_comprehend.detect_sentiment(feedback, 'en') s 
    ORDER BY s.confidence DESC;
 feedback                         | sentiment | confidence
 ----------------------------------------------------------+-----------+------------
 Thank you for the excellent customer support!            | POSITIVE  |   0.999771
 The latest version of this product stinks!               | NEGATIVE  |   0.999184
 Your support team is just awesome! I am blown away.      | POSITIVE  |   0.997774
 Your product is too complex, but your support is great.  | MIXED     |   0.957958
 Your support tech helped me in fifteen minutes.          | POSITIVE  |   0.949491
 My problem was never resolved!                           | NEGATIVE  |   0.920644
 When will the new version of this product be released?   | NEUTRAL   |   0.902706
 I cannot stand that chatbot.                             | NEGATIVE  |   0.895219
 Your support tech talked down to me.                     | NEGATIVE  |   0.868598
 It took me way too long to get a real person.            | NEGATIVE  |   0.481805
 
 (10 rows)

To avoid being charged for sentiment detection more than once per table row, you can materialize the results. Do this on the rows of interest. For example, the clinician's notes are being updated so that only those in French (fr) use the sentiment detection function.

UPDATE clinician_notes
SET sentiment = (aws_comprehend.detect_sentiment (french_notes, 'fr')).sentiment,
    confidence = (aws_comprehend.detect_sentiment (french_notes, 'fr')).confidence
WHERE
    clinician_notes.french_notes IS NOT NULL AND
    LENGTH(TRIM(clinician_notes.french_notes)) > 0 AND
    clinician_notes.sentiment IS NULL;

For more information on optimizing your function calls, see Performance considerations for using Aurora machine learning with Aurora PostgreSQL.

Using SageMaker with your Aurora PostgreSQL DB cluster

After setting up your SageMaker environment and integrating with Aurora PostgreSQL as outlined in Setting up Aurora PostgreSQL to use Amazon SageMaker , you can invoke operations by using the aws_sagemaker.invoke_endpoint function. The aws_sagemaker.invoke_endpoint function connects only to a model endpoint in the same AWS Region. If your database instance has replicas in multiple AWS Regions be sure that you setup and deploy each SageMaker model to every AWS Region.

Calls to aws_sagemaker.invoke_endpoint are authenticated using the IAM role that you set up to associated your Aurora PostgreSQL DB cluster with the SageMaker service and the endpoint that you provided during the setup process. SageMaker model endpoints are scoped to an individual account and are not public. The endpoint_name URL doesn't contain the account ID. SageMaker determines the account ID from the authentication token that is supplied by the SageMaker IAM role of the database instance.

The inputs and outputs of this function are as follows.

Creating a user-defined function to invoke a SageMaker model

Create a separate user-defined function to call aws_sagemaker.invoke_endpoint for each of your SageMaker models. Your user-defined function represents the SageMaker endpoint hosting the model. The aws_sagemaker.invoke_endpoint function runs within the user-defined function. User-defined functions provide many advantages:

To specify a user-defined function, use the SQL data definition language (DDL) statement CREATE FUNCTION. When you define the function, you specify the following:

The user-defined function returns the inference computed by the SageMaker endpoint after running the model on the input parameters. The following example creates a user-defined function for an SageMaker model with two input parameters.

CREATE FUNCTION classify_event (IN arg1 INT, IN arg2 DATE, OUT category INT)
AS $$
    SELECT aws_sagemaker.invoke_endpoint (
        'sagemaker_model_endpoint_name', NULL,
        arg1, arg2                        -- model inputs are separate arguments
        )::INT                            -- cast the output to INT
$$ LANGUAGE SQL PARALLEL SAFE COST 5000;

Note the following:

Passing an array as input to a SageMaker model

The aws_sagemaker.invoke_endpoint function can have up to 100 input parameters, which is the limit for PostgreSQL functions. If the SageMaker model requires more than 100 parameters of the same type, pass the model parameters as an array.

The following example defines a function that passes an array as input to the SageMaker regression model. The output is cast to a REAL value.

CREATE FUNCTION regression_model (params REAL[], OUT estimate REAL)
AS $$
    SELECT aws_sagemaker.invoke_endpoint (
      'sagemaker_model_endpoint_name',
      NULL,
      params                            
      )::REAL
$$ LANGUAGE SQL PARALLEL SAFE COST 5000;

Specifying batch size when invoking a SageMaker model

The following example creates a user-defined function for a SageMaker model that sets the batch size default to NULL. The function also allows you to provide a different batch size when you invoke it.

CREATE FUNCTION classify_event (
    IN event_type INT, IN event_day DATE, IN amount REAL, -- model inputs
    max_rows_per_batch INT DEFAULT NULL,  -- optional batch size limit
    OUT category INT)                     -- model output
AS $$
    SELECT aws_sagemaker.invoke_endpoint (
        'sagemaker_model_endpoint_name', max_rows_per_batch,
        event_type, event_day, COALESCE(amount, 0.0)
        )::INT              -- casts output to type INT
$$ LANGUAGE SQL PARALLEL SAFE COST 5000;

Note the following:

Invoking a SageMaker model that has multiple outputs

The following example creates a user-defined function for a SageMaker model that returns multiple outputs. Your function needs to cast the output of the aws_sagemaker.invoke_endpoint function to a corresponding data type. For example, you could use the built-in PostgreSQL point type for (x,y) pairs or a user-defined composite type.

This user-defined function returns values from a model that returns multiple outputs by using a composite type for the outputs.

CREATE TYPE company_forecasts AS ( 
    six_month_estimated_return real,
    one_year_bankruptcy_probability float);
CREATE FUNCTION analyze_company (
    IN free_cash_flow NUMERIC(18, 6),
    IN debt NUMERIC(18,6),
    IN max_rows_per_batch INT DEFAULT NULL,
    OUT prediction company_forecasts) 
AS $$
    SELECT (aws_sagemaker.invoke_endpoint('endpt_name',
       max_rows_per_batch,free_cash_flow, debt))::company_forecasts;                                                                                                                  
$$ LANGUAGE SQL PARALLEL SAFE COST 5000;

For the composite type, use fields in the same order as they appear in the model output and cast the output of aws_sagemaker.invoke_endpoint to your composite type. The caller can extract the individual fields either by name or with PostgreSQL ".*" notation.

Exporting data to Amazon S3 for SageMaker model training (Advanced)

We recommend that you become familiar with Aurora machine learning and SageMaker by using the provided algorithms and examples rather than trying to train your own models. For more information, see Get Started with Amazon SageMaker

To train SageMaker models, you export data to an Amazon S3 bucket. The Amazon S3 bucket is used by SageMaker to train your model before it is deployed. You can query data from an Aurora PostgreSQL DB cluster and save it directly into text files stored in an Amazon S3 bucket. Then SageMaker consumes the data from the Amazon S3 bucket for training. For more about SageMaker model training, see Train a model with Amazon SageMaker.

For more information about exporting your data, see Exporting data from an Aurora PostgreSQL DB cluster to Amazon S3.

Performance considerations for using Aurora machine learning with Aurora PostgreSQL

The Amazon Comprehend and SageMaker services do most of the work when invoked by an Aurora machine learning function. That means that you can scale those resources as needed, independently. For your Aurora PostgreSQL DB cluster, you can make your function calls as efficient as possible. Following, you can find some performance considerations to note when working with Aurora machine learning from Aurora PostgreSQL.

Understanding batch mode and Aurora machine learning functions

Typically, PostgreSQL runs functions one row at a time. Aurora machine learning can reduce this overhead by combining the calls to the external Aurora machine learning service for many rows into batches with an approach called batch-mode execution. In batch mode, Aurora machine learning receives the responses for a batch of input rows, and then delivers the responses back to the running query one row at a time. This optimization improves the throughput of your Aurora queries without limiting the PostgreSQL query optimizer.

Aurora automatically uses batch mode if the function is referenced from the SELECT list, a WHERE clause, or a HAVING clause. Note that top-level simple CASE expressions are eligible for batch-mode execution. Top-level searched CASE expressions are also eligible for batch-mode execution provided that the first WHEN clause is a simple predicate with a batch-mode function call.

Your user-defined function must be a LANGUAGE SQL function and should specify PARALLEL SAFE and COST 5000.

Function migration from the SELECT statement to the FROM clause

Usually, an aws_ml function that is eligible for batch-mode execution is automatically migrated by Aurora to the FROM clause.

The migration of eligible batch-mode functions to the FROM clause can be examined manually on a per-query level. To do this, you use EXPLAIN statements (and ANALYZE and VERBOSE) and find the "Batch Processing" information below each batch-mode Function Scan. You can also use EXPLAIN (with VERBOSE) without running the query. You then observe whether the calls to the function appear as a Function Scan under a nested loop join that was not specified in the original statement.

In the following example, the nested loop join operator in the plan shows that Aurora migrated the anomaly_score function. It migrated this function from the SELECT list to the FROM clause, where it's eligible for batch-mode execution.

EXPLAIN (VERBOSE, COSTS false) 
SELECT anomaly_score(ts.R.description) from ts.R;
                         QUERY PLAN                          
-------------------------------------------------------------
 Nested Loop
   Output: anomaly_score((r.description)::text)
   ->  Seq Scan on ts.r
         Output: r.id, r.description, r.score
   ->  Function Scan on public.anomaly_score
         Output: anomaly_score.anomaly_score
         Function Call: anomaly_score((r.description)::text)

To disable batch-mode execution, set the apg_enable_function_migration parameter to false. This prevents the migration of aws_ml functions from the SELECT to the FROM clause. The following shows how.

SET apg_enable_function_migration = false;

The apg_enable_function_migration parameter is a Grand Unified Configuration (GUC) parameter that is recognized by the Aurora PostgreSQL apg_plan_mgmt extension for query plan management. To disable function migration in a session, use query plan management to save the resulting plan as an approved plan. At runtime, query plan management enforces the approved plan with its apg_enable_function_migration setting. This enforcement occurs regardless of the apg_enable_function_migration GUC parameter setting. For more information, see Managing query execution plans for Aurora PostgreSQL.

Using the max_rows_per_batch parameter

Both the aws_comprehend.detect_sentiment and the aws_sagemaker.invoke_endpoint functions have a max_rows_per_batch parameter. This parameter specifies the number of rows that can be sent to the Aurora machine learning service. The larger the dataset processed by your function, the larger you can make the batch size.

Batch-mode functions improve efficiency by building batches of rows that spread the cost of the Aurora machine learning function calls over a large number of rows. However, if a SELECT statement finishes early due to a LIMIT clause, then the batch can be constructed over more rows than the query uses. This approach can result in additional charges to your AWS account. To gain the benefits of batch-mode execution but avoid building batches that are too large, use a smaller value for the max_rows_per_batch parameter in your function calls.

If you do an EXPLAIN (VERBOSE, ANALYZE) of a query that uses batch-mode execution, you see a FunctionScan operator that is below a nested loop join. The number of loops reported by EXPLAIN equals the number of times a row was fetched from the FunctionScan operator. If a statement uses a LIMIT clause, the number of fetches is consistent. To optimize the size of the batch, set the max_rows_per_batch parameter to this value. However, if the batch-mode function is referenced in a predicate in the WHERE clause or HAVING clause, then you probably can't know the number of fetches in advance. In this case, use the loops as a guideline and experiment with max_rows_per_batch to find a setting that optimizes performance.

Verifying batch-mode execution

To see if a function ran in batch mode, use EXPLAIN ANALYZE. If batch-mode execution was used, then the query plan will include the information in a "Batch Processing" section.

EXPLAIN ANALYZE SELECT user-defined-function();
 Batch Processing: num batches=1 avg/min/max batch size=3333.000/3333.000/3333.000
                                avg/min/max batch call time=146.273/146.273/146.273

In this example, there was 1 batch that contained 3,333 rows, which took 146.273 ms to process. The "Batch Processing" section shows the following:

Typically the final batch is smaller than the rest, which often results in a minimum batch size that is much smaller than the average.

To return the first few rows more quickly, set the max_rows_per_batch parameter to a smaller value.

To reduce the number of batch mode calls to the ML service when you use a LIMIT in your user-defined function, set the max_rows_per_batch parameter to a smaller value.

Improving response times with parallel query processing

To get results as fast as possible from a large number of rows, you can combine parallel query processing with batch mode processing. You can use parallel query processing for SELECT, CREATE TABLE AS SELECT, and CREATE MATERIALIZED VIEW statements.

Parallel query processing occurs both within the database and within the ML service. The number of cores in the instance class of the database limits the degree of parallelism that can be used when running a query. The database server can construct a parallel query execution plan that partitions the task among a set of parallel workers. Then each of these workers can build batched requests containing tens of thousands of rows (or as many as are allowed by each service).

The batched requests from all of the parallel workers are sent to the SageMaker endpoint. The degree of parallelism that the endpoint can support is constrained by the number and type of instances that support it. For K degrees of parallelism, you need a database instance class that has at least K cores. You also need to configure the SageMaker endpoint for your model to have K initial instances of a sufficiently high-performing instance class.

To use parallel query processing, you can set the parallel_workers storage parameter of the table that contains the data that you plan to pass. You set parallel_workers to a batch-mode function such as aws_comprehend.detect_sentiment. If the optimizer chooses a parallel query plan, the AWS ML services can be called both in batch and in parallel.

You can use the following parameters with the aws_comprehend.detect_sentiment function to get a plan with four-way parallelism. If you change either of the following two parameters, you must restart the database instance for the changes to take effect

-- SET max_worker_processes to 8;  -- default value is 8
-- SET max_parallel_workers to 8;  -- not greater than max_worker_processes
SET max_parallel_workers_per_gather to 4;  -- not greater than max_parallel_workers

-- You can set the parallel_workers storage parameter on the table that the data
-- for the Aurora machine learning function is coming from in order to manually override the degree of
-- parallelism that would otherwise be chosen by the query optimizer
--
ALTER TABLE yourTable SET (parallel_workers = 4);

-- Example query to exploit both batch-mode execution and parallel query
EXPLAIN (verbose, analyze, buffers, hashes) 
SELECT aws_comprehend.detect_sentiment(description, 'en')).*
FROM yourTable
WHERE id < 100;

For more information about controlling parallel query, see Parallel plans in the PostgreSQL documentation.

Using materialized views and materialized columns

When you invoke an AWS service such as SageMaker or Amazon Comprehend from your database, your account is charged according to the pricing policy of that service. To minimize charges to your account, you can materialize the result of calling the AWS service into a materialized column so that the AWS service is not called more than once per input row. If desired, you can add a materializedAt timestamp column to record the time at which the columns were materialized.

The latency of an ordinary single-row INSERT statement is typically much less than the latency of calling a batch-mode function. Thus, you might not be able to meet the latency requirements of your application if you invoke the batch-mode function for every single-row INSERT that your application performs. To materialize the result of calling an AWS service into a materialized column, high-performance applications generally need to populate the materialized columns. To do this, they periodically issue an UPDATE statement that operates on a large batch of rows at the same time.

UPDATE takes a row-level lock that can impact a running application. So you might need to use SELECT ... FOR UPDATE SKIP LOCKED, or use MATERIALIZED VIEW.

Analytic queries that operate on a large number of rows in real time can combine batch-mode materialization with real-time processing. To do this, these queries assemble a UNION ALL of the pre-materialized results with a query over the rows that don't yet have materialized results. In some cases, such a UNION ALL is needed in multiple places, or the query is generated by a third-party application. If so, you can create a VIEW to encapsulate the UNION ALL operation so this detail isn't exposed to the rest of the SQL application.

You can use a materialized view to materialize the results of an arbitrary SELECT statement at a snapshot in time. You can also use it to refresh the materialized view at any time in the future. Currently PostgreSQL doesn't support incremental refresh, so each time the materialized view is refreshed the materialized view is fully recomputed.

You can refresh materialized views with the CONCURRENTLY option, which updates the contents of the materialized view without taking an exclusive lock. Doing this allows a SQL application to read from the materialized view while it's being refreshed.

Monitoring Aurora machine learning

You can monitor the aws_ml functions by setting the track_functions parameter in your custom DB cluster parameter group to all. By default, this parameter is set to pl which means that only procedure-language functions are tracked. By changing this to all, the aws_ml functions are also tracked. For more information, see Run-time Statistics in the PostgreSQL documentation.

For information about monitoring the performance of the SageMaker operations called from Aurora machine learning functions, see Monitor Amazon SageMaker in the Amazon SageMaker Developer Guide.

With track_functions set to all, you can query the pg_stat_user_functions view to get statistics about the functions that you define and use to invoke Aurora machine learning services. For each function, the view provides the number of calls, total_time, and self_time.

To view the statistics for the aws_sagemaker.invoke_endpoint and the aws_comprehend.detect_sentiment functions, you can filter results by schema name using the following query.

SELECT * FROM pg_stat_user_functions 
    WHERE schemaname 
    LIKE 'aws_%';

To clear the statistics, do as follows.

SELECT pg_stat_reset();

You can get the names of your SQL functions that call the aws_sagemaker.invoke_endpoint function by querying the PostgreSQL pg_proc system catalog. This catalog stores information about functions, procedures, and more. For more information, see pg_proc in the PostgreSQL documentation. Following is an example of querying the table to get the names of functions (proname) whose source (prosrc) includes the text invoke_endpoint.

SELECT proname FROM pg_proc WHERE prosrc LIKE '%invoke_endpoint%';