def register_recipe(recipe: BaseRecipe):
fs_remote = AzureBlobFileSystem(
connection_string=os.environ["FLOW_STORAGE_CONNECTION_STRING"]
)
target = FSSpecTarget(
fs_remote,
root_path=f"abfs://{os.environ['FLOW_CACHE_CONTAINER']}/azurerecipetest/",
)
recipe.target = target
recipe.input_cache = CacheFSSpecTarget(
fs_remote,
root_path=(
f"abfs://{os.environ['FLOW_CACHE_CONTAINER']}/azurerecipetestcache/"
),
)
recipe.metadata_cache = target
executor = PrefectPipelineExecutor()
pipeline = recipe.to_pipelines()
flow = executor.pipelines_to_plan(pipeline)
job_template = yaml.safe_load(
"""
apiVersion: batch/v1
kind: Job
metadata:
annotations:
"cluster-autoscaler.kubernetes.io/safe-to-evict": "false"
spec:
template:
spec:
containers:
- name: flow
"""
)
flow_name = "test-noaa-flow"
flow.storage = storage.Azure(
container=os.environ["FLOW_STORAGE_CONTAINER"],
connection_string=os.environ["FLOW_STORAGE_CONNECTION_STRING"],
)
flow.run_config = KubernetesRun(
job_template=job_template,
image=os.environ["BAKERY_IMAGE"],
env={
"AZURE_STORAGE_CONNECTION_STRING": os.environ[
"FLOW_STORAGE_CONNECTION_STRING"
],
},
labels=json.loads(os.environ["PREFECT__CLOUD__AGENT__LABELS"]),
cpu_request="1000m",
memory_request="3Gi",
)
worker_spec = make_pod_spec(
image=os.environ["BAKERY_IMAGE"],
labels={"flow": flow_name},
memory_limit="1Gi",
memory_request="500Mi",
cpu_limit="512m",
cpu_request="256m",
env={
"AZURE_STORAGE_CONNECTION_STRING": os.environ[
"FLOW_STORAGE_CONNECTION_STRING"
]
},
)
scheduler_spec = make_pod_spec(
image=os.environ["BAKERY_IMAGE"],
labels={"flow": flow_name},
memory_request="500Mi",
cpu_request="256m",
)
scheduler_spec.spec.containers[0].args = ["dask-scheduler"]
scheduler_spec = clean_pod_template(scheduler_spec, pod_type="scheduler")
flow.executor = DaskExecutor(
cluster_class="dask_kubernetes.KubeCluster",
cluster_kwargs={
"pod_template": worker_spec,
"scheduler_pod_template": scheduler_spec,
},
adapt_kwargs={"maximum": 10},
)
for flow_task in flow.tasks:
flow_task.run = set_log_level(flow_task.run)
flow.name = flow_name
project_name = os.environ["PREFECT_PROJECT"]
flow_id = flow.register(project_name=project_name)
return flow_id
def clean_pod_spec(pod_spec):
yield clean_pod_template(pod_spec)
def configure_dask_executor(cluster: Cluster, recipe_bakery: RecipeBakery,
recipe_name: str, secrets: Dict):
if cluster.type == FARGATE_CLUSTER:
worker_cpu = recipe_bakery.resources.cpu if recipe_bakery.resources is not None else 1024
worker_mem = recipe_bakery.resources.memory if recipe_bakery.resources is not None else 4096
dask_executor = DaskExecutor(
cluster_class="dask_cloudprovider.aws.FargateCluster",
cluster_kwargs={
"image": cluster.worker_image,
"vpc": cluster.cluster_options.vpc,
"cluster_arn": cluster.cluster_options.cluster_arn,
"task_role_arn": cluster.cluster_options.task_role_arn,
"execution_role_arn":
cluster.cluster_options.execution_role_arn,
"security_groups": cluster.cluster_options.security_groups,
"scheduler_cpu": 2048,
"scheduler_mem": 16384,
"worker_cpu": worker_cpu,
"worker_mem": worker_mem,
"scheduler_timeout": "15 minutes",
"environment": {
"PREFECT__LOGGING__EXTRA_LOGGERS":
"['pangeo_forge_recipes']",
"MALLOC_TRIM_THRESHOLD_": "0",
},
"tags": {
"Project": "pangeo-forge",
"Recipe": recipe_name,
},
},
adapt_kwargs={
"minimum": 5,
"maximum": cluster.max_workers
},
)
return dask_executor
elif cluster.type == AKS_CLUSTER:
worker_cpu = recipe_bakery.resources.cpu if recipe_bakery.resources is not None else 250
worker_mem = recipe_bakery.resources.memory if recipe_bakery.resources is not None else 512
scheduler_spec = make_pod_spec(
image=cluster.worker_image,
labels={
"Recipe": recipe_name,
"Project": "pangeo-forge"
},
memory_request="10000Mi",
cpu_request="2048m",
)
scheduler_spec.spec.containers[0].args = ["dask-scheduler"]
scheduler_spec = clean_pod_template(scheduler_spec,
pod_type="scheduler")
dask_executor = DaskExecutor(
cluster_class="dask_kubernetes.KubeCluster",
cluster_kwargs={
"pod_template":
make_pod_spec(
image=cluster.worker_image,
labels={
"Recipe": recipe_name,
"Project": "pangeo-forge"
},
memory_request=f"{worker_mem}Mi",
cpu_request=f"{worker_cpu}m",
env={
"AZURE_STORAGE_CONNECTION_STRING":
secrets[cluster.flow_storage_options.secret]
},
),
"scheduler_pod_template":
scheduler_spec,
},
adapt_kwargs={
"minimum": 5,
"maximum": cluster.max_workers
},
)
return dask_executor
else:
raise UnsupportedClusterType