toil 6.1.0__py3-none-any.whl → 7.0.0__py3-none-any.whl

toil/jobStores/aws/utils.py

@@ -17,26 +17,25 @@ import logging
 import os
 import types
 from ssl import SSLError
-from typing import Optional, cast, TYPE_CHECKING
+from typing import Optional, cast, TYPE_CHECKING, Dict, List, Tuple
 
 from boto3.s3.transfer import TransferConfig
-from boto.exception import SDBResponseError
 from botocore.client import Config
 from botocore.exceptions import ClientError
+from mypy_boto3_sdb.type_defs import ItemTypeDef, AttributeTypeDef
 
-from toil.lib.aws import session
-from toil.lib.aws.utils import connection_reset, get_bucket_region
+from toil.lib.aws import session, AWSServerErrors
+from toil.lib.aws.utils import connection_error, get_bucket_region
 from toil.lib.compatibility import compat_bytes
 from toil.lib.retry import (DEFAULT_DELAYS,
                             DEFAULT_TIMEOUT,
-                            ErrorCondition,
                             get_error_code,
                             get_error_message,
                             get_error_status,
                             old_retry,
                             retry)
 if TYPE_CHECKING:
-    from mypy_boto3_s3 import S3Client, S3ServiceResource
+    from mypy_boto3_s3 import S3ServiceResource
 
 logger = logging.getLogger(__name__)
 
@@ -126,11 +125,11 @@ class SDBHelper:
         return cls._maxChunks() * cls.maxValueSize
 
     @classmethod
-    def binaryToAttributes(cls, binary):
+    def binaryToAttributes(cls, binary) -> Dict[str, str]:
         """
         Turn a bytestring, or None, into SimpleDB attributes.
         """
-        if binary is None: return {'numChunks': 0}
+        if binary is None: return {'numChunks': '0'}
         assert isinstance(binary, bytes)
         assert len(binary) <= cls.maxBinarySize()
         # The use of compression is just an optimization. We can't include it in the maxValueSize
@@ -144,10 +143,41 @@ class SDBHelper:
         assert len(encoded) <= cls._maxEncodedSize()
         n = cls.maxValueSize
         chunks = (encoded[i:i + n] for i in range(0, len(encoded), n))
-        attributes = {cls._chunkName(i): chunk for i, chunk in enumerate(chunks)}
-        attributes.update({'numChunks': len(attributes)})
+        attributes = {cls._chunkName(i): chunk.decode("utf-8") for i, chunk in enumerate(chunks)}
+        attributes.update({'numChunks': str(len(attributes))})
         return attributes
 
+    @classmethod
+    def attributeDictToList(cls, attributes: Dict[str, str]) -> List[AttributeTypeDef]:
+        """
+        Convert the attribute dict (ex: from binaryToAttributes) into a list of attribute typed dicts
+        to be compatible with boto3 argument syntax
+        :param attributes: Dict[str, str], attribute in object form
+        :return: List[AttributeTypeDef], list of attributes in typed dict form
+        """
+        return [{"Name": name, "Value": value} for name, value in attributes.items()]
+
+    @classmethod
+    def attributeListToDict(cls, attributes: List[AttributeTypeDef]) -> Dict[str, str]:
+        """
+        Convert the attribute boto3 representation of list of attribute typed dicts
+        back to a dictionary with name, value pairs
+        :param attribute: List[AttributeTypeDef, attribute in typed dict form
+        :return: Dict[str, str], attribute in dict form
+        """
+        return {attribute["Name"]: attribute["Value"] for attribute in attributes}
+
+    @classmethod
+    def get_attributes_from_item(cls, item: ItemTypeDef, keys: List[str]) -> List[Optional[str]]:
+        return_values: List[Optional[str]] = [None for _ in keys]
+        mapped_indices: Dict[str, int] = {name: index for index, name in enumerate(keys)}
+        for attribute in item["Attributes"]:
+            name = attribute["Name"]
+            value = attribute["Value"]
+            if name in mapped_indices:
+                return_values[mapped_indices[name]] = value
+        return return_values
+
     @classmethod
     def _chunkName(cls, i):
         return str(i).zfill(3)
@@ -166,14 +196,21 @@ class SDBHelper:
         return 'numChunks'
 
     @classmethod
-    def attributesToBinary(cls, attributes):
+    def attributesToBinary(cls, attributes: List[AttributeTypeDef]) -> Tuple[bytes, int]:
         """
         :rtype: (str|None,int)
         :return: the binary data and the number of chunks it was composed from
         """
-        chunks = [(int(k), v) for k, v in attributes.items() if cls._isValidChunkName(k)]
+        chunks = []
+        numChunks: int = 0
+        for attribute in attributes:
+            name = attribute["Name"]
+            value = attribute["Value"]
+            if cls._isValidChunkName(name):
+                chunks.append((int(name), value))
+            if name == "numChunks":
+                numChunks = int(value)
         chunks.sort()
-        numChunks = int(attributes['numChunks'])
         if numChunks:
             serializedJob = b''.join(v.encode() for k, v in chunks)
             compressed = base64.b64decode(serializedJob)
@@ -193,10 +230,7 @@ def fileSizeAndTime(localFilePath):
     return file_stat.st_size, file_stat.st_mtime
 
 
-@retry(errors=[ErrorCondition(
-    error=ClientError,
-    error_codes=[404, 500, 502, 503, 504]
-)])
+@retry(errors=[AWSServerErrors])
 def uploadFromPath(localFilePath: str,
                    resource,
                    bucketName: str,
@@ -232,10 +266,7 @@ def uploadFromPath(localFilePath: str,
     return version
 
 
-@retry(errors=[ErrorCondition(
-    error=ClientError,
-    error_codes=[404, 500, 502, 503, 504]
-)])
+@retry(errors=[AWSServerErrors])
 def uploadFile(readable,
                resource,
                bucketName: str,
@@ -287,10 +318,7 @@ class ServerSideCopyProhibitedError(RuntimeError):
     insists that you pay to download and upload the data yourself instead.
     """
 
-@retry(errors=[ErrorCondition(
-    error=ClientError,
-    error_codes=[404, 500, 502, 503, 504]
-)])
+@retry(errors=[AWSServerErrors])
 def copyKeyMultipart(resource: "S3ServiceResource",
                      srcBucketName: str,
                      srcKeyName: str,
@@ -439,9 +467,9 @@ def sdb_unavailable(e):
 
 
 def no_such_sdb_domain(e):
-    return (isinstance(e, SDBResponseError)
-            and e.error_code
-            and e.error_code.endswith('NoSuchDomain'))
+    return (isinstance(e, ClientError)
+            and get_error_code(e)
+            and get_error_code(e).endswith('NoSuchDomain'))
 
 
 def retryable_ssl_error(e):
@@ -452,7 +480,7 @@ def retryable_ssl_error(e):
 def retryable_sdb_errors(e):
     return (sdb_unavailable(e)
             or no_such_sdb_domain(e)
-            or connection_reset(e)
+            or connection_error(e)
             or retryable_ssl_error(e))
 
 

toil/jobStores/conftest.py

@@ -17,7 +17,7 @@
 collect_ignore = []
 
 try:
-    import boto
-    print(boto.__file__)  # prevent this import from being removed
+    import boto3
+    print(boto3.__file__)  # prevent this import from being removed
 except ImportError:
     collect_ignore.append("aws")

toil/jobStores/fileJobStore.py

@@ -113,7 +113,7 @@ class FileJobStore(AbstractJobStore):
             os.mkdir(self.jobStoreDir)
         except OSError as e:
             if e.errno == errno.EEXIST:
-                raise JobStoreExistsException(self.jobStoreDir)
+                raise JobStoreExistsException(self.jobStoreDir, "file")
             else:
                 raise
         os.makedirs(self.jobsDir, exist_ok=True)
@@ -127,7 +127,7 @@ class FileJobStore(AbstractJobStore):
 
     def resume(self):
         if not os.path.isdir(self.jobStoreDir):
-            raise NoSuchJobStoreException(self.jobStoreDir)
+            raise NoSuchJobStoreException(self.jobStoreDir, "file")
         super().resume()
 
     def destroy(self):

toil/jobStores/googleJobStore.py

@@ -164,7 +164,7 @@ class GoogleJobStore(AbstractJobStore):
         try:
             self.bucket = self.storageClient.create_bucket(self.bucketName)
         except exceptions.Conflict:
-            raise JobStoreExistsException(self.locator)
+            raise JobStoreExistsException(self.locator, "google")
         super().initialize(config)
 
         # set up sever side encryption after we set up config in super
@@ -178,7 +178,7 @@ class GoogleJobStore(AbstractJobStore):
         try:
             self.bucket = self.storageClient.get_bucket(self.bucketName)
         except exceptions.NotFound:
-            raise NoSuchJobStoreException(self.locator)
+            raise NoSuchJobStoreException(self.locator, "google")
         super().resume()
 
     @google_retry
@@ -209,8 +209,7 @@ class GoogleJobStore(AbstractJobStore):
 
     def assign_job_id(self, job_description):
         jobStoreID = self._new_job_id()
-        log.debug("Assigning ID to job %s for '%s'",
-                  jobStoreID, '<no command>' if job_description.command is None else job_description.command)
+        log.debug("Assigning ID to job %s", jobStoreID)
         job_description.jobStoreID = jobStoreID
 
     @contextmanager

toil/leader.py

@@ -119,7 +119,12 @@ class Leader:
         if self.config.write_messages is None:
             # The user hasn't specified a place for the message bus so we
             # should make one.
-            self.config.write_messages = gen_message_bus_path()
+            # pass in coordination_dir for toil-cwl-runner; we want to obey --tmpdir-prefix
+            # from cwltool and we change the coordination_dir when detected. we don't want
+            # to make another config attribute so put the message bus in the already prefixed dir
+            # if a coordination_dir is provided normally, we can still put the bus in there
+            # as the coordination dir should serve a similar purpose to the tmp directory
+            self.config.write_messages = gen_message_bus_path(config.coordination_dir)
 
         # Message bus messages need to go to the given file.
         # Keep a reference to the return value so the listener stays alive.
@@ -289,7 +294,11 @@ class Leader:
                 for job_id in self.toilState.totalFailedJobs:
                     # Refresh all the failed jobs to get e.g. the log file IDs that the workers wrote
                     self.toilState.reset_job(job_id)
-                    failed_jobs.append(self.toilState.get_job(job_id))
+                    try:
+                        failed_jobs.append(self.toilState.get_job(job_id))
+                    except NoSuchJobException:
+                        # Job actually finished and was removed
+                        pass
 
                 logger.info("Failed jobs at end of the run: %s", ' '.join(str(j) for j in failed_jobs))
                 raise FailedJobsException(self.jobStore, failed_jobs, exit_code=self.recommended_fail_exit_code)
@@ -522,10 +531,10 @@ class Leader:
                          "manager: %s", readyJob.jobStoreID)
         elif readyJob.jobStoreID in self.toilState.hasFailedSuccessors:
             self._processFailedSuccessors(job_id)
-        elif readyJob.command is not None or result_status != 0:
-            # The job has a command it must be run before any successors.
+        elif readyJob.has_body() or result_status != 0:
+            # The job has a body it must be run before any successors.
             # Similarly, if the job previously failed we rerun it, even if it doesn't have a
-            # command to run, to eliminate any parts of the stack now completed.
+            # body to run, to eliminate any parts of the stack now completed.
             isServiceJob = readyJob.jobStoreID in self.toilState.service_to_client
 
             # We want to run the job, and expend one of its "tries" (possibly
@@ -551,6 +560,7 @@ class Leader:
                 for serviceID in serviceJobList:
                     if serviceID in self.toilState.service_to_client:
                         raise RuntimeError(f"The ready service ID: {serviceID} was already added.")
+                    # TODO: Why do we refresh here?
                     self.toilState.reset_job(serviceID)
                     serviceHost = self.toilState.get_job(serviceID)
                     self.toilState.service_to_client[serviceID] = readyJob.jobStoreID
@@ -896,11 +906,6 @@ class Leader:
             workerCommand.append('--context')
             workerCommand.append(base64.b64encode(pickle.dumps(context)).decode('utf-8'))
 
-        # We locally override the command. This shouldn't get persisted back to
-        # the job store, or we will detach the job body from the job
-        # description. TODO: Don't do it this way! It's weird!
-        jobNode.command = ' '.join(workerCommand)
-
         omp_threads = os.environ.get('OMP_NUM_THREADS') \
             or str(max(1, int(jobNode.cores)))  # make sure OMP_NUM_THREADS is a positive integer
 
@@ -910,7 +915,7 @@ class Leader:
         }
 
         # jobBatchSystemID is an int for each job
-        jobBatchSystemID = self.batchSystem.issueBatchJob(jobNode, job_environment=job_environment)
+        jobBatchSystemID = self.batchSystem.issueBatchJob(' '.join(workerCommand), jobNode, job_environment=job_environment)
         # Record the job by the ID the batch system will use to talk about it with us
         self.issued_jobs_by_batch_system_id[jobBatchSystemID] = jobNode.jobStoreID
         # Record that this job is issued right now and shouldn't e.g. be issued again.
@@ -1048,7 +1053,7 @@ class Leader:
             jobs = [job for job in jobs if job.preemptible == preemptible]
         return jobs
 
-    def killJobs(self, jobsToKill):
+    def killJobs(self, jobsToKill, exit_reason: BatchJobExitReason = BatchJobExitReason.KILLED):
         """
         Kills the given set of jobs and then sends them for processing.
 
@@ -1062,7 +1067,7 @@ class Leader:
             self.batchSystem.killBatchJobs(jobsToKill)
             for jobBatchSystemID in jobsToKill:
                 # Reissue immediately, noting that we killed the job
-                willRerun = self.process_finished_job(jobBatchSystemID, 1, exit_reason=BatchJobExitReason.KILLED)
+                willRerun = self.process_finished_job(jobBatchSystemID, 1, exit_reason=exit_reason)
 
                 if willRerun:
                     # Compose a list of all the jobs that will run again
@@ -1092,7 +1097,7 @@ class Leader:
                                 str(runningJobs[jobBatchSystemID]),
                                 str(maxJobDuration))
                     jobsToKill.append(jobBatchSystemID)
-            reissued = self.killJobs(jobsToKill)
+            reissued = self.killJobs(jobsToKill, exit_reason=BatchJobExitReason.MAXJOBDURATION)
             if len(jobsToKill) > 0:
                 # Summarize our actions
                 logger.info("Killed %d over long jobs and reissued %d of them", len(jobsToKill), len(reissued))
@@ -1130,7 +1135,7 @@ class Leader:
             if timesMissing == killAfterNTimesMissing:
                 self.reissueMissingJobs_missingHash.pop(jobBatchSystemID)
                 jobsToKill.append(jobBatchSystemID)
-        self.killJobs(jobsToKill)
+        self.killJobs(jobsToKill, exit_reason=BatchJobExitReason.MISSING)
         return len( self.reissueMissingJobs_missingHash ) == 0 #We use this to inform
         #if there are missing jobs
 
@@ -1168,7 +1173,7 @@ class Leader:
                                          exit_reason: Optional[BatchJobExitReason] = None,
                                          batch_system_id: Optional[int] = None) -> bool:
         """
-        Process a finished JobDescription based upon its succees or failure.
+        Process a finished JobDescription based upon its success or failure.
 
         If wall-clock time is available, informs the cluster scaler about the
         job finishing.
@@ -1191,19 +1196,62 @@ class Leader:
             logger.debug("Job %s continues to exist (i.e. has more to do)", finished_job)
             try:
                 # Reload the job as modified by the worker
-                self.toilState.reset_job(job_store_id)
-                replacement_job = self.toilState.get_job(job_store_id)
+                if finished_job.has_body():
+                    # The worker was expected to do some work. We expect the
+                    # worker to have updated the job description.
+
+                    # If the job succeeded, we wait around to see the update
+                    # and fail the job if we don't see it.
+                    if result_status == 0:
+                        timeout = self.config.job_store_timeout
+                        complaint = (
+                            f"has no new version available after {timeout} "
+                            "seconds. Either worker updates to "
+                            "the job store are delayed longer than your "
+                            "--jobStoreTimeout, or the worker trying to run the "
+                            "job was killed (or never started)."
+                        )
+                    else:
+                        timeout = 0
+                        complaint = (
+                            "has no new version available immediately. The "
+                            "batch system may have killed (or never started) "
+                            "the Toil worker."
+                        )
+                    change_detected = self.toilState.reset_job_expecting_change(job_store_id, timeout)
+                    replacement_job = self.toilState.get_job(job_store_id)
+
+                    if not change_detected:
+                        logger.warning(
+                            'Job %s %s',
+                            replacement_job,
+                            complaint
+                        )
+                        if result_status == 0:
+                            # Make the job fail because we ran it and it finished
+                            # and we never heard back.
+                            logger.error(
+                                'Marking ostensibly successful job %s that did '
+                                'not report in to the job store before '
+                                '--jobStoreTimeout as having been partitioned '
+                                'from us.',
+                                replacement_job
+                            )
+                            result_status = EXIT_STATUS_UNAVAILABLE_VALUE
+                            exit_reason = BatchJobExitReason.PARTITION
+                else:
+                    # If there was no body sent, the worker won't commit any
+                    # changes to the job description. So don't wait around for
+                    # any and don't complain if we don't see them.
+                    self.toilState.reset_job(job_store_id)
+                    replacement_job = self.toilState.get_job(job_store_id)
+
             except NoSuchJobException:
                 # We have a ghost job - the job has been deleted but a stale
                 # read from e.g. a non-POSIX-compliant filesystem gave us a
                 # false positive when we checked for its existence. Process the
                 # job from here as any other job removed from the job store.
-                # This is a hack until we can figure out how to actually always
-                # have a strongly-consistent communications channel. See
-                # https://github.com/BD2KGenomics/toil/issues/1091
-                logger.warning('Got a stale read for job %s; caught its '
-                'completion in time, but other jobs may try to run twice! Fix '
-                'the consistency of your job store storage!', finished_job)
+                logger.debug("Job %s is actually complete upon closer inspection", finished_job)
                 self.processRemovedJob(finished_job, result_status)
                 return False
             if replacement_job.logJobStoreFileID is not None:

toil/lib/aws/__init__.py

@@ -16,11 +16,25 @@ import logging
 import os
 import re
 import socket
+import toil.lib.retry
 from http.client import HTTPException
-from typing import Dict, MutableMapping, Optional
+from typing import Dict, MutableMapping, Optional, Union, Literal
 from urllib.error import URLError
 from urllib.request import urlopen
 
+from botocore.exceptions import ClientError
+
+from mypy_boto3_s3.literals import BucketLocationConstraintType
+
+AWSRegionName = Union[BucketLocationConstraintType, Literal["us-east-1"]]
+
+# These are errors where we think something randomly
+# went wrong on the AWS side and we ought to retry.
+AWSServerErrors = toil.lib.retry.ErrorCondition(
+    error=ClientError,
+    error_codes=[404, 500, 502, 503, 504]
+)
+
 logger = logging.getLogger(__name__)
 
 # This file isn't allowed to import anything that depends on Boto or Boto3,
@@ -67,11 +81,10 @@ def get_aws_zone_from_metadata() -> Optional[str]:
         # metadata.
         try:
             # Use the EC2 metadata service
-            import boto
-            str(boto)  # to prevent removal of the import
-            from boto.utils import get_instance_metadata
+            from ec2_metadata import ec2_metadata
+
             logger.debug("Fetch AZ from EC2 metadata")
-            return get_instance_metadata()['placement']['availability-zone']
+            return ec2_metadata.availability_zone
         except ImportError:
             # This is expected to happen a lot
             logger.debug("No boto to fetch ECS metadata")
@@ -82,12 +95,15 @@ def get_aws_zone_from_metadata() -> Optional[str]:
 
 def get_aws_zone_from_boto() -> Optional[str]:
     """
-    Get the AWS zone from the Boto config file, if it is configured and the
-    boto module is available.
+    Get the AWS zone from the Boto3 config file or from AWS_DEFAULT_REGION, if it is configured and the
+    boto3 module is available.
     """
     try:
-        import boto
-        zone = boto.config.get('Boto', 'ec2_region_name')
+        import boto3
+        from session import client
+        boto3_session = boto3.session.Session()
+        # this should check AWS_DEFAULT_REGION and ~/.aws/config
+        zone = boto3_session.region_name
         if zone is not None:
             zone += 'a'  # derive an availability zone in the region
         return zone
@@ -128,7 +144,7 @@ def get_current_aws_zone() -> Optional[str]:
         get_aws_zone_from_environment_region() or \
         get_aws_zone_from_boto()
 
-def zone_to_region(zone: str) -> str:
+def zone_to_region(zone: str) -> AWSRegionName:
     """Get a region (e.g. us-west-2) from a zone (e.g. us-west-1c)."""
     # re.compile() caches the regex internally so we don't have to
     availability_zone = re.compile(r'^([a-z]{2}-[a-z]+-[1-9][0-9]*)([a-z])$')

toil/lib/aws/iam.py

@@ -257,8 +257,8 @@ def get_policy_permissions(region: str) -> AllowedActionCollection:
     :param zone: AWS zone to connect to
     """
 
-    iam: IAMClient = cast(IAMClient, get_client('iam', region))
-    sts: STSClient = cast(STSClient, get_client('sts', region))
+    iam: IAMClient = get_client('iam', region)
+    sts: STSClient = get_client('sts', region)
     #TODO Condider effect: deny at some point
     allowed_actions: AllowedActionCollection = defaultdict(lambda: {'Action': [], 'NotAction': []})
     try:

toil/lib/aws/session.py

@@ -15,16 +15,21 @@ import collections
 import logging
 import os
 import threading
-from typing import Dict, Optional, Tuple, cast
+from typing import Dict, Optional, Tuple, cast, Union, Literal, overload, TypeVar
 
 import boto3
 import boto3.resources.base
-import boto.connection
 import botocore
 from boto3 import Session
 from botocore.client import Config
 from botocore.session import get_session
 from botocore.utils import JSONFileCache
+from mypy_boto3_autoscaling import AutoScalingClient
+from mypy_boto3_ec2 import EC2Client, EC2ServiceResource
+from mypy_boto3_iam import IAMClient, IAMServiceResource
+from mypy_boto3_s3 import S3Client, S3ServiceResource
+from mypy_boto3_sdb import SimpleDBClient
+from mypy_boto3_sts import STSClient
 
 logger = logging.getLogger(__name__)
 
@@ -120,6 +125,13 @@ class AWSConnectionManager:
             storage.item = _new_boto3_session(region_name=region)
         return cast(boto3.session.Session, storage.item)
 
+    @overload
+    def resource(self, region: Optional[str], service_name: Literal["s3"], endpoint_url: Optional[str] = None) -> S3ServiceResource: ...
+    @overload
+    def resource(self, region: Optional[str], service_name: Literal["iam"], endpoint_url: Optional[str] = None) -> IAMServiceResource: ...
+    @overload
+    def resource(self, region: Optional[str], service_name: Literal["ec2"], endpoint_url: Optional[str] = None) -> EC2ServiceResource: ...
+
     def resource(self, region: Optional[str], service_name: str, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
         """
         Get the Boto3 Resource to use with the given service (like 'ec2') in the given region.
@@ -146,7 +158,28 @@ class AWSConnectionManager:
 
         return cast(boto3.resources.base.ServiceResource, storage.item)
 
-    def client(self, region: Optional[str], service_name: str, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> botocore.client.BaseClient:
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["ec2"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> EC2Client: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["iam"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> IAMClient: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["s3"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> S3Client: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["sts"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> STSClient: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["sdb"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> SimpleDBClient: ...
+    @overload
+    def client(self, region: Optional[str], service_name: Literal["autoscaling"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> AutoScalingClient: ...
+
+
+    def client(self, region: Optional[str], service_name: Literal["ec2", "iam", "s3", "sts", "sdb", "autoscaling"], endpoint_url: Optional[str] = None,
+               config: Optional[Config] = None) -> botocore.client.BaseClient:
         """
         Get the Boto3 Client to use with the given service (like 'ec2') in the given region.
 
@@ -159,9 +192,9 @@ class AWSConnectionManager:
             # Don't try and memoize if a custom config is used
             with _init_lock:
                 if endpoint_url is not None:
-                    return self.session(region).client(service_name, endpoint_url=endpoint_url, config=config) # type: ignore
+                    return self.session(region).client(service_name, endpoint_url=endpoint_url, config=config)
                 else:
-                    return self.session(region).client(service_name, config=config) # type: ignore
+                    return self.session(region).client(service_name, config=config)
 
         key = (region, service_name, endpoint_url)
         storage = self.client_cache[key]
@@ -172,25 +205,12 @@ class AWSConnectionManager:
                 if endpoint_url is not None:
                     # The Boto3 stubs are probably missing an overload here too. See:
                     # <https://github.com/vemel/mypy_boto3_builder/issues/121#issuecomment-1011322636>
-                    storage.item = self.session(region).client(service_name, endpoint_url=endpoint_url) # type: ignore
+                    storage.item = self.session(region).client(service_name, endpoint_url=endpoint_url)
                 else:
                     # We might not be able to pass None to Boto3 and have it be the same as no argument.
-                    storage.item = self.session(region).client(service_name) # type: ignore
+                    storage.item = self.session(region).client(service_name)
         return cast(botocore.client.BaseClient , storage.item)
 
-    def boto2(self, region: Optional[str], service_name: str) -> boto.connection.AWSAuthConnection:
-        """
-        Get the connected boto2 connection for the given region and service.
-        """
-        if service_name == 'iam':
-            # IAM connections are regionless
-            region = 'universal'
-        key = (region, service_name)
-        storage = self.boto2_cache[key]
-        if not hasattr(storage, 'item'):
-            with _init_lock:
-                storage.item = getattr(boto, service_name).connect_to_region(region, profile_name=os.environ.get("TOIL_AWS_PROFILE", None))
-        return cast(boto.connection.AWSAuthConnection, storage.item)
 
 # If you don't want your own AWSConnectionManager, we have a global one and some global functions
 _global_manager = AWSConnectionManager()
@@ -205,7 +225,20 @@ def establish_boto3_session(region_name: Optional[str] = None) -> Session:
     # Just use a global version of the manager. Note that we change the argument order!
     return _global_manager.session(region_name)
 
-def client(service_name: str, region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> botocore.client.BaseClient:
+@overload
+def client(service_name: Literal["ec2"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> EC2Client: ...
+@overload
+def client(service_name: Literal["iam"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> IAMClient: ...
+@overload
+def client(service_name: Literal["s3"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> S3Client: ...
+@overload
+def client(service_name: Literal["sts"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> STSClient: ...
+@overload
+def client(service_name: Literal["sdb"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> SimpleDBClient: ...
+@overload
+def client(service_name: Literal["autoscaling"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> AutoScalingClient: ...
+
+def client(service_name: Literal["ec2", "iam", "s3", "sts", "sdb", "autoscaling"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None, config: Optional[Config] = None) -> botocore.client.BaseClient:
     """
     Get a Boto 3 client for a particular AWS service, usable by the current thread.
 
@@ -215,7 +248,14 @@ def client(service_name: str, region_name: Optional[str] = None, endpoint_url: O
     # Just use a global version of the manager. Note that we change the argument order!
     return _global_manager.client(region_name, service_name, endpoint_url=endpoint_url, config=config)
 
-def resource(service_name: str, region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
+@overload
+def resource(service_name: Literal["s3"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> S3ServiceResource: ...
+@overload
+def resource(service_name: Literal["iam"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> IAMServiceResource: ...
+@overload
+def resource(service_name: Literal["ec2"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> EC2ServiceResource: ...
+
+def resource(service_name: Literal["s3", "iam", "ec2"], region_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> boto3.resources.base.ServiceResource:
     """
     Get a Boto 3 resource for a particular AWS service, usable by the current thread.