import json import logging import time import airflow import boto3 import redis import requests from airflow.exceptions import AirflowSkipException from airflow.providers.amazon.aws.hooks.batch_client import AwsBatchClientHook from airflow.providers.amazon.aws.hooks.secrets_manager import SecretsManagerHook from fansifter.operators.base import FanSifterBaseOperator, FanSifterLambdaOperator ML_ENGINE_URL = "https://ml_engine.fansifter.cloud/train_models" ML_ENRICHMENT_ECS_VS_BATCH_THRESHOLD = 10000 AWS_BATCH_FARGATE_JOB_QUEUE_ARN = "arn:aws:batch:eu-west-1:776891437216:job-queue/ml-enrichment-queue" AWS_BATCH_FARGATE_JOB_DEFINITION_ARN = "arn:aws:batch:eu-west-1:776891437216:job-definition/batch-ml-enrichment-job:6" class MLECSOperator(FanSifterBaseOperator): def __init__( self, enrichment_name, input_task_id, input_xcom_key_suffix="preprocessing", output_xcom_key_suffix="processing", **kwargs, ) -> None: super().__init__(**kwargs) self.enrichment_name = enrichment_name self.input_task_id = input_task_id self.input_xcom_key_suffix = input_xcom_key_suffix self.output_xcom_key_suffix = output_xcom_key_suffix def execute(self, context): ti: airflow.models.TaskInstance = context["task_instance"] [preprocessing_results] = ti.xcom_pull( key=f"{self.enrichment_name}-{self.input_xcom_key_suffix}", task_ids=[self.input_task_id] ) dataset_size = preprocessing_results["dataset_size"] schema = preprocessing_results["schema"] if dataset_size < 5: logging.info( f"At least 4 fans are required to perform PCA (present: {dataset_size}, terminating the task..." ) raise AirflowSkipException elif dataset_size >= ML_ENRICHMENT_ECS_VS_BATCH_THRESHOLD: logging.info("Dataset is too big for ECS ML Engine, terminating the task...") raise AirflowSkipException response = requests.post(ML_ENGINE_URL, json=preprocessing_results, timeout=30 * 60) if response.status_code != 200: raise RuntimeError(f"{schema} ML Engine response: {response.text}") response = response.json() logging.info(response) mlflow_run_id = response.get("run_id") if not mlflow_run_id: raise RuntimeError( f"{schema} - ML Engine response does not contain Mlflow run_id, but contains: {response}" ) result = dict(mlflow_run_id=mlflow_run_id) ti.xcom_push(key=f"{self.enrichment_name}-{self.output_xcom_key_suffix}", value=result) class MLBatchOperator(FanSifterBaseOperator): def __init__( self, enrichment_name, input_task_id, input_xcom_key_suffix="preprocessing", output_xcom_key_suffix="processing", **kwargs, ) -> None: super().__init__(**kwargs) self.enrichment_name = enrichment_name self.input_task_id = input_task_id self.input_xcom_key_suffix = input_xcom_key_suffix self.output_xcom_key_suffix = output_xcom_key_suffix def execute(self, context): ti: airflow.models.TaskInstance = context["task_instance"] [preprocessing_results] = ti.xcom_pull( key=f"{self.enrichment_name}-{self.input_xcom_key_suffix}", task_ids=[self.input_task_id] ) dataset_size = preprocessing_results["dataset_size"] bucket = preprocessing_results["bucket"] file_key = preprocessing_results["file_key"] if dataset_size < ML_ENRICHMENT_ECS_VS_BATCH_THRESHOLD: logging.info("Dataset is too small for AWS Batch, terminating the task...") raise AirflowSkipException conf = self._get_config() schema = conf.get("workspace_schema") or conf.get("alliance_schema") collection_id = conf["collection_id"] aws_batch_response = boto3.client("batch").submit_job( jobName=f'ml-enrichment-{schema}-{time.strftime("%Y%m%d-%H%M%S")}', jobQueue=AWS_BATCH_FARGATE_JOB_QUEUE_ARN, jobDefinition=AWS_BATCH_FARGATE_JOB_DEFINITION_ARN, parameters=dict( # these are injected into a CLI call to ML Engine inside its container - see job definition collection_ids=str(collection_id), bucket=bucket, file_key=file_key, schema=schema, ), containerOverrides=dict( resourceRequirements=[dict(type="VCPU", value="4"), dict(type="MEMORY", value="30720"),] ), arrayProperties={}, tags=dict(schema=schema, dataset_length=str(dataset_size)), ) aws_batch_job_id = aws_batch_response["jobId"] logging.info(aws_batch_response) if "jobId" not in aws_batch_response: raise RuntimeError("No jobId present in AWS Batch Response") logging.info(f"Submitted AWS Batch Job (id: {aws_batch_job_id})") batch_hook = AwsBatchClientHook() batch_hook.wait_for_job(aws_batch_job_id, delay=10) # 10 seconds between status checks batch_hook.check_job_success(aws_batch_job_id) redis_config = SecretsManagerHook().get_secret_as_dict("fansifter-redis") redis_client = redis.Redis(**redis_config, db=0) ml_response = json.loads(redis_client.get(f"batch-ml-result:{aws_batch_job_id}")) mlflow_run_id = ml_response.get("run_id") if mlflow_run_id is None: raise RuntimeError( f"{schema} - ML Engine response does not contain Mlflow run_id, but contains: {ml_response}" ) redis_client.close() result = dict(mlflow_run_id=mlflow_run_id) ti.xcom_push(key=f"{self.enrichment_name}-{self.output_xcom_key_suffix}", value=result) class MLEnrichmentPostprocessingOperator(FanSifterLambdaOperator): def __init__( self, enrichment_name, input_preprocessing_task_id, input_ecs_task_id, input_batch_task_id, input_xcom_preprocessing_key_suffix="preprocessing", input_xcom_processing_key_suffix="processing", output_xcom_key_suffix="enrichment", **kwargs, ) -> None: super().__init__(**kwargs) self.enrichment_name = enrichment_name self.input_preprocessing_task_id = input_preprocessing_task_id self.input_ecs_task_id = input_ecs_task_id self.input_batch_task_id = input_batch_task_id self.input_xcom_preprocessing_key_suffix = input_xcom_preprocessing_key_suffix self.input_xcom_processing_key_suffix = input_xcom_processing_key_suffix self.output_xcom_key_suffix = output_xcom_key_suffix def execute(self, context): ti: airflow.models.TaskInstance = context["task_instance"] conf = self._get_config() profile = conf["profile"] workspace_schema = conf.get("workspace_schema") alliance_schema = conf.get("alliance_schema") user_id = conf["user_id"] source_collection_id = conf["collection_id"] [preprocessing_results] = ti.xcom_pull( key=f"{self.enrichment_name}-{self.input_xcom_preprocessing_key_suffix}", task_ids=[self.input_preprocessing_task_id], ) grouped_df_csv_path = preprocessing_results["grouped_df_csv_path"] enrichment_collection_id = preprocessing_results["enrichment_collection_id"] [processing_results_ecs, processing_results_batch] = ti.xcom_pull( key=f"{self.enrichment_name}-{self.input_xcom_processing_key_suffix}", task_ids=[self.input_ecs_task_id, self.input_batch_task_id], ) processing_results = processing_results_ecs or processing_results_batch if processing_results is None: logging.info("Model training has not returned any results, terminating the task...") raise AirflowSkipException mlflow_run_id = processing_results["mlflow_run_id"] payload = { "endpoint": f"enrich/{self.enrichment_name}", "arguments": { "user_id": user_id, "collection_id": source_collection_id, "enrichment_config": { "phase": "postprocessing", "grouped_df_csv_path": grouped_df_csv_path, "enrichment_collection_id": enrichment_collection_id, "mlflow_run_id": mlflow_run_id, }, }, } if workspace_schema is not None: payload["arguments"].update({"workspace_schema": workspace_schema}) elif alliance_schema is not None: payload["arguments"].update({"alliance_schema": alliance_schema}) try: lambda_result = self.invoke_lambda( function_name=f"frontend-api-{profile}-enrichment-requests", payload=payload ) if lambda_result is None: logging.info("Postprocessing has not created any data, terminating the task...") return None except Exception as e: logging.error(e) raise e ti.xcom_push(key=f"{self.enrichment_name}-{self.output_xcom_key_suffix}", value=lambda_result)