ARM64 support for ECS and Fargate

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

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)