diff --git a/dask_cloudprovider/aws/ecs.py b/dask_cloudprovider/aws/ecs.py index 02db1d55..1999a601 100644 --- a/dask_cloudprovider/aws/ecs.py +++ b/dask_cloudprovider/aws/ecs.py @@ -484,6 +484,10 @@ class ECSCluster(SpecCluster, ConfigMixin): The docker image to use for the scheduler and worker tasks. Defaults to ``daskdev/dask:latest`` or ``rapidsai/rapidsai:latest`` if ``worker_gpu`` is set. + cpu_architecture: str (optional) + Runtime platform CPU architecture + + Defaults to ``X86_64``. scheduler_cpu: int (optional) The amount of CPU to request for the scheduler in milli-cpu (1/1024). @@ -712,6 +716,7 @@ def __init__( fargate_workers=None, fargate_spot=None, image=None, + cpu_architecture="X86_64", scheduler_cpu=None, scheduler_mem=None, scheduler_port=8786, @@ -758,6 +763,7 @@ def __init__( self._fargate_workers = fargate_workers self._fargate_spot = fargate_spot self.image = image + self._cpu_architecture = cpu_architecture.upper() self._scheduler_cpu = scheduler_cpu self._scheduler_mem = scheduler_mem self._scheduler_port = scheduler_port @@ -1223,6 +1229,7 @@ async def _create_scheduler_task_definition_arn(self): if self._volumes and self._mount_volumes_on_scheduler else [], requiresCompatibilities=["FARGATE"] if self._fargate_scheduler else [], + runtimePlatform={"cpuArchitecture": self._cpu_architecture}, cpu=str(self._scheduler_cpu), memory=str(self._scheduler_mem), tags=dict_to_aws(self.tags), @@ -1297,6 +1304,7 @@ async def _create_worker_task_definition_arn(self): ], volumes=self._volumes if self._volumes else [], requiresCompatibilities=["FARGATE"] if self._fargate_workers else [], + runtimePlatform={"cpuArchitecture": self._cpu_architecture}, cpu=str(self._worker_cpu), memory=str(self._worker_mem), tags=dict_to_aws(self.tags), diff --git a/dask_cloudprovider/cloudprovider.yaml b/dask_cloudprovider/cloudprovider.yaml index 1a17ff45..aa99b102 100755 --- a/dask_cloudprovider/cloudprovider.yaml +++ b/dask_cloudprovider/cloudprovider.yaml @@ -15,6 +15,7 @@ cloudprovider: scheduler_timeout: "5 minutes" # Length of inactivity to wait before closing the cluster image: "daskdev/dask:latest" # Docker image to use for non GPU tasks + cpu_architecture: "X86_64" # Runtime platform CPU architecture gpu_image: "rapidsai/rapidsai:latest" # Docker image to use for GPU tasks cluster_name_template: "dask-{uuid}" # Template to use when creating a cluster cluster_arn: "" # ARN of existing ECS cluster to use (if not set one will be created)