def register_workflow( self, environment: str, prefect_execution_environment: str, prefect_workflow_register_token: str, workflow_cpu_configuration: int, workflow_memory_configuration: int, ) -> None: """ Registers the workflow to Prefect Cloud Parameters: environment [string] -- environment the workflow should be pushed to prefect_execution_environment [string] -- e.g: ecs_fargate, kubernetes prefect_workflow_register_token [string] -- Prefect API token to register workflows workflow_cpu_configuration [int] -- e.g: 256,512,1024,2048,4096 workflow_memory_configuration [int] -- e.g: 512,30720 """ # set flow properties flow_module, flow_name = self.set_workflow_properties( environment, prefect_execution_environment, workflow_cpu_configuration, workflow_memory_configuration, ) # Authenticate to ECR as the registration process pushes the image to AWS self.aws_conn_helpers.ecr_authenticate() # Create ECR repository self.create_workflow_ecr_repository(flow_name=flow_name) try: # Instantiate the prefect client prefect_client = Client(api_token=prefect_workflow_register_token) # Register the Workflow prefect_client.register( flow=flow_module.flow, project_name=f"{environment}_dataflow_automation") except Exception as e: raise e else: logging.info(f"Workflow '{flow_name}' registered successfully!")
def register_workflow(prefect_register_token_secret_name: str): """ Registers the workflow to Prefect Cloud Parameters: prefect_register_token_secret_name [str] -- name of aws secrets manager secret where prefect register token is stored """ flow_module = __import__("flow") flow_name = f"{env}_{flow_module.flow.name}" flow_module.flow.name = flow_name flow_module.flow.environment = FargateTaskEnvironment( requiresCompatibilities=["FARGATE"], region=aws_region, labels=[f"{env}_dataflow_automation"], taskDefinition=flow_name, family=flow_name, cpu="512", memory="3072", networkMode="awsvpc", networkConfiguration={ "awsvpcConfiguration": { "assignPublicIp": "ENABLED", "subnets": subnets, "securityGroups": [], } }, containerDefinitions=[{ "logConfiguration": { "logDriver": "awslogs", "options": { "awslogs-region": aws_region, "awslogs-group": f"{env}_dataflow_automation_workflows", "awslogs-stream-prefix": flow_name, }, } }], executionRoleArn=execution_role_arn, taskRoleArn=task_role_arn, cluster=f"{env}_dataflow_automation_workflows", ) # Set the flow storage. Where to get the code from flow_module.flow.storage = Docker( registry_url=f"{account_id}.dkr.ecr.{aws_region}.amazonaws.com", image_name=flow_name, image_tag="latest", python_dependencies=["boto3"], env_vars={"PYTHONPATH": "/opt/prefect/flows"}, ) # Authenticate to ECR as the registration process pushes the image to AWS ecr_authenticate() # Instantiate the prefect client prefect_client = Client(api_token=get_prefect_token( secret_name=prefect_register_token_secret_name)) # Create ECR repository create_ecr_repository(flow_name=flow_name) # Register the Workflow prefect_client.register(flow=flow_module.flow, project_name=f"{env}_dataflow_automation")