async def _create_worker_task_definition_arn(self): response = await self._clients["ecs"].register_task_definition( family="{}-{}".format(self.cluster_name, "worker"), taskRoleArn=self._task_role_arn, executionRoleArn=self._execution_role_arn, networkMode="awsvpc", containerDefinitions=[{ "name": "dask-worker", "image": self.image, "cpu": self._worker_cpu, "memory": self._worker_mem, "memoryReservation": self._worker_mem, "essential": True, "environment": dict_to_aws(self.environment, key_string="name"), "command": [ "dask-worker", "--nthreads", "{}".format(int(self._worker_cpu / 1024)), "--memory-limit", "{}GB".format(int(self._worker_mem / 1024)), "--death-timeout", "60", ], "logConfiguration": { "logDriver": "awslogs", "options": { "awslogs-region": self._clients["ecs"].meta.region_name, "awslogs-group": self.cloudwatch_logs_group, "awslogs-stream-prefix": self._cloudwatch_logs_stream_prefix, "awslogs-create-group": "true", }, }, }], volumes=[], requiresCompatibilities=["FARGATE"] if self._fargate else [], cpu=str(self._worker_cpu), memory=str(self._worker_mem), tags=dict_to_aws(self.tags), ) weakref.finalize(self, self.sync, self._delete_worker_task_definition_arn) return response["taskDefinition"]["taskDefinitionArn"]
async def _create_task_role(self): response = await self._clients["iam"].create_role( RoleName=self._task_role_name, AssumeRolePolicyDocument="""{ "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": "ecs-tasks.amazonaws.com" }, "Action": "sts:AssumeRole" } ] }""", Description="A role for dask tasks to use when executing", Tags=dict_to_aws(self.tags, upper=True), ) for policy in self._task_role_policies: await self._clients["iam"].attach_role_policy( RoleName=self._task_role_name, PolicyArn=policy) weakref.finalize(self, self.sync, self._delete_role, self._task_role_name) return response["Role"]["Arn"]
async def _create_execution_role(self): response = await self._clients["iam"].create_role( RoleName=self._execution_role_name, AssumeRolePolicyDocument="""{ "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Principal": { "Service": "ecs-tasks.amazonaws.com" }, "Action": "sts:AssumeRole" } ] }""", Description="A role for ECS to use when executing", Tags=dict_to_aws(self.tags, upper=True), ) await self._clients["iam"].attach_role_policy( RoleName=self._execution_role_name, PolicyArn= "arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly", ) await self._clients["iam"].attach_role_policy( RoleName=self._execution_role_name, PolicyArn="arn:aws:iam::aws:policy/CloudWatchLogsFullAccess", ) await self._clients["iam"].attach_role_policy( RoleName=self._execution_role_name, PolicyArn= "arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceRole", ) weakref.finalize(self, self.sync, self._delete_role, self._execution_role_name) return response["Role"]["Arn"]
async def start(self): timeout = Timeout( 60, "Unable to start %s after 60 seconds" % self.task_type) while timeout.run(): try: kwargs = ( { "tags": dict_to_aws(self.tags) } if await self._is_long_arn_format_enabled() else {} ) # Tags are only supported if you opt into long arn format so we need to check for that [self.task] = (await self._clients["ecs"].run_task( cluster=self.cluster_arn, taskDefinition=self.task_definition_arn, overrides=self.overrides, count=1, launchType="FARGATE" if self.fargate else "EC2", networkConfiguration={ "awsvpcConfiguration": { "subnets": self._vpc_subnets, "securityGroups": self._security_groups, "assignPublicIp": "ENABLED" if self._use_public_ip else "DISABLED", } }, **kwargs))["tasks"] break except Exception as e: timeout.set_exception(e) await asyncio.sleep(1) self.task_arn = self.task["taskArn"] while self.task["lastStatus"] in ["PENDING", "PROVISIONING"]: await asyncio.sleep(1) await self._update_task() if not await self._task_is_running(): raise RuntimeError("%s failed to start" % type(self).__name__) [eni] = [ attachment for attachment in self.task["attachments"] if attachment["type"] == "ElasticNetworkInterface" ] [network_interface_id] = [ detail["value"] for detail in eni["details"] if detail["name"] == "networkInterfaceId" ] eni = await self._clients["ec2"].describe_network_interfaces( NetworkInterfaceIds=[network_interface_id]) [interface] = eni["NetworkInterfaces"] self.public_ip = interface["Association"]["PublicIp"] self.private_ip = interface["PrivateIpAddresses"][0][ "PrivateIpAddress"] self.address = await self._get_address_from_logs() self.status = "running"
async def _create_cluster(self): if not self._fargate: raise RuntimeError( "You must specify a cluster when not using Fargate.") self.cluster_name = dask.config.expand_environment_variables( self._cluster_name_template) self.cluster_name = self.cluster_name.format( uuid=str(uuid.uuid4())[:10]) response = await self._clients["ecs"].create_cluster( clusterName=self.cluster_name, tags=dict_to_aws(self.tags)) weakref.finalize(self, self.sync, self._delete_cluster) return response["cluster"]["clusterArn"]
async def _create_security_groups(self): response = await self._clients["ec2"].create_security_group( Description="A security group for dask-ecs", GroupName=self.cluster_name, VpcId=self._vpc, DryRun=False, ) await self._clients["ec2"].authorize_security_group_ingress( GroupId=response["GroupId"], IpPermissions=[ { "IpProtocol": "TCP", "FromPort": 8786, "ToPort": 8787, "IpRanges": [{ "CidrIp": "0.0.0.0/0", "Description": "Anywhere" }], "Ipv6Ranges": [{ "CidrIpv6": "::/0", "Description": "Anywhere" }], }, { "IpProtocol": "TCP", "FromPort": 0, "ToPort": 65535, "UserIdGroupPairs": [{ "GroupName": self.cluster_name }], }, ], DryRun=False, ) await self._clients["ec2"].create_tags(Resources=[response["GroupId"]], Tags=dict_to_aws(self.tags, upper=True)) weakref.finalize(self, self.sync, self._delete_security_groups) return [response["GroupId"]]
def test_aws_to_dict_and_back(): from dask_cloud.providers.aws.helper import aws_to_dict, dict_to_aws aws_dict = [{"key": "hello", "value": "world"}] aws_upper_dict = [{"Key": "hello", "Value": "world"}] py_dict = {"hello": "world"} assert dict_to_aws(py_dict) == aws_dict assert dict_to_aws(py_dict, upper=True) == aws_upper_dict assert aws_to_dict(aws_dict) == py_dict assert aws_to_dict(dict_to_aws(py_dict, upper=True)) == py_dict assert aws_to_dict(dict_to_aws(py_dict)) == py_dict assert dict_to_aws(aws_to_dict(aws_dict)) == aws_dict assert dict_to_aws(aws_to_dict(aws_upper_dict), upper=True) == aws_upper_dict