Merge pull request #691 from NOAA-OWP/viz-rds-parameter-group-tuning

- Upgraded viz-processing RDS Instance Type from "db.m5.2xlarge" to "db.m5.4xlarge".
- Adjusted various parameters in RDS Parameter group for viz-processing RDS Instance to better optimize memory usage.
- Added concurrency limit to the "Input Data Files" step of the viz processing pipeline Step Function.
- Updated all Lambda Functions and Viz Lambda Layer to better utilize and dispose of RDS database connections.
  - Updated all instances of DB connections and cursors to use context managers.
  - Explicitly closed all DB connections after they are no longer used.
  - Removed unneeded .commit() calls.
  - Removed any redundant and unused DB connections.
- Temporarily commented out the coastal services in the initialize pipeline Lambda function to stop them from updating.
- Fixed S3 Object downloading issue in hand fim processing Lambda by changing it to stream directly from the bucket and added retries for when it occasionally failed.

Core/LAMBDA/layers/viz_lambda_shared_funcs/python/viz_classes.py

@@ -91,14 +91,14 @@ def load_df_into_db(self, table_name, df, drop_first=True):
     ###################################
     def run_sql_file_in_db(self, sql_file):
         sql = open(sql_file, 'r').read()
-        with self.connection as db_connection:
+        with self.connection:
             try:
-                cur = db_connection.cursor()
-                print(f"---> Running {sql_file}")
-                cur.execute(sql)
-                db_connection.commit()
+                with self.connection.cursor() as cur:
+                    print(f"---> Running {sql_file}")
+                    cur.execute(sql)
             except Exception as e:
                 raise e
+        self.connection.close()
 
     ###################################                
     def run_sql_in_db(self, sql, return_geodataframe=False):
@@ -119,21 +119,23 @@ def run_sql_in_db(self, sql, return_geodataframe=False):
     ###################################
     def get_est_row_count_in_table(self, table):
         print(f"Getting estimated total rows in {table}.")
-        with self.connection as db_connection:
+        with self.connection:
             try:
-                cur = db_connection.cursor()
-                sql = f"""
-                SELECT (CASE WHEN c.reltuples < 0 THEN NULL -- never vacuumed
-                            WHEN c.relpages = 0 THEN float8 '0' -- empty table
-                            ELSE c.reltuples / c.relpages END
-                    * (pg_catalog.pg_relation_size(c.oid) / pg_catalog.current_setting('block_size')::int))::bigint
-                FROM   pg_catalog.pg_class c
-                WHERE  c.oid = '{table}'::regclass; -- schema-qualified table here
-                """
-                cur.execute(sql)
-                rows = cur.fetchone()[0]
+                with self.connection.cursor() as cur:
+                    sql = f"""
+                    SELECT (CASE WHEN c.reltuples < 0 THEN NULL -- never vacuumed
+                                WHEN c.relpages = 0 THEN float8 '0' -- empty table
+                                ELSE c.reltuples / c.relpages END
+                        * (pg_catalog.pg_relation_size(c.oid) / pg_catalog.current_setting('block_size')::int))::bigint
+                    FROM   pg_catalog.pg_class c
+                    WHERE  c.oid = '{table}'::regclass; -- schema-qualified table here
+                    """
+                    cur.execute(sql)
+                    rows = cur.fetchone()[0]
             except Exception as e:
                 raise e
+        self.connection.close()
+
         return rows
 
     ###################################

Core/LAMBDA/layers/viz_lambda_shared_funcs/python/viz_lambda_shared_funcs.py

@@ -456,12 +456,14 @@ def get_service_metadata(include_ingest_sources=True, include_latest_ref_time=Fa
 		                ON a.ingest_table = b.target group by service) AS ref_times
                         on admin.services.service = ref_times.service2"""
 
+
     connection = get_db_connection("viz")
-    with connection.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-        cur.execute(f"SELECT * FROM admin.services{extra_sql};")
-        column_names = [desc[0] for desc in cur.description]
-        response = cur.fetchall()
-        cur.close()
+    with connection:
+        with connection.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
+            cur.execute(f"SELECT * FROM admin.services{extra_sql};")
+            column_names = [desc[0] for desc in cur.description]
+            response = cur.fetchall()
+            # cur.close()
     connection.close()
     return list(map(lambda x: dict(zip(column_names, x)), response))
 
@@ -493,17 +495,16 @@ def load_df_into_db(table_name, db_engine, df):
 def run_sql_file_in_db(db_type, sql_file):
     print("Getting connection to run sql files")
     sql = open(sql_file, 'r').read()
-    db_connection = get_db_connection(db_type)
 
-    try:
-        cur = db_connection.cursor()
-        print(f"Running {sql_file}")
-        cur.execute(sql)
-        db_connection.commit()
-    except Exception as e:
-        raise e
-    finally:
-        db_connection.close()
+    connection = get_db_connection(db_type)
+    with connection:
+        try:
+            with connection.cursor() as cur:
+                print(f"Running {sql_file}")
+                cur.execute(sql)
+        except Exception as e:
+            raise e
+    connection.close()
 
 
 def move_data_to_another_db(origin_db, dest_db, origin_table, dest_table, stage=True, add_oid=True,

Core/LAMBDA/rnr_functions/rnr_domain_generator/lambda_function.py

@@ -33,11 +33,12 @@ def lambda_handler(event, context):
     if step == 'domain':
         print(f'Executing {step}.sql...')
 
-        with viz_db.get_db_connection() as connection:
-            cur = connection.cursor()
-            cur.execute(sql)
-            result = cur.fetchone()
-            connection.commit()
+        connection = viz_db.get_db_connection()
+        with connection:
+            with connection.cursor() as cur:
+                cur.execute(sql)
+                result = cur.fetchone()
+        connection.close()
 
         reference_time = dt.datetime.strptime(result[0], '%Y-%m-%d %H:%M:%S UTC').strftime(DT_FORMAT)
         run_time_obj = dt.datetime.strptime(event['run_time'], '%Y-%m-%dT%H:%M:%SZ')

Core/LAMBDA/viz_functions/image_based/viz_hand_fim_processing/lambda_function.py

@@ -1,7 +1,9 @@
 import boto3
+from botocore.exceptions import ResponseStreamingError
 import rasterio
 import numpy as np
 import pandas as pd
+import awswrangler as wr
 import geopandas as gpd
 import os
 import time
@@ -14,8 +16,6 @@
 FIM_PREFIX = os.environ['FIM_PREFIX']
 FIM_VERSION = re.findall("[/_]?(\d*_\d*_\d*_\d*)/?", FIM_PREFIX)[0]
 
-s3 = boto3.client("s3")
-
 class HANDDatasetReadError(Exception):
     """ my custom exception class """
 
@@ -436,15 +436,19 @@ def process(window):
 
     return df_final
 
-def s3_csv_to_df(bucket, key):
-    basename = os.path.basename(key)
-    local_file = f"/tmp/{basename}"
+def s3_csv_to_df(bucket, key):    
+    # Allow retrying a few times before failing
+    for i in range(5):
+        try:
+            # Read S3 csv file into Pandas DataFrame
+            print(f"Reading {key} from {bucket} into DataFrame")
+            df = wr.s3.read_csv(path=f"s3://{bucket}/{key}")
+            print("DataFrame creation Successful")
+        except ResponseStreamingError:
+            if i == 4: print("Failed to read from S3")
+            continue
+        break
 
-    print(f"Downloading {key} from {bucket}")
-    s3.download_file(bucket, key, local_file)
-    df = pd.read_csv(local_file)
-    os.remove(local_file)
-
     return df
 
 def calculate_stage_values(hydrotable_key, subsetted_streams_bucket, subsetted_streams, huc8_branch):

Core/LAMBDA/viz_functions/image_based/viz_hand_fim_processing/requirements.txt

@@ -3,3 +3,5 @@ geopandas==0.11.1
 SQLAlchemy==1.4.40
 GeoAlchemy2==0.12.3
 psycopg2-binary==2.9.3
+pandas==2.1.4
+awswrangler

Core/LAMBDA/viz_functions/image_based/viz_schism_fim_processing/lambda_function.py

@@ -543,30 +543,32 @@ def setup_db_table(db_fim_table, viz_db, sql_replace=None):
     db_schema = db_fim_table.split('.')[0]
 
     print(f"Setting up {db_fim_table}")
-
-    with viz_db.get_db_connection() as connection:
-        cur = connection.cursor()
 
-         # See if the target table exists #TODO: Ensure table exists would make a good helper function
-        cur.execute(f"SELECT EXISTS (SELECT FROM pg_tables WHERE schemaname = '{db_fim_table.split('.')[0]}' AND tablename = '{db_fim_table.split('.')[1]}');")
-        table_exists = cur.fetchone()[0]
-
-        # If the target table doesn't exist, create one basd on the sql_replace dict.
-        if not table_exists:
-            print(f"--> {db_fim_table} does not exist. Creating now.")
-            original_table = list(sql_replace.keys())[list(sql_replace.values()).index(db_fim_table)] #ToDo: error handling if not in list
-            cur.execute(f"DROP TABLE IF EXISTS {db_fim_table}; CREATE TABLE {db_fim_table} (LIKE {original_table})")
+    connection = viz_db.get_db_connection()
+    with connection:
+        with connection.cursor() as cur:
+            # See if the target table exists #TODO: Ensure table exists would make a good helper function
+            cur.execute(f"SELECT EXISTS (SELECT FROM pg_tables WHERE schemaname = '{db_fim_table.split('.')[0]}' AND tablename = '{db_fim_table.split('.')[1]}');")
+            table_exists = cur.fetchone()[0]
+
+            # If the target table doesn't exist, create one basd on the sql_replace dict.
+            if not table_exists:
+                print(f"--> {db_fim_table} does not exist. Creating now.")
+                original_table = list(sql_replace.keys())[list(sql_replace.values()).index(db_fim_table)] #ToDo: error handling if not in list
+                cur.execute(f"DROP TABLE IF EXISTS {db_fim_table}; CREATE TABLE {db_fim_table} (LIKE {original_table})")
+                connection.commit()
+
+            # Drop the existing index on the target table
+            print("Dropping target table index (if exists).")
+            SQL = f"DROP INDEX IF EXISTS {db_schema}.{index_name};"
+            cur.execute(SQL)
+
+            # Truncate all records.
+            print("Truncating target table.")
+            SQL = f"TRUNCATE TABLE {db_fim_table};"
+            cur.execute(SQL)
             connection.commit()
 
-        # Drop the existing index on the target table
-        print("Dropping target table index (if exists).")
-        SQL = f"DROP INDEX IF EXISTS {db_schema}.{index_name};"
-        cur.execute(SQL)
-
-        # Truncate all records.
-        print("Truncating target table.")
-        SQL = f"TRUNCATE TABLE {db_fim_table};"
-        cur.execute(SQL)
-        connection.commit()
+    connection.close()
 
     return db_fim_table

Core/LAMBDA/viz_functions/viz_db_ingest/lambda_function.py

@@ -56,80 +56,80 @@ def lambda_handler(event, context):
         return json.dumps(dump_dict)
 
     viz_db = database(db_type="viz")
-    with viz_db.get_db_connection() as connection:
-        cursor = connection.cursor()
-        try:
-            nwm_version = 0
+    try:
+        nwm_version = 0
 
-            if file.endswith('.nc'):
-                ds = xr.open_dataset(download_path)
-                ds_vars = [var for var in ds.variables]
+        if file.endswith('.nc'):
+            ds = xr.open_dataset(download_path)
+            ds_vars = [var for var in ds.variables]
 
-                if not target_cols:
-                    target_cols = ds_vars
+            if not target_cols:
+                target_cols = ds_vars
 
-                try:
-                    if "hawaii" in file:
-                        ds['forecast_hour'] = int(int(re.findall("(\d{8})/[a-z0-9_]*/.*t(\d{2})z.*[ftm](\d*)\.", file)[0][-1])/100)
-                    else:
-                        ds['forecast_hour'] = int(re.findall("(\d{8})/[a-z0-9_]*/.*t(\d{2})z.*[ftm](\d*)\.", file)[0][-1])
-                    if 'forecast_hour' not in target_cols:
-                        target_cols.append('forecast_hour')
-                except:
-                    print("Regex pattern for the forecast hour didn't match the netcdf file")
-                
-                try:
-                    ds['nwm_vers'] = float(ds.NWM_version_number.replace("v",""))
-                    if 'nwm_vers' not in target_cols:
-                        target_cols.append('nwm_vers')
-                except:
-                    print("NWM_version_number property is not available in the netcdf file")
+            try:
+                if "hawaii" in file:
+                    ds['forecast_hour'] = int(int(re.findall("(\d{8})/[a-z0-9_]*/.*t(\d{2})z.*[ftm](\d*)\.", file)[0][-1])/100)
+                else:
+                    ds['forecast_hour'] = int(re.findall("(\d{8})/[a-z0-9_]*/.*t(\d{2})z.*[ftm](\d*)\.", file)[0][-1])
+                if 'forecast_hour' not in target_cols:
+                    target_cols.append('forecast_hour')
+            except:
+                print("Regex pattern for the forecast hour didn't match the netcdf file")
+
+            try:
+                ds['nwm_vers'] = float(ds.NWM_version_number.replace("v",""))
+                if 'nwm_vers' not in target_cols:
+                    target_cols.append('nwm_vers')
+            except:
+                print("NWM_version_number property is not available in the netcdf file")
 
-                drop_vars = [var for var in ds_vars if var not in target_cols]
-                df = ds.to_dataframe().reset_index()
-                df = df.drop(columns=drop_vars)
-                ds.close()
-                if 'streamflow' in target_cols:
-                    df = df.loc[df['streamflow'] >= keep_flows_at_or_above].round({'streamflow': 2}).copy()  # noqa
-                df = df[target_cols]
+            drop_vars = [var for var in ds_vars if var not in target_cols]
+            df = ds.to_dataframe().reset_index()
+            df = df.drop(columns=drop_vars)
+            ds.close()
+            if 'streamflow' in target_cols:
+                df = df.loc[df['streamflow'] >= keep_flows_at_or_above].round({'streamflow': 2}).copy()  # noqa
+            df = df[target_cols]
 
-            elif file.endswith('.csv'):
-                df = pd.read_csv(download_path)
-                for column in df:  # Replace any 'None' strings with nulls
-                    df[column].replace('None', np.nan, inplace=True)
-                df = df.copy()
-            else:
-                print("File format not supported.")
-                exit()
+        elif file.endswith('.csv'):
+            df = pd.read_csv(download_path)
+            for column in df:  # Replace any 'None' strings with nulls
+                df[column].replace('None', np.nan, inplace=True)
+            df = df.copy()
+        else:
+            print("File format not supported.")
+            exit()
 
-            print(f"--> Preparing and Importing {file}")
-            f = StringIO()  # Use StringIO to store the temporary text file in memory (faster than on disk)
-            df.to_csv(f, sep='\t', index=False, header=False)
-            f.seek(0)
-            try:
-                with viz_db.get_db_connection() as connection:
-                    cursor = connection.cursor()
-                    cursor.copy_expert(f"COPY {target_table} FROM STDIN WITH DELIMITER E'\t' null as ''", f)
-                    connection.commit()
-            except (UndefinedTable, BadCopyFileFormat, InvalidTextRepresentation):
-                if not create_table:
-                    raise
+        print(f"--> Preparing and Importing {file}")
+        f = StringIO()  # Use StringIO to store the temporary text file in memory (faster than on disk)
+        df.to_csv(f, sep='\t', index=False, header=False)
+        f.seek(0)
+        try:
+            connection = viz_db.get_db_connection()
+            with connection:
+                with connection.cursor() as cur:
+                    cur.copy_expert(f"COPY {target_table} FROM STDIN WITH DELIMITER E'\t' null as ''", f)
+            connection.close()
+        except (UndefinedTable, BadCopyFileFormat, InvalidTextRepresentation):
+            if not create_table:
+                raise
 
-                print("Error encountered. Recreating table now and retrying import...")
-                create_table_df = df.head(0)
-                schema, table = target_table.split('.')
-                create_table_df.to_sql(con=viz_db.engine, schema=schema, name=table, index=False, if_exists='replace')
-                with viz_db.get_db_connection() as connection:
-                    cursor = connection.cursor()
-                    cursor.copy_expert(f"COPY {target_table} FROM STDIN WITH DELIMITER E'\t' null as ''", f)
-                    connection.commit()
+            print("Error encountered. Recreating table now and retrying import...")
+            create_table_df = df.head(0)
+            schema, table = target_table.split('.')
+            create_table_df.to_sql(con=viz_db.engine, schema=schema, name=table, index=False, if_exists='replace')
+            connection = viz_db.get_db_connection()
+            with connection:
+                with connection.cursor() as cur:
+                    cur.copy_expert(f"COPY {target_table} FROM STDIN WITH DELIMITER E'\t' null as ''", f)
+            connection.close()
 
-            print(f"--> Import of {len(df)} rows Complete. Removing {download_path} and closing db connection.")
-            os.remove(download_path)
-    
-        except Exception as e:
-            print(f"Error: {e}")
-            raise e
+        print(f"--> Import of {len(df)} rows Complete. Removing {download_path} and closing db connection.")
+        os.remove(download_path)
+
+    except Exception as e:
+        print(f"Error: {e}")
+        raise e
 
     dump_dict = {
                         "file": file,