from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from operators.bedrock_agent import BedrockInvokeOperator
from airflow.models import BaseOperator
from util.base import BaseUtil
from util.data_lake import DataLakeUtil
from util import Util
[docs]
class MlflowModelInvokeOperator(BaseOperator):
ui_color = "#A6E6A6"
FIELD_MAPPINGS = {
"BYTES": "string",
"FP32": "double",
"FP64": "double",
"INT32": "int",
"INT64": "bigint",
"INT8": "int",
}
@apply_defaults
def __init__(
self,
model_id,
output_s3_prefix,
output_athena_dataset_name,
input_sql,
datalake_db,
output_fields: dict = None,
model_params={},
use_ecs=False,
fallback_to_ecs_on_timeout=False,
mode="overwrite",
kserve_timeout=3600,
ecs_size="",
*args,
**kwargs,
):
self.map_name = model_id
self.model_id = model_id
self.input_sql = input_sql
self.output_athena_dataset_name = output_athena_dataset_name
self.output_s3_prefix = output_s3_prefix
self.datalake_db = datalake_db
self.output_fields = output_fields
self.mode = mode
self.model_params = model_params
self.use_ecs = use_ecs
self.kserve_timeout = kserve_timeout
self.ecs_size = ecs_size
self.fallback_to_ecs_on_timeout = fallback_to_ecs_on_timeout
print("MLFLOW MODEL ID: ", self.model_id)
print("MLFLOW MODEL OUTPUT S3 PREFIX: ", self.output_s3_prefix)
print("MLFLOW MODEL OUTPUT ATHENA DATASET NAME: ", self.output_athena_dataset_name)
print("MLFLOW MODEL INPUT SQL: ", self.input_sql)
print("MLFLOW MODEL DATALAKE DB: ", self.datalake_db)
print("MLFLOW MODEL OUTPUT FIELDS: ", self.output_fields)
print("MLFLOW MODEL PARAMS: ", self.model_params)
print("MLFLOW MODEL USE ECS: ", self.use_ecs)
print("MLFLOW MODEL MODE: ", self.mode)
print("MLFLOW MODEL KServe TIMEOUT: ", self.kserve_timeout)
print("MLFLOW MODEL ECS SIZE: ", self.ecs_size)
super().__init__(*args, **kwargs)
[docs]
def get_model_base_url(self):
model_id = self.model_id.split("/")[-1]
return f"http://{model_id}.default.models.ready-internal.net"
[docs]
def get_model_invocation_url(self):
return f"{self.get_model_base_url()}/invocations"
[docs]
def get_model_schema(self):
model_metadata = self.get_model_metadata()
input_schema = {
field["name"]: self.FIELD_MAPPINGS.get(field["datatype"], field["datatype"])
for field in model_metadata.get("inputs", [])
}
output_schema = {
field["name"]: self.FIELD_MAPPINGS.get(field["datatype"], field["datatype"])
for field in model_metadata.get("outputs", [])
}
return input_schema, output_schema
[docs]
def create_output_table(self, all_df_fields):
import awswrangler as wr
if self.mode == "overwrite":
tgt_prefix = (
self.output_s3_prefix.replace(f"s3://{Config.DATA_LAKE_BUCKET}/", "")
if self.output_s3_prefix.startswith("s3://")
else self.output_s3_prefix
)
DataLakeUtil.remove_datalake_table(
tgt_table=self.output_athena_dataset_name,
tgt_prefix=tgt_prefix,
tgt_db=self.datalake_db,
)
table_exists = wr.catalog.does_table_exist(
database=self.datalake_db, table=self.output_athena_dataset_name
)
if not table_exists:
wr.catalog.create_parquet_table(
database=self.datalake_db,
table=self.output_athena_dataset_name,
path=self.output_s3_prefix,
columns_types={
f: self.output_fields.get(f, "string") for f in all_df_fields
},
description="Table created from DataFrame",
)
print(f"Created table '{self.output_athena_dataset_name}'")
[docs]
def run_ecs_model(self, input):
input_dict = input.to_dict()
print("RUNNING ECS MODEL")
print("INPUT: ", input_dict)
print("MODEL ID: ", self.model_id)
Util.run_mlflow_model(
model_uri=self.model_id,
input=input.to_dict(),
s3_root=self.output_s3_prefix,
output_athena_dataset_name=self.output_athena_dataset_name,
wr_database=self.datalake_db,
params=self.model_params,
task_suffix=self.ecs_size
)
[docs]
def run_using_kserve(self, input, url, all_df_fields,does_table_exist):
import awswrangler as wr
import pandas as pd
id = input["id"]
# try:
data = {
"dataframe_records": [input.to_dict()],
"params": self.model_params,
}
results = BaseUtil.post_request(
url=url,
data=data,
headers={"Content-Type": "application/json"},
params=self.model_params,
timeout=self.kserve_timeout,
)
model_results = results.json()
output_df = pd.DataFrame(model_results.get("predictions", []))
if not output_df.empty:
output_df["id"] = id
output_df["api_call_status"] = "success"
output_df = output_df[all_df_fields]
if does_table_exist:
wr.s3.to_parquet(
df=output_df,
path=self.output_s3_prefix,
database=self.datalake_db
)
else:
wr.s3.to_parquet(
df=output_df,
path=self.output_s3_prefix,
database=self.datalake_db,
table=self.output_athena_dataset_name,
dataset=True,
mode="append",
)
# except Exception as ex:
# # we want to know which geoids are failing
# print(f"Error processing row {step}- {id}: {ex}")
# output_df = pd.DataFrame(columns=all_df_fields)
# new_row = {"api_call_status": "failed", "id": id}
# output_df.loc[len(output_df)] = new_row
# finally:
# if not output_df.empty:
# output_df = output_df[all_df_fields]
# wr.s3.to_parquet(
# df=output_df,
# path=self.output_s3_prefix,
# database=self.datalake_db,
# table=self.output_athena_dataset_name,
# dataset=True,
# mode="append",
# )
[docs]
def execute(self, context):
import awswrangler as wr
import pandas as pd
from requests.exceptions import ReadTimeout
# add map_index_template
context["params"]["map_name"] = self.map_name
context["map_index_template"] = "{{ params.map_name }}"
input_schema, output_schema = self.get_model_schema()
print("INPUT SCHEMA: ", input_schema)
print("OUTPUT SCHEMA: ", output_schema)
all_df_fields = ["id", "api_call_status"] + list(self.output_fields.keys())
print("ALL DF FIELDS: ", self.output_fields)
self.create_output_table(all_df_fields)
url = self.get_model_invocation_url()
print("READING ATHENA INPUT SQL: ", self.input_sql)
input_df = wr.athena.read_sql_query(
sql=self.input_sql,
database=self.datalake_db,
)
if input_df.empty:
raise Exception("Input DataFrame is empty. Please check the input SQL.")
total_rows = len(input_df)
does_table_exist=wr.catalog.does_table_exist(database=self.datalake_db, table=self.output_athena_dataset_name)
print(f"Total rows to process: {total_rows}")
print(f"Does table exist: {does_table_exist}")
for idx, row in input_df.iterrows():
step=f"{idx+1}/{total_rows}"
print(f"Processing row {step}- {row['id']}")
print(row.to_dict())
id = row["id"]
try:
if self.use_ecs:
self.run_ecs_model(row)
else:
try:
self.run_using_kserve(row, url, all_df_fields,does_table_exist)
except Exception as e:
print(f"Read timeout processing row {step}- {id}: {e}")
if self.fallback_to_ecs_on_timeout:
print(f"Fallback to ECS processing row {step}- {id}")
self.run_ecs_model(row)
except Exception as ex:
print(f"Error processing row {step}- {id}: {ex}")
output_df = pd.DataFrame(columns=all_df_fields)
new_row = {"api_call_status": "failed", "id": id}
output_df.loc[len(output_df)] = new_row
wr.s3.to_parquet(
df=output_df,
path=self.output_s3_prefix,
database=self.datalake_db,
table=self.output_athena_dataset_name,
dataset=True,
mode="append",
)
finally:
does_table_exist = True
wr.athena.repair_table(
database=self.datalake_db,
table=self.output_athena_dataset_name,
)