def test_extra_container_config_merge(image_name, loop):
"""
Test that our container config merging process works recursively fine
"""
with KubeCluster(
make_pod_spec(
image_name,
extra_container_config={
"env": [{
"name": "BOO",
"value": "FOO"
}],
"args": ["last-item"],
},
),
loop=loop,
n_workers=0,
env={"TEST": "HI"},
) as cluster:
pod = cluster.pod_template
assert pod.spec.containers[0].env == [
{
"name": "TEST",
"value": "HI"
},
{
"name": "BOO",
"value": "FOO"
},
]
assert pod.spec.containers[0].args[-1] == "last-item"
def test_extra_pod_config(docker_image, loop):
"""
Test that our pod config merging process works fine
"""
with KubeCluster(
make_pod_spec(
docker_image, extra_pod_config={"automountServiceAccountToken": False}
),
loop=loop,
n_workers=0,
) as cluster:
pod = cluster.pod_template
assert pod.spec.automount_service_account_token is False
def test_container_resources_config(docker_image, loop):
"""
Test container resource requests / limits being set properly
"""
with KubeCluster(
make_pod_spec(
docker_image, memory_request="0.5G", memory_limit="1G", cpu_limit="1"
),
loop=loop,
n_workers=0,
) as cluster:
pod = cluster.pod_template
assert pod.spec.containers[0].resources.requests["memory"] == "0.5G"
assert pod.spec.containers[0].resources.limits["memory"] == "1G"
assert pod.spec.containers[0].resources.limits["cpu"] == "1"
assert "cpu" not in pod.spec.containers[0].resources.requests
def test_extra_pod_config(image_name, loop):
"""
Test that our pod config merging process works fine
"""
cluster = KubeCluster(
make_pod_spec(
image_name,
extra_pod_config={
'automountServiceAccountToken': False
}
),
loop=loop,
n_workers=0,
)
pod = cluster.pod_template
assert pod.spec.automount_service_account_token is False
def test_extra_container_config(docker_image, loop):
"""
Test that our container config merging process works fine
"""
with KubeCluster(
make_pod_spec(
docker_image,
extra_container_config={
"imagePullPolicy": "IfNotPresent",
"securityContext": {"runAsUser": 0},
},
),
loop=loop,
n_workers=0,
) as cluster:
pod = cluster.pod_template
assert pod.spec.containers[0].image_pull_policy == "IfNotPresent"
assert pod.spec.containers[0].security_context == {"runAsUser": 0}
def test_container_resources_config(image_name, loop):
"""
Test container resource requests / limits being set properly
"""
cluster = KubeCluster(
make_pod_spec(
image_name,
memory_request="1G",
memory_limit="2G",
cpu_limit="2",
),
loop=loop,
n_workers=0,
)
pod = cluster.pod_template
assert pod.spec.containers[0].resources.requests['memory'] == '1G'
assert pod.spec.containers[0].resources.limits['memory'] == '2G'
assert pod.spec.containers[0].resources.limits['cpu'] == '2'
assert "cpu" not in pod.spec.containers[0].resources.requests
def test_extra_container_config(image_name, loop):
"""
Test that our container config merging process works fine
"""
cluster = KubeCluster(
make_pod_spec(
image_name,
extra_container_config={
'imagePullPolicy': 'IfNotReady',
'securityContext': {
'runAsUser': 0
}
}
),
loop=loop,
n_workers=0,
)
pod = cluster.pod_template
assert pod.spec.containers[0].image_pull_policy == 'IfNotReady'
assert pod.spec.containers[0].security_context == {
'runAsUser': 0
}
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 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
