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

toil/lib/accelerators.py

@@ -15,8 +15,9 @@
 """Accelerator (i.e. GPU) utilities for Toil"""
 
 import os
+import string
 import subprocess
-from typing import Dict, List, Set, Union, cast
+from typing import Union, cast
 from xml.dom import minidom
 
 from toil.job import AcceleratorRequirement
@@ -33,13 +34,20 @@ def have_working_nvidia_smi() -> bool:
     it can fulfill a CUDARequirement.
     """
     try:
-        subprocess.check_call(['nvidia-smi'])
-    except (FileNotFoundError, subprocess.CalledProcessError):
+        subprocess.check_call(["nvidia-smi"])
+    except (
+        FileNotFoundError,
+        PermissionError,
+        subprocess.CalledProcessError,
+        OSError,
+        UnicodeDecodeError,
+    ):
         return False
     return True
 
+
 @memoize
-def get_host_accelerator_numbers() -> List[int]:
+def get_host_accelerator_numbers() -> list[int]:
     """
     Work out what accelerator is what.
 
@@ -51,7 +59,12 @@ def get_host_accelerator_numbers() -> List[int]:
     right GPUs as seen from a Docker daemon.
     """
 
-    for number_list_var in ['SLURM_STEP_GPUS', 'SLURM_JOB_GPUS', 'CUDA_VISIBLE_DEVICES', 'NVIDIA_VISIBLE_DEVICES']:
+    for number_list_var in [
+        "SLURM_STEP_GPUS",
+        "SLURM_JOB_GPUS",
+        "CUDA_VISIBLE_DEVICES",
+        "NVIDIA_VISIBLE_DEVICES",
+    ]:
         # Any of these can have a list of GPU numbers, but the CUDA/NVIDIA ones
         # also support a system of GPU GUIDs that we don't support.
         # TODO: If Slurm confinement is set we ignore any attempt to further
@@ -61,7 +74,9 @@ def get_host_accelerator_numbers() -> List[int]:
         if number_list_var in os.environ:
             device_string = os.environ[number_list_var]
             # Parse all the numbers we have
-            device_numbers = [int(part) for part in device_string.split(',') if part.isnumeric()]
+            device_numbers = [
+                int(part) for part in device_string.split(",") if part.isnumeric()
+            ]
             if len(device_numbers) > 0:
                 # We found some numbers, so use those
                 return device_numbers
@@ -69,6 +84,7 @@ def get_host_accelerator_numbers() -> List[int]:
     # If we don't see a set of limits we understand, say we have all nvidia GPUs
     return list(range(count_nvidia_gpus()))
 
+
 @memoize
 def have_working_nvidia_docker_runtime() -> bool:
     """
@@ -76,11 +92,30 @@ def have_working_nvidia_docker_runtime() -> bool:
     """
     try:
         # The runtime injects nvidia-smi; it doesn't seem to have to be in the image we use here
-        subprocess.check_call(['docker', 'run', '--rm', '--runtime', 'nvidia', '--gpus', 'all', 'ubuntu:20.04', 'nvidia-smi'])
-    except (FileNotFoundError, subprocess.CalledProcessError):
+        subprocess.check_call(
+            [
+                "docker",
+                "run",
+                "--rm",
+                "--runtime",
+                "nvidia",
+                "--gpus",
+                "all",
+                "ubuntu:20.04",
+                "nvidia-smi",
+            ]
+        )
+    except (
+        FileNotFoundError,
+        PermissionError,
+        subprocess.CalledProcessError,
+        OSError,
+        UnicodeDecodeError,
+    ):
         return False
     return True
 
+
 @memoize
 def count_nvidia_gpus() -> int:
     """
@@ -100,7 +135,7 @@ def count_nvidia_gpus() -> int:
                 .firstChild,
             ).data
         )
-    except: 
+    except:
         return 0
 
     # TODO: Parse each gpu > product_name > text content and convert to some
@@ -108,7 +143,56 @@ def count_nvidia_gpus() -> int:
 
 
 @memoize
-def get_individual_local_accelerators() -> List[AcceleratorRequirement]:
+def count_amd_gpus() -> int:
+    """
+    Return the number of amd GPUs seen by rocm-smi, or 0 if it is not working.
+    :return:
+    """
+    try:
+        # see if the amd-smi CLI tool is installed
+        # we believe this is the expected output for amd-smi, but we don't actually have and amd gpu to test against
+        # so we assume the output from the amd-smi documentation:
+        # https://rocm.docs.amd.com/projects/amdsmi/en/latest/how-to/using-AMD-SMI-CLI-tool.html
+        out = subprocess.check_output(["amd-smi", "static"])
+        gpu_count = len(
+            [line for line in out.decode("utf-8").split("\n") if line.startswith("gpu")]
+        )
+        return gpu_count
+    except (
+        FileNotFoundError,
+        PermissionError,
+        subprocess.SubprocessError,
+        OSError,
+        UnicodeDecodeError,
+    ):
+        # if the amd-smi command fails, try rocm-smi
+        # if a different exception is raised, something other than the subprocess call is wrong
+        pass
+    try:
+        # similarly, since we don't have an AMD gpu to test against, assume the output from the rocm-smi documentation:
+        # https://rocm.blogs.amd.com/software-tools-optimization/affinity/part-2/README.html#gpu-numa-configuration-rocm-smi-showtoponuma
+        out = subprocess.check_output(["rocm-smi"])
+        gpu_count = len(
+            [
+                line
+                for line in out.decode("utf-8").split("\n")
+                if len(line) > 0 and line[0] in string.digits
+            ]
+        )
+        return gpu_count
+    except (
+        FileNotFoundError,
+        PermissionError,
+        subprocess.SubprocessError,
+        OSError,
+        UnicodeDecodeError,
+    ):
+        pass
+    return 0
+
+
+@memoize
+def get_individual_local_accelerators() -> list[AcceleratorRequirement]:
     """
     Determine all the local accelerators available. Report each with count 1,
     in the order of the number that can be used to assign them.
@@ -117,10 +201,22 @@ def get_individual_local_accelerators() -> List[AcceleratorRequirement]:
     accelerator assignment API.
     """
 
-    # For now we only know abput nvidia GPUs
-    return [{'kind': 'gpu', 'brand': 'nvidia', 'api': 'cuda', 'count': 1} for _ in range(count_nvidia_gpus())]
-
-def get_restrictive_environment_for_local_accelerators(accelerator_numbers : Union[Set[int], List[int]]) -> Dict[str, str]:
+    gpus: list[AcceleratorRequirement] = [
+        {"kind": "gpu", "brand": "nvidia", "api": "cuda", "count": 1}
+        for _ in range(count_nvidia_gpus())
+    ]
+    gpus.extend(
+        [
+            {"kind": "gpu", "brand": "amd", "api": "rocm", "count": 1}
+            for _ in range(count_amd_gpus())
+        ]
+    )
+    return gpus
+
+
+def get_restrictive_environment_for_local_accelerators(
+    accelerator_numbers: Union[set[int], list[int]]
+) -> dict[str, str]:
     """
     Get environment variables which can be applied to a process to restrict it
     to using only the given accelerator numbers.
@@ -131,11 +227,12 @@ def get_restrictive_environment_for_local_accelerators(accelerator_numbers : Uni
 
     # Since we only know about nvidia GPUs right now, we can just say our
     # accelerator numbering space is the same as nvidia's GPU numbering space.
-    gpu_list = ','.join(str(i) for i in accelerator_numbers)
+    gpu_list = ",".join(str(i) for i in accelerator_numbers)
 
     # Put this in several places: CUDA_VISIBLE_DEVICES for controlling
     # processes right here, and SINGULARITYENV_CUDA_VISIBLE_DEVICES for
     # propagating to Singularity containers.
-    return {'CUDA_VISIBLE_DEVICES': gpu_list, 'SINGULARITYENV_CUDA_VISIBLE_DEVICES': gpu_list}
-
-
+    return {
+        "CUDA_VISIBLE_DEVICES": gpu_list,
+        "SINGULARITYENV_CUDA_VISIBLE_DEVICES": gpu_list,
+    }

toil/lib/aws/__init__.py

@@ -16,23 +16,25 @@ import logging
 import os
 import re
 import socket
-import toil.lib.retry
+from collections.abc import MutableMapping
 from http.client import HTTPException
-from typing import Dict, MutableMapping, Optional, Union, Literal
+from typing import TYPE_CHECKING, Literal, Optional, Union
 from urllib.error import URLError
 from urllib.request import urlopen
 
 from botocore.exceptions import ClientError
 
-from mypy_boto3_s3.literals import BucketLocationConstraintType
+import toil.lib.retry
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3.literals import BucketLocationConstraintType
 
-AWSRegionName = Union[BucketLocationConstraintType, Literal["us-east-1"]]
+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]
+    error=ClientError, error_codes=[404, 500, 502, 503, 504]
 )
 
 logger = logging.getLogger(__name__)
@@ -40,6 +42,7 @@ logger = logging.getLogger(__name__)
 # This file isn't allowed to import anything that depends on Boto or Boto3,
 # which may not be installed, because it has to be importable everywhere.
 
+
 def get_current_aws_region() -> Optional[str]:
     """
     Return the AWS region that the currently configured AWS zone (see
@@ -49,11 +52,13 @@ def get_current_aws_region() -> Optional[str]:
     aws_zone = get_current_aws_zone()
     return zone_to_region(aws_zone) if aws_zone else None
 
+
 def get_aws_zone_from_environment() -> Optional[str]:
     """
     Get the AWS zone from TOIL_AWS_ZONE if set.
     """
-    return os.environ.get('TOIL_AWS_ZONE', None)
+    return os.environ.get("TOIL_AWS_ZONE", None)
+
 
 def get_aws_zone_from_metadata() -> Optional[str]:
     """
@@ -68,11 +73,15 @@ def get_aws_zone_from_metadata() -> Optional[str]:
         # Use the ECS metadata service
         logger.debug("Fetch AZ from ECS metadata")
         try:
-            resp = json.load(urlopen(os.environ['ECS_CONTAINER_METADATA_URI_V4'] + '/task', timeout=1))
+            resp = json.load(
+                urlopen(
+                    os.environ["ECS_CONTAINER_METADATA_URI_V4"] + "/task", timeout=1
+                )
+            )
             logger.debug("ECS metadata: %s", resp)
             if isinstance(resp, dict):
                 # We found something. Go with that.
-                return resp.get('AvailabilityZone')
+                return resp.get("AvailabilityZone")
         except (json.decoder.JSONDecodeError, KeyError, URLError) as e:
             # We're on ECS but can't get the metadata. That's odd.
             logger.warning("Skipping ECS metadata due to error: %s", e)
@@ -93,6 +102,7 @@ def get_aws_zone_from_metadata() -> Optional[str]:
             logger.warning("Skipping EC2 metadata due to error: %s", e)
     return None
 
+
 def get_aws_zone_from_boto() -> Optional[str]:
     """
     Get the AWS zone from the Boto3 config file or from AWS_DEFAULT_REGION, if it is configured and the
@@ -100,28 +110,30 @@ def get_aws_zone_from_boto() -> Optional[str]:
     """
     try:
         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
+            zone += "a"  # derive an availability zone in the region
         return zone
     except ImportError:
         pass
     return None
 
+
 def get_aws_zone_from_environment_region() -> Optional[str]:
     """
     Pick an AWS zone in the region defined by TOIL_AWS_REGION, if it is set.
     """
-    aws_region = os.environ.get('TOIL_AWS_REGION')
+    aws_region = os.environ.get("TOIL_AWS_REGION")
     if aws_region is not None:
         # If a region is specified, use the first zone in the region.
-        return aws_region + 'a'
+        return aws_region + "a"
     # Otherwise, don't pick a region and let us fall back on the next method.
     return None
 
+
 def get_current_aws_zone() -> Optional[str]:
     """
     Get the currently configured or occupied AWS zone to use.
@@ -137,63 +149,78 @@ def get_current_aws_zone() -> Optional[str]:
     Finally, if we have boto2, and a default region is configured in Boto 2,
     chooses a zone in that region.
 
-    Returns None if no method can produce a zone to use.
+    Returns 'us-east-1a' if no method can produce a zone to use.
     """
-    return get_aws_zone_from_environment() or \
-        get_aws_zone_from_metadata() or \
-        get_aws_zone_from_environment_region() or \
-        get_aws_zone_from_boto()
+    return (
+        get_aws_zone_from_environment()
+        or get_aws_zone_from_metadata()
+        or get_aws_zone_from_environment_region()
+        or get_aws_zone_from_boto()
+        or "us-east-1a"
+    )  # AWS's native default
+
 
 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])$')
+    availability_zone = re.compile(r"^([a-z]{2}-[a-z]+-[1-9][0-9]*)([a-z])$")
     m = availability_zone.match(zone)
     if not m:
         raise ValueError(f"Can't extract region from availability zone '{zone}'")
     return m.group(1)
 
+
 def running_on_ec2() -> bool:
     """
     Return True if we are currently running on EC2, and false otherwise.
     """
+
     # TODO: Move this to toil.lib.ec2 and make toil.lib.ec2 importable without boto?
     def file_begins_with(path, prefix):
         with open(path) as f:
             return f.read(len(prefix)) == prefix
 
-    hv_uuid_path = '/sys/hypervisor/uuid'
-    if os.path.exists(hv_uuid_path) and file_begins_with(hv_uuid_path, 'ec2'):
+    hv_uuid_path = "/sys/hypervisor/uuid"
+    if os.path.exists(hv_uuid_path) and file_begins_with(hv_uuid_path, "ec2"):
         return True
     # Some instances do not have the /sys/hypervisor/uuid file, so check the identity document instead.
     # See https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html
     try:
-        urlopen('http://169.254.169.254/latest/dynamic/instance-identity/document', timeout=1)
+        urlopen(
+            "http://169.254.169.254/latest/dynamic/instance-identity/document",
+            timeout=1,
+        )
         return True
     except (URLError, socket.timeout, HTTPException):
         return False
 
+
 def running_on_ecs() -> bool:
     """
     Return True if we are currently running on Amazon ECS, and false otherwise.
     """
     # We only care about relatively current ECS
-    return 'ECS_CONTAINER_METADATA_URI_V4' in os.environ
+    return "ECS_CONTAINER_METADATA_URI_V4" in os.environ
+
 
-def build_tag_dict_from_env(environment: MutableMapping[str, str] = os.environ) -> Dict[str, str]:
+def build_tag_dict_from_env(
+    environment: MutableMapping[str, str] = os.environ
+) -> dict[str, str]:
     tags = dict()
-    owner_tag = environment.get('TOIL_OWNER_TAG')
+    owner_tag = environment.get("TOIL_OWNER_TAG")
     if owner_tag:
-        tags.update({'Owner': owner_tag})
+        tags.update({"Owner": owner_tag})
 
-    user_tags = environment.get('TOIL_AWS_TAGS')
+    user_tags = environment.get("TOIL_AWS_TAGS")
     if user_tags:
         try:
             json_user_tags = json.loads(user_tags)
             if isinstance(json_user_tags, dict):
                 tags.update(json.loads(user_tags))
             else:
-                logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')
+                logger.error(
+                    'TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}'
+                )
                 exit(1)
         except json.decoder.JSONDecodeError:
             logger.error('TOIL_AWS_TAGS must be in JSON format: {"key" : "value", ...}')