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")
