Skip to content

Commit

Permalink
ARM64 support for ECS and Fargate (#425)
Browse files Browse the repository at this point in the history
* Support explicit runtime platform architecture selection for ECS and Fargate clusters

* Include CPU architecture into yaml config

* Capitalise CPU architecture, as AWS requires upper case, but some libraries output lower vase
  • Loading branch information
dmitry-livchak committed Mar 15, 2024
1 parent 1eb17ff commit 35e4d1d
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 0 deletions.
8 changes: 8 additions & 0 deletions dask_cloudprovider/aws/ecs.py
Expand Up @@ -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).
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down
1 change: 1 addition & 0 deletions dask_cloudprovider/cloudprovider.yaml
Expand Up @@ -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)
Expand Down

0 comments on commit 35e4d1d

Please sign in to comment.