toil 6.1.0a1__py3-none-any.whl → 8.0.0__py3-none-any.whl

toil/test/provisioners/aws/awsProvisionerTest.py

@@ -19,22 +19,37 @@ import time
 from abc import abstractmethod
 from inspect import getsource
 from textwrap import dedent
+from typing import TYPE_CHECKING, Optional
 from uuid import uuid4
 
 import pytest
 
 from toil.provisioners import cluster_factory
 from toil.provisioners.aws.awsProvisioner import AWSProvisioner
-from toil.test import (ToilTest,
-                       integrative,
-                       needs_aws_ec2,
-                       needs_fetchable_appliance,
-                       needs_mesos,
-                       slow,
-                       timeLimit)
+from toil.test import (
+    ToilTest,
+    integrative,
+    needs_aws_ec2,
+    needs_fetchable_appliance,
+    needs_mesos,
+    slow,
+    timeLimit,
+)
 from toil.test.provisioners.clusterTest import AbstractClusterTest
 from toil.version import exactPython
 
+if TYPE_CHECKING:
+    from mypy_boto3_ec2 import EC2Client
+    from mypy_boto3_ec2.type_defs import (
+        DescribeVolumesResultTypeDef,
+        EbsInstanceBlockDeviceTypeDef,
+        FilterTypeDef,
+        InstanceBlockDeviceMappingTypeDef,
+        InstanceTypeDef,
+        VolumeTypeDef,
+    )
+
+
 log = logging.getLogger(__name__)
 
 
@@ -46,11 +61,13 @@ class AWSProvisionerBenchTest(ToilTest):
     # Needs to talk to EC2 for image discovery
     @needs_aws_ec2
     def test_AMI_finding(self):
-        for zone in ['us-west-2a', 'eu-central-1a', 'sa-east-1b']:
-            provisioner = AWSProvisioner('fakename', 'mesos', zone, 10000, None, None)
+        for zone in ["us-west-2a", "eu-central-1a", "sa-east-1b"]:
+            provisioner = AWSProvisioner(
+                "fakename", "mesos", zone, 10000, None, None, enable_fuse=False
+            )
             ami = provisioner._discoverAMI()
             # Make sure we got an AMI and it looks plausible
-            assert(ami.startswith('ami-'))
+            assert ami.startswith("ami-")
 
     @needs_aws_ec2
     def test_read_write_global_files(self):
@@ -58,8 +75,16 @@ class AWSProvisionerBenchTest(ToilTest):
         Make sure the `_write_file_to_cloud()` and `_read_file_from_cloud()`
         functions of the AWS provisioner work as intended.
         """
-        provisioner = AWSProvisioner(f'aws-provisioner-test-{uuid4()}', 'mesos', 'us-west-2a', 50, None, None)
-        key = 'config/test.txt'
+        provisioner = AWSProvisioner(
+            f"aws-provisioner-test-{uuid4()}",
+            "mesos",
+            "us-west-2a",
+            50,
+            None,
+            None,
+            enable_fuse=False,
+        )
+        key = "config/test.txt"
         contents = b"Hello, this is a test."
 
         try:
@@ -80,23 +105,23 @@ class AbstractAWSAutoscaleTest(AbstractClusterTest):
     def __init__(self, methodName):
         super().__init__(methodName=methodName)
         self.instanceTypes = ["m5a.large"]
-        self.clusterName = 'aws-provisioner-test-' + str(uuid4())
-        self.numWorkers = ['2']
+        self.clusterName = "aws-provisioner-test-" + str(uuid4())
+        self.numWorkers = ["2"]
         self.numSamples = 2
         self.spotBid = 0.15
         # We can't dump our user script right in /tmp or /home, because hot
         # deploy refuses to zip up those whole directories. So we make sure to
         # have a subdirectory to upload the script to.
-        self.scriptDir = '/tmp/t'
+        self.scriptDir = "/tmp/t"
         # Where should we put our virtualenv?
-        self.venvDir = '/tmp/venv'
+        self.venvDir = "/tmp/venv"
         # Where should we put our data to work on?
         # Must exist in the Toil container; the leader will try to rsync to it
         # (for the SSE key) and not create it.
-        self.dataDir = '/tmp'
+        self.dataDir = "/tmp"
         # What filename should we use for our script (without path)?
         # Can be changed by derived tests.
-        self.scriptName = 'test_script.py'
+        self.scriptName = "test_script.py"
 
     def script(self):
         """
@@ -111,17 +136,37 @@ class AbstractAWSAutoscaleTest(AbstractClusterTest):
         return os.path.join(self.dataDir, filename)
 
     def rsyncUtil(self, src, dest):
-        subprocess.check_call(['toil', 'rsync-cluster', '--insecure', '-p=aws', '-z', self.zone, self.clusterName] + [src, dest])
-
-    def getRootVolID(self):
-        instances = self.cluster._get_nodes_in_cluster()
-        instances.sort(key=lambda x: x.launch_time)
-        leader = instances[0]  # assume leader was launched first
-
-        from boto.ec2.blockdevicemapping import BlockDeviceType
-        rootBlockDevice = leader.block_device_mapping["/dev/xvda"]
-        assert isinstance(rootBlockDevice, BlockDeviceType)
-        return rootBlockDevice.volume_id
+        subprocess.check_call(
+            [
+                "toil",
+                "rsync-cluster",
+                "--insecure",
+                "-p=aws",
+                "-z",
+                self.zone,
+                self.clusterName,
+            ]
+            + [src, dest]
+        )
+
+    def getRootVolID(self) -> str:
+        instances: list["InstanceTypeDef"] = self.cluster._get_nodes_in_cluster_boto3()
+        instances.sort(key=lambda x: x.get("LaunchTime"))
+        leader: "InstanceTypeDef" = instances[0]  # assume leader was launched first
+
+        bdm: Optional[list["InstanceBlockDeviceMappingTypeDef"]] = leader.get(
+            "BlockDeviceMappings"
+        )
+        assert bdm is not None
+        root_block_device: Optional["EbsInstanceBlockDeviceTypeDef"] = None
+        for device in bdm:
+            if device["DeviceName"] == "/dev/xvda":
+                root_block_device = device["Ebs"]
+        assert (
+            root_block_device is not None
+        )  # There should be a device named "/dev/xvda"
+        assert root_block_device.get("VolumeId") is not None
+        return root_block_device["VolumeId"]
 
     @abstractmethod
     def _getScript(self):
@@ -132,18 +177,19 @@ class AbstractAWSAutoscaleTest(AbstractClusterTest):
         """
         Helper method for _getScript to inject a script file at the configured script path, from text.
         """
-        cluster = cluster_factory(provisioner='aws', zone=self.zone, clusterName=self.clusterName)
+        cluster = cluster_factory(
+            provisioner="aws", zone=self.zone, clusterName=self.clusterName
+        )
         leader = cluster.getLeader()
 
-        self.sshUtil(['mkdir', '-p', self.scriptDir])
+        self.sshUtil(["mkdir", "-p", self.scriptDir])
 
-        with tempfile.NamedTemporaryFile(mode='w') as t:
+        with tempfile.NamedTemporaryFile(mode="w") as t:
             # use appliance ssh method instead of sshutil so we can specify input param
             t.write(content)
             # This works to make writes visible on non-Windows
             t.flush()
-            leader.injectFile(t.name, self.script(), 'toil_leader')
-
+            leader.injectFile(t.name, self.script(), "toil_leader")
 
     @abstractmethod
     def _runScript(self, toilOptions):
@@ -162,51 +208,60 @@ class AbstractAWSAutoscaleTest(AbstractClusterTest):
         self.launchCluster()
         # get the leader so we know the IP address - we don't need to wait since create cluster
         # already insures the leader is running
-        self.cluster = cluster_factory(provisioner='aws', zone=self.zone, clusterName=self.clusterName)
+        self.cluster = cluster_factory(
+            provisioner="aws", zone=self.zone, clusterName=self.clusterName
+        )
         self.leader = self.cluster.getLeader()
-        self.sshUtil(['mkdir', '-p', self.scriptDir])
-        self.sshUtil(['mkdir', '-p', self.dataDir])
+        self.sshUtil(["mkdir", "-p", self.scriptDir])
+        self.sshUtil(["mkdir", "-p", self.dataDir])
 
         assert len(self.cluster._getRoleNames()) == 1
         # --never-download prevents silent upgrades to pip, wheel and setuptools
-        venv_command = ['virtualenv', '--system-site-packages', '--python', exactPython, '--never-download', self.venvDir]
+        venv_command = [
+            "virtualenv",
+            "--system-site-packages",
+            "--python",
+            exactPython,
+            "--never-download",
+            self.venvDir,
+        ]
         self.sshUtil(venv_command)
 
-        log.info('Set up script...')
+        log.info("Set up script...")
         self._getScript()
 
-        toilOptions = [self.jobStore,
-                       '--workDir=/var/lib/toil',
-                       '--clean=always',
-                       '--retryCount=2',
-                       '--logDebug',
-                       '--logFile=' + os.path.join(self.scriptDir, 'sort.log')
-                       ]
+        toilOptions = [
+            self.jobStore,
+            "--workDir=/var/lib/toil",
+            "--clean=always",
+            "--retryCount=2",
+            "--logDebug",
+            "--logFile=" + os.path.join(self.scriptDir, "sort.log"),
+        ]
 
         if preemptibleJobs:
-            toilOptions.extend(['--defaultPreemptible'])
+            toilOptions.extend(["--defaultPreemptible"])
 
-        log.info('Run script...')
+        log.info("Run script...")
         self._runScript(toilOptions)
 
         assert len(self.cluster._getRoleNames()) == 1
 
-        from boto.exception import EC2ResponseError
         volumeID = self.getRootVolID()
         self.cluster.destroyCluster()
+        boto3_ec2: "EC2Client" = self.aws.client(region=self.region, service_name="ec2")
+        volume_filter: "FilterTypeDef" = {"Name": "volume-id", "Values": [volumeID]}
+        volumes: Optional[list["VolumeTypeDef"]] = None
         for attempt in range(6):
             # https://github.com/BD2KGenomics/toil/issues/1567
             # retry this for up to 1 minute until the volume disappears
-            try:
-                self.boto2_ec2.get_all_volumes(volume_ids=[volumeID])
-                time.sleep(10)
-            except EC2ResponseError as e:
-                if e.status == 400 and 'InvalidVolume.NotFound' in e.code:
-                    break
-                else:
-                    raise
-        else:
-            self.fail('Volume with ID %s was not cleaned up properly' % volumeID)
+            volumes = boto3_ec2.describe_volumes(Filters=[volume_filter])["Volumes"]
+            if len(volumes) == 0:
+                # None are left, so they have been properly deleted
+                break
+            time.sleep(10)
+        if volumes is None or len(volumes) > 0:
+            self.fail("Volume with ID %s was not cleaned up properly" % volumeID)
 
         assert len(self.cluster._getRoleNames()) == 0
 
@@ -217,66 +272,86 @@ class AbstractAWSAutoscaleTest(AbstractClusterTest):
 class AWSAutoscaleTest(AbstractAWSAutoscaleTest):
     def __init__(self, name):
         super().__init__(name)
-        self.clusterName = 'provisioner-test-' + str(uuid4())
+        self.clusterName = "provisioner-test-" + str(uuid4())
         self.requestedLeaderStorage = 80
-        self.scriptName = 'sort.py'
+        self.scriptName = "sort.py"
 
     def setUp(self):
         super().setUp()
-        self.jobStore = f'aws:{self.awsRegion()}:autoscale-{uuid4()}'
+        self.jobStore = f"aws:{self.awsRegion()}:autoscale-{uuid4()}"
 
     def _getScript(self):
         fileToSort = os.path.join(os.getcwd(), str(uuid4()))
-        with open(fileToSort, 'w') as f:
+        with open(fileToSort, "w") as f:
             # Fixme: making this file larger causes the test to hang
-            f.write('01234567890123456789012345678901')
-        self.rsyncUtil(os.path.join(self._projectRootPath(), 'src/toil/test/sort/sort.py'), ':' + self.script())
-        self.rsyncUtil(fileToSort, ':' + self.data('sortFile'))
+            f.write("01234567890123456789012345678901")
+        self.rsyncUtil(
+            os.path.join(self._projectRootPath(), "src/toil/test/sort/sort.py"),
+            ":" + self.script(),
+        )
+        self.rsyncUtil(fileToSort, ":" + self.data("sortFile"))
         os.unlink(fileToSort)
 
     def _runScript(self, toilOptions):
-        toilOptions.extend(['--provisioner=aws', '--batchSystem=mesos',
-                            '--nodeTypes=' + ",".join(self.instanceTypes),
-                            '--maxNodes=' + ",".join(self.numWorkers)])
-        runCommand = [self.python(), self.script(), '--fileToSort=' + self.data('sortFile'), '--sseKey=' + self.data('sortFile')]
+        toilOptions.extend(
+            [
+                "--provisioner=aws",
+                "--batchSystem=mesos",
+                "--nodeTypes=" + ",".join(self.instanceTypes),
+                "--maxNodes=" + ",".join(self.numWorkers),
+            ]
+        )
+        runCommand = [
+            self.python(),
+            self.script(),
+            "--fileToSort=" + self.data("sortFile"),
+            "--sseKey=" + self.data("sortFile"),
+        ]
         runCommand.extend(toilOptions)
         self.sshUtil(runCommand)
 
     def launchCluster(self):
         # add arguments to test that we can specify leader storage
-        self.createClusterUtil(args=['--leaderStorage', str(self.requestedLeaderStorage)])
+        self.createClusterUtil(
+            args=["--leaderStorage", str(self.requestedLeaderStorage)]
+        )
 
-    def getRootVolID(self):
+    def getRootVolID(self) -> str:
         """
         Adds in test to check that EBS volume is build with adequate size.
         Otherwise is functionally equivalent to parent.
         :return: volumeID
         """
         volumeID = super().getRootVolID()
-        rootVolume = self.boto2_ec2.get_all_volumes(volume_ids=[volumeID])[0]
+        boto3_ec2: "EC2Client" = self.aws.client(region=self.region, service_name="ec2")
+        volume_filter: "FilterTypeDef" = {"Name": "volume-id", "Values": [volumeID]}
+        volumes: "DescribeVolumesResultTypeDef" = boto3_ec2.describe_volumes(
+            Filters=[volume_filter]
+        )
+        root_volume: "VolumeTypeDef" = volumes["Volumes"][0]  # should be first
         # test that the leader is given adequate storage
-        self.assertGreaterEqual(rootVolume.size, self.requestedLeaderStorage)
+        self.assertGreaterEqual(root_volume["Size"], self.requestedLeaderStorage)
         return volumeID
 
     @integrative
     @needs_aws_ec2
     def testAutoScale(self):
         self.instanceTypes = ["m5a.large"]
-        self.numWorkers = ['2']
+        self.numWorkers = ["2"]
         self._test()
 
     @integrative
     @needs_aws_ec2
     def testSpotAutoScale(self):
         self.instanceTypes = ["m5a.large:%f" % self.spotBid]
-        self.numWorkers = ['2']
+        self.numWorkers = ["2"]
         self._test(preemptibleJobs=True)
 
     @integrative
     @needs_aws_ec2
     def testSpotAutoScaleBalancingTypes(self):
         self.instanceTypes = ["m5.large/m5a.large:%f" % self.spotBid]
-        self.numWorkers = ['2']
+        self.numWorkers = ["2"]
         self._test(preemptibleJobs=True)
 
 
@@ -285,26 +360,36 @@ class AWSAutoscaleTest(AbstractAWSAutoscaleTest):
 @pytest.mark.timeout(2400)
 class AWSStaticAutoscaleTest(AWSAutoscaleTest):
     """Runs the tests on a statically provisioned cluster with autoscaling enabled."""
+
     def __init__(self, name):
         super().__init__(name)
         self.requestedNodeStorage = 20
 
     def launchCluster(self):
-        from boto.ec2.blockdevicemapping import BlockDeviceType
-
         from toil.lib.ec2 import wait_instances_running
-        self.createClusterUtil(args=['--leaderStorage', str(self.requestedLeaderStorage),
-                                     '--nodeTypes', ",".join(self.instanceTypes),
-                                     '-w', ",".join(self.numWorkers),
-                                     '--nodeStorage', str(self.requestedLeaderStorage)])
 
-        self.cluster = cluster_factory(provisioner='aws', zone=self.zone, clusterName=self.clusterName)
+        self.createClusterUtil(
+            args=[
+                "--leaderStorage",
+                str(self.requestedLeaderStorage),
+                "--nodeTypes",
+                ",".join(self.instanceTypes),
+                "-w",
+                ",".join(self.numWorkers),
+                "--nodeStorage",
+                str(self.requestedLeaderStorage),
+            ]
+        )
+
+        self.cluster = cluster_factory(
+            provisioner="aws", zone=self.zone, clusterName=self.clusterName
+        )
         # We need to wait a little bit here because the workers might not be
         # visible to EC2 read requests immediately after the create returns,
         # which is the last thing that starting the cluster does.
         time.sleep(10)
-        nodes = self.cluster._get_nodes_in_cluster()
-        nodes.sort(key=lambda x: x.launch_time)
+        nodes: list["InstanceTypeDef"] = self.cluster._get_nodes_in_cluster_boto3()
+        nodes.sort(key=lambda x: x.get("LaunchTime"))
         # assuming that leader is first
         workers = nodes[1:]
         # test that two worker nodes were created
@@ -312,18 +397,49 @@ class AWSStaticAutoscaleTest(AWSAutoscaleTest):
         # test that workers have expected storage size
         # just use the first worker
         worker = workers[0]
-        worker = next(wait_instances_running(self.boto2_ec2, [worker]))
-        rootBlockDevice = worker.block_device_mapping["/dev/xvda"]
-        self.assertTrue(isinstance(rootBlockDevice, BlockDeviceType))
-        rootVolume = self.boto2_ec2.get_all_volumes(volume_ids=[rootBlockDevice.volume_id])[0]
-        self.assertGreaterEqual(rootVolume.size, self.requestedNodeStorage)
+        boto3_ec2: "EC2Client" = self.aws.client(region=self.region, service_name="ec2")
+
+        worker: "InstanceTypeDef" = next(wait_instances_running(boto3_ec2, [worker]))
+
+        bdm: Optional[list["InstanceBlockDeviceMappingTypeDef"]] = worker.get(
+            "BlockDeviceMappings"
+        )
+        assert bdm is not None
+        root_block_device: Optional["EbsInstanceBlockDeviceTypeDef"] = None
+        for device in bdm:
+            if device["DeviceName"] == "/dev/xvda":
+                root_block_device = device["Ebs"]
+        assert root_block_device is not None
+        assert (
+            root_block_device.get("VolumeId") is not None
+        )  # TypedDicts cannot have runtime type checks
+
+        volume_filter: "FilterTypeDef" = {
+            "Name": "volume-id",
+            "Values": [root_block_device["VolumeId"]],
+        }
+        root_volume: "VolumeTypeDef" = boto3_ec2.describe_volumes(
+            Filters=[volume_filter]
+        )["Volumes"][
+            0
+        ]  # should be first
+        self.assertGreaterEqual(root_volume.get("Size"), self.requestedNodeStorage)
 
     def _runScript(self, toilOptions):
         # Autoscale even though we have static nodes
-        toilOptions.extend(['--provisioner=aws', '--batchSystem=mesos',
-                            '--nodeTypes=' + ",".join(self.instanceTypes),
-                            '--maxNodes=' + ",".join(self.numWorkers)])
-        runCommand = [self.python(), self.script(), '--fileToSort=' + self.data('sortFile')]
+        toilOptions.extend(
+            [
+                "--provisioner=aws",
+                "--batchSystem=mesos",
+                "--nodeTypes=" + ",".join(self.instanceTypes),
+                "--maxNodes=" + ",".join(self.numWorkers),
+            ]
+        )
+        runCommand = [
+            self.python(),
+            self.script(),
+            "--fileToSort=" + self.data("sortFile"),
+        ]
         runCommand.extend(toilOptions)
         self.sshUtil(runCommand)
 
@@ -332,26 +448,39 @@ class AWSStaticAutoscaleTest(AWSAutoscaleTest):
 @pytest.mark.timeout(1200)
 class AWSManagedAutoscaleTest(AWSAutoscaleTest):
     """Runs the tests on a self-scaling Kubernetes cluster."""
+
     def __init__(self, name):
         super().__init__(name)
         self.requestedNodeStorage = 20
 
     def launchCluster(self):
-        from boto.ec2.blockdevicemapping import BlockDeviceType  # noqa
-
-        from toil.lib.ec2 import wait_instances_running  # noqa
-        self.createClusterUtil(args=['--leaderStorage', str(self.requestedLeaderStorage),
-                                     '--nodeTypes', ",".join(self.instanceTypes),
-                                     '--workers', ",".join([f'0-{c}' for c in self.numWorkers]),
-                                     '--nodeStorage', str(self.requestedLeaderStorage),
-                                     '--clusterType', 'kubernetes'])
-
-        self.cluster = cluster_factory(provisioner='aws', zone=self.zone, clusterName=self.clusterName)
+        self.createClusterUtil(
+            args=[
+                "--leaderStorage",
+                str(self.requestedLeaderStorage),
+                "--nodeTypes",
+                ",".join(self.instanceTypes),
+                "--workers",
+                ",".join([f"0-{c}" for c in self.numWorkers]),
+                "--nodeStorage",
+                str(self.requestedLeaderStorage),
+                "--clusterType",
+                "kubernetes",
+            ]
+        )
+
+        self.cluster = cluster_factory(
+            provisioner="aws", zone=self.zone, clusterName=self.clusterName
+        )
 
     def _runScript(self, toilOptions):
         # Don't use the provisioner, and use Kubernetes instead of Mesos
-        toilOptions.extend(['--batchSystem=kubernetes'])
-        runCommand = [self.python(), self.script(), '--fileToSort=' + self.data('sortFile')]
+        toilOptions.extend(["--batchSystem=kubernetes"])
+        runCommand = [
+            self.python(),
+            self.script(),
+            "--fileToSort=" + self.data("sortFile"),
+        ]
         runCommand.extend(toilOptions)
         self.sshUtil(runCommand)
 
@@ -362,37 +491,51 @@ class AWSManagedAutoscaleTest(AWSAutoscaleTest):
 class AWSAutoscaleTestMultipleNodeTypes(AbstractAWSAutoscaleTest):
     def __init__(self, name):
         super().__init__(name)
-        self.clusterName = 'provisioner-test-' + str(uuid4())
+        self.clusterName = "provisioner-test-" + str(uuid4())
 
     def setUp(self):
         super().setUp()
-        self.jobStore = f'aws:{self.awsRegion()}:autoscale-{uuid4()}'
+        self.jobStore = f"aws:{self.awsRegion()}:autoscale-{uuid4()}"
 
     def _getScript(self):
-        sseKeyFile = os.path.join(os.getcwd(), 'keyFile')
-        with open(sseKeyFile, 'w') as f:
-            f.write('01234567890123456789012345678901')
-        self.rsyncUtil(os.path.join(self._projectRootPath(), 'src/toil/test/sort/sort.py'), ':' + self.script())
-        self.rsyncUtil(sseKeyFile, ':' + self.data('keyFile'))
+        sseKeyFile = os.path.join(os.getcwd(), "keyFile")
+        with open(sseKeyFile, "w") as f:
+            f.write("01234567890123456789012345678901")
+        self.rsyncUtil(
+            os.path.join(self._projectRootPath(), "src/toil/test/sort/sort.py"),
+            ":" + self.script(),
+        )
+        self.rsyncUtil(sseKeyFile, ":" + self.data("keyFile"))
         os.unlink(sseKeyFile)
 
     def _runScript(self, toilOptions):
         # Set memory requirements so that sort jobs can be run
         # on small instances, but merge jobs must be run on large
         # instances
-        toilOptions.extend(['--provisioner=aws', '--batchSystem=mesos',
-                            '--nodeTypes=' + ",".join(self.instanceTypes),
-                            '--maxNodes=' + ",".join(self.numWorkers)])
-        runCommand = [self.python(), self.script(), '--fileToSort=/home/s3am/bin/asadmin', '--sortMemory=0.6G', '--mergeMemory=3.0G']
+        toilOptions.extend(
+            [
+                "--provisioner=aws",
+                "--batchSystem=mesos",
+                "--nodeTypes=" + ",".join(self.instanceTypes),
+                "--maxNodes=" + ",".join(self.numWorkers),
+            ]
+        )
+        runCommand = [
+            self.python(),
+            self.script(),
+            "--fileToSort=/home/s3am/bin/asadmin",
+            "--sortMemory=0.6G",
+            "--mergeMemory=3.0G",
+        ]
         runCommand.extend(toilOptions)
-        runCommand.append('--sseKey=' + self.data('keyFile'))
+        runCommand.append("--sseKey=" + self.data("keyFile"))
         self.sshUtil(runCommand)
 
     @integrative
     @needs_aws_ec2
     def testAutoScale(self):
         self.instanceTypes = ["t2.small", "m5a.large"]
-        self.numWorkers = ['2', '1']
+        self.numWorkers = ["2", "1"]
         self._test()
 
 
@@ -401,16 +544,17 @@ class AWSAutoscaleTestMultipleNodeTypes(AbstractAWSAutoscaleTest):
 @pytest.mark.timeout(1200)
 class AWSRestartTest(AbstractAWSAutoscaleTest):
     """This test insures autoscaling works on a restarted Toil run."""
+
     def __init__(self, name):
         super().__init__(name)
-        self.clusterName = 'restart-test-' + str(uuid4())
-        self.scriptName = 'restartScript.py'
+        self.clusterName = "restart-test-" + str(uuid4())
+        self.scriptName = "restartScript.py"
 
     def setUp(self):
         super().setUp()
-        self.instanceTypes = ['t2.small']
-        self.numWorkers = ['1']
-        self.jobStore = f'aws:{self.awsRegion()}:restart-{uuid4()}'
+        self.instanceTypes = ["t2.small"]
+        self.numWorkers = ["1"]
+        self.jobStore = f"aws:{self.awsRegion()}:restart-{uuid4()}"
 
     def _getScript(self):
         def restartScript():
@@ -421,38 +565,56 @@ class AWSRestartTest(AbstractAWSAutoscaleTest):
             from toil.job import Job
 
             def f0(job):
-                if 'FAIL' in os.environ:
-                    raise RuntimeError('failed on purpose')
+                if "FAIL" in os.environ:
+                    raise RuntimeError("failed on purpose")
 
-            if __name__ == '__main__':
+            if __name__ == "__main__":
                 parser = ArgumentParser()
                 Job.Runner.addToilOptions(parser)
                 options = parser.parse_args()
-                rootJob = Job.wrapJobFn(f0, cores=0.5, memory='50 M', disk='50 M')
+                rootJob = Job.wrapJobFn(f0, cores=0.5, memory="50 M", disk="50 M")
                 Job.Runner.startToil(rootJob, options)
 
-        script = dedent('\n'.join(getsource(restartScript).split('\n')[1:]))
+        script = dedent("\n".join(getsource(restartScript).split("\n")[1:]))
         self.putScript(script)
 
     def _runScript(self, toilOptions):
         # Use the provisioner in the workflow
-        toilOptions.extend(['--provisioner=aws', '--batchSystem=mesos',
-                            '--nodeTypes=' + ",".join(self.instanceTypes),
-                            '--maxNodes=' + ",".join(self.numWorkers)])
+        toilOptions.extend(
+            [
+                "--provisioner=aws",
+                "--batchSystem=mesos",
+                "--nodeTypes=" + ",".join(self.instanceTypes),
+                "--maxNodes=" + ",".join(self.numWorkers),
+            ]
+        )
         # clean = onSuccess
-        disallowedOptions = ['--clean=always', '--retryCount=2']
-        newOptions = [option for option in toilOptions if option not in disallowedOptions]
+        disallowedOptions = ["--clean=always", "--retryCount=2"]
+        newOptions = [
+            option for option in toilOptions if option not in disallowedOptions
+        ]
         try:
             # include a default memory - on restart the minimum memory requirement is the default, usually 2 GB
-            command = [self.python(), self.script(), '--setEnv', 'FAIL=true', '--defaultMemory=50000000']
+            command = [
+                self.python(),
+                self.script(),
+                "--setEnv",
+                "FAIL=true",
+                "--defaultMemory=50000000",
+            ]
             command.extend(newOptions)
             self.sshUtil(command)
         except subprocess.CalledProcessError:
             pass
         else:
-            self.fail('Command succeeded when we expected failure')
+            self.fail("Command succeeded when we expected failure")
         with timeLimit(600):
-            command = [self.python(), self.script(), '--restart', '--defaultMemory=50000000']
+            command = [
+                self.python(),
+                self.script(),
+                "--restart",
+                "--defaultMemory=50000000",
+            ]
             command.extend(toilOptions)
             self.sshUtil(command)
 
@@ -466,14 +628,17 @@ class AWSRestartTest(AbstractAWSAutoscaleTest):
 class PreemptibleDeficitCompensationTest(AbstractAWSAutoscaleTest):
     def __init__(self, name):
         super().__init__(name)
-        self.clusterName = 'deficit-test-' + str(uuid4())
-        self.scriptName = 'userScript.py'
+        self.clusterName = "deficit-test-" + str(uuid4())
+        self.scriptName = "userScript.py"
 
     def setUp(self):
         super().setUp()
-        self.instanceTypes = ['m5a.large:0.01', "m5a.large"]  # instance needs to be available on the spot market
-        self.numWorkers = ['1', '1']
-        self.jobStore = f'aws:{self.awsRegion()}:deficit-{uuid4()}'
+        self.instanceTypes = [
+            "m5a.large:0.01",
+            "m5a.large",
+        ]  # instance needs to be available on the spot market
+        self.numWorkers = ["1", "1"]
+        self.jobStore = f"aws:{self.awsRegion()}:deficit-{uuid4()}"
 
     def test(self):
         self._test(preemptibleJobs=True)
@@ -490,10 +655,10 @@ class PreemptibleDeficitCompensationTest(AbstractAWSAutoscaleTest):
             # we will observe a deficit of preemptible nodes that the non-preemptible scaler will
             # compensate for by spinning up non-preemptible nodes instead.
             #
-            def job(job, disk='10M', cores=1, memory='10M', preemptible=True):
+            def job(job, disk="10M", cores=1, memory="10M", preemptible=True):
                 pass
 
-            if __name__ == '__main__':
+            if __name__ == "__main__":
                 options = Job.Runner.getDefaultArgumentParser().parse_args()
                 with Toil(options) as toil:
                     if toil.config.restart:
@@ -501,14 +666,19 @@ class PreemptibleDeficitCompensationTest(AbstractAWSAutoscaleTest):
                     else:
                         toil.start(Job.wrapJobFn(job))
 
-        script = dedent('\n'.join(getsource(userScript).split('\n')[1:]))
+        script = dedent("\n".join(getsource(userScript).split("\n")[1:]))
         self.putScript(script)
 
     def _runScript(self, toilOptions):
-        toilOptions.extend(['--provisioner=aws', '--batchSystem=mesos',
-                            '--nodeTypes=' + ",".join(self.instanceTypes),
-                            '--maxNodes=' + ",".join(self.numWorkers)])
-        toilOptions.extend(['--preemptibleCompensation=1.0'])
+        toilOptions.extend(
+            [
+                "--provisioner=aws",
+                "--batchSystem=mesos",
+                "--nodeTypes=" + ",".join(self.instanceTypes),
+                "--maxNodes=" + ",".join(self.numWorkers),
+            ]
+        )
+        toilOptions.extend(["--preemptibleCompensation=1.0"])
         command = [self.python(), self.script()]
         command.extend(toilOptions)
         self.sshUtil(command)