Merge pull request #85 from joeschmid/fargate-private-ip

Support avoiding use of a public IP with Fargate

dask_cloudprovider/providers/aws/ecs.py

@@ -73,6 +73,13 @@ class Task:
         Configurable timeout in seconds for finding the task IP from the
         cloudwatch logs.
 
+    name: str (optional)
+        Name for the task. Currently used for the --namecommand line argument to dask-worker.
+
+    fargate_use_private_ip: bool (optional)
+        Whether to use a private IP (if True) or public IP (if False) with Fargate.
+        Defaults to False, i.e. public IP.
+
     kwargs:
         Any additional kwargs which may need to be stored for later use.
 
@@ -96,6 +103,7 @@ def __init__(
         tags,
         find_address_timeout,
         name=None,
+        fargate_use_private_ip=False,
         **kwargs
     ):
         self.lock = asyncio.Lock()
@@ -120,6 +128,7 @@ def __init__(
         self.environment = environment or {}
         self.tags = tags
         self._find_address_timeout = find_address_timeout
+        self._fargate_use_private_ip = fargate_use_private_ip
         self.kwargs = kwargs
         self.status = "created"
 
@@ -135,14 +144,7 @@ async def _():
 
     @property
     def _use_public_ip(self):
-        # Fargate needs public IP for image pull, EC2 doesn't support public IP, therefore
-        # we will assume for now that we will use a public IP when in Fargate mode and not
-        # when in EC2 mode.
-
-        # TODO Fargate can also use a NAT to pull the image so we could allow this to be false
-        # when in Fargate provided there is a NAT
-
-        return self.fargate
+        return self.fargate and not self._fargate_use_private_ip
 
     async def _is_long_arn_format_enabled(self):
         async with self._client("ecs") as ecs:
@@ -524,6 +526,10 @@ class ECSCluster(SpecCluster):
         Skip cleaning up of stale resources. Useful if you have lots of resources
         and this operation takes a while.
 
+        Default ``False``.
+    fargate_use_private_ip: bool (optional)
+        Whether to use a private IP (if True) or public IP (if False) with Fargate.
+        
         Default ``False``.
     **kwargs: dict
         Additional keyword arguments to pass to ``SpecCluster``.
@@ -563,6 +569,7 @@ def __init__(
         aws_access_key_id=None,
         aws_secret_access_key=None,
         region_name=None,
+        fargate_use_private_ip=False,
         **kwargs
     ):
         self._fargate_scheduler = fargate_scheduler
@@ -591,6 +598,7 @@ def __init__(
         self._tags = tags
         self._find_address_timeout = find_address_timeout
         self._skip_cleanup = skip_cleanup
+        self._fargate_use_private_ip = fargate_use_private_ip
         self._aws_access_key_id = aws_access_key_id
         self._aws_secret_access_key = aws_secret_access_key
         self._region_name = region_name
@@ -744,6 +752,7 @@ async def _start(self,):
             "environment": self._environment,
             "tags": self.tags,
             "find_address_timeout": self._find_address_timeout,
+            "fargate_use_private_ip": self._fargate_use_private_ip,
         }
         scheduler_options = {
             "task_definition_arn": self.scheduler_task_definition_arn,