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