toil 8.1.0b1__py3-none-any.whl → 9.0.0__py3-none-any.whl

toil/__init__.py

@@ -480,38 +480,3 @@ def logProcessContext(config: "Config") -> None:
     log.info("Running Toil version %s on host %s.", version, socket.gethostname())
     log.debug("Configuration: %s", config.__dict__)
 
-
-try:
-    cache_path = "~/.cache/aws/cached_temporary_credentials"
-    datetime_format = (
-        "%Y-%m-%dT%H:%M:%SZ"  # incidentally the same as the format used by AWS
-    )
-    log = logging.getLogger(__name__)
-
-    # But in addition to our manual cache, we also are going to turn on boto3's
-    # new built-in caching layer.
-
-    def datetime_to_str(dt):
-        """
-        Convert a naive (implicitly UTC) datetime object into a string, explicitly UTC.
-
-        >>> datetime_to_str(datetime(1970, 1, 1, 0, 0, 0))
-        '1970-01-01T00:00:00Z'
-        """
-        return dt.strftime(datetime_format)
-
-    def str_to_datetime(s):
-        """
-        Convert a string, explicitly UTC into a naive (implicitly UTC) datetime object.
-
-        >>> str_to_datetime( '1970-01-01T00:00:00Z' )
-        datetime.datetime(1970, 1, 1, 0, 0)
-
-        Just to show that the constructor args for seconds and microseconds are optional:
-        >>> datetime(1970, 1, 1, 0, 0, 0)
-        datetime.datetime(1970, 1, 1, 0, 0)
-        """
-        return datetime.strptime(s, datetime_format)
-
-except ImportError:
-    pass

toil/batchSystems/abstractBatchSystem.py

@@ -310,7 +310,7 @@ class BatchSystemSupport(AbstractBatchSystem):
     """Partial implementation of AbstractBatchSystem, support methods."""
 
     def __init__(
-        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: int
     ) -> None:
         """
         Initialize initial state of the object.

toil/batchSystems/abstractGridEngineBatchSystem.py

@@ -421,7 +421,7 @@ class AbstractGridEngineBatchSystem(BatchSystemCleanupSupport):
             raise NotImplementedError()
 
     def __init__(
-        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: float
     ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
         self.config = config

toil/batchSystems/awsBatch.py

@@ -83,7 +83,7 @@ class AWSBatchBatchSystem(BatchSystemCleanupSupport):
         return True
 
     def __init__(
-        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: int
     ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
 

toil/batchSystems/cleanup_support.py

@@ -45,7 +45,7 @@ class BatchSystemCleanupSupport(BatchSystemLocalSupport):
         return contexts
 
     def __init__(
-        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: int
     ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
 

toil/batchSystems/kubernetes.py

@@ -21,6 +21,7 @@ cannot yet be launched. That functionality will need to wait for user-mode
 Docker
 """
 import datetime
+import io
 import logging
 import math
 import os
@@ -50,7 +51,8 @@ else:
 from typing import Protocol, TypedDict, runtime_checkable
 
 import urllib3
-import yaml
+import ruamel.yaml as yaml
+import json
 
 # The Right Way to use the Kubernetes module is to `import kubernetes` and then you get all your stuff as like ApiClient. But this doesn't work for the stubs: the stubs seem to only support importing things from the internal modules in `kubernetes` where they are actually defined. See for example <https://github.com/MaterializeInc/kubernetes-stubs/issues/9 and <https://github.com/MaterializeInc/kubernetes-stubs/issues/10>. So we just import all the things we use into our global namespace here.
 from kubernetes.client import (
@@ -77,6 +79,7 @@ from kubernetes.client import (
     V1ResourceRequirements,
     V1SecretVolumeSource,
     V1SecurityContext,
+    V1PodSecurityContext,
     V1Toleration,
     V1Volume,
     V1VolumeMount,
@@ -148,7 +151,7 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
         customObjects: NotRequired[CustomObjectsApi]
 
     def __init__(
-        self, config: Config, maxCores: int, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: int, maxMemory: float, maxDisk: int
     ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
 
@@ -314,7 +317,10 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
                 del here[k]
 
         drop_boring(root_dict)
-        return yaml.dump(root_dict)
+        s = io.StringIO()
+        YAML = yaml.YAML(typ='safe')
+        YAML.dump(root_dict, s)
+        return s.getvalue()
 
     @overload
     def _api(
@@ -803,6 +809,26 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
                     ],
                 )
 
+    class FakeResponse:
+        data: str
+
+    T = TypeVar('T')
+    def _load_kubernetes_object(self, file: str, cls: type[T]) -> T:
+        """
+        Deserialize a YAML representation into a Kubernetes object
+        :param file: Path to YAML file
+        :param cls: Kubernetes API model type for deserialized object
+        :return: Deserialized object
+        """
+        YAML = yaml.YAML(typ='safe')
+        object_def = YAML.load(open('container.yaml').read())
+        # The kubernetes API does not have an actual deserializer, so this is a workaround
+        # See: https://github.com/kubernetes-client/python/issues/977
+        faked_response = self.FakeResponse()
+        faked_response.data = json.dumps(object_def)
+        return ApiClient().deserialize(faked_response, cls)
+
+
     def _create_pod_spec(
         self,
         command: str,
@@ -946,17 +972,24 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
             volume_mounts=mounts,
         )
 
+        if self.config.kubernetes_security_context:
+            container.security_context = self._load_kubernetes_object(self.config.kubernetes_security_context, V1SecurityContext)
+
         # In case security context rules are not allowed to be set, we only apply
         # a security context at all if we need to turn on privileged mode.
         if self.config.kubernetes_privileged:
-            container.security_context = V1SecurityContext(
-                privileged=self.config.kubernetes_privileged
-            )
+            if container.security_context is None:
+                container.security_context = V1SecurityContext()
+            container.security_context.privileged = self.config.kubernetes_privileged
 
         # Wrap the container in a spec
         pod_spec = V1PodSpec(
             containers=[container], volumes=volumes, restart_policy="Never"
         )
+
+        if self.config.kubernetes_pod_security_context:
+            pod_spec.security_context = self._load_kubernetes_object(self.config.kubernetes_pod_security_context, V1PodSecurityContext)
+
         # Tell the spec where to land
         placement.apply(pod_spec)
 
@@ -2126,7 +2159,18 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
             "privileged operations, such as FUSE. On Toil-managed clusters with --enableFuse, "
             "this is set to True. (default: %(default)s)",
         )
-
+        parser.add_argument("--kubernetesPodSecurityContext",
+                            dest='kubernetes_pod_security_context',
+                            type=str,
+                            env_var="TOIL_KUBERNETES_POD_SECURITY_CONTEXT",
+                            default=None,
+                            help="Path to a YAML defining a pod security context to apply to all pods.")
+        parser.add_argument("--kubernetesSecurityContext",
+                            dest='kubernetes_security_context',
+                            type=str,
+                            env_var="TOIL_KUBERNETES_SECURITY_CONTEXT",
+                            default=None,
+                            help="Path to a YAML defining a security context to apply to all containers.")
     OptionType = TypeVar("OptionType")
 
     @classmethod
@@ -2138,3 +2182,5 @@ class KubernetesBatchSystem(BatchSystemCleanupSupport):
         )
         setOption("kubernetes_pod_timeout")
         setOption("kubernetes_privileged")
+        setOption("kubernetes_pod_security_context")
+        setOption("kubernetes_security_context")

toil/batchSystems/local_support.py

@@ -30,7 +30,7 @@ class BatchSystemLocalSupport(BatchSystemSupport):
     """Adds a local queue for helper jobs, useful for CWL & others."""
 
     def __init__(
-        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: int
     ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
         max_local_jobs = (

toil/batchSystems/mesos/batchSystem.py

@@ -40,6 +40,7 @@ from toil.batchSystems.abstractBatchSystem import (
 from toil.batchSystems.local_support import BatchSystemLocalSupport
 from toil.batchSystems.mesos import JobQueue, MesosShape, TaskData, ToilJob
 from toil.batchSystems.options import OptionSetter
+from toil.common import Config
 from toil.job import JobDescription
 from toil.lib.conversions import b_to_mib, mib_to_b
 from toil.lib.memoize import strict_bool
@@ -61,22 +62,26 @@ class MesosBatchSystem(BatchSystemLocalSupport, AbstractScalableBatchSystem, Sch
     """
 
     @classmethod
-    def supportsAutoDeployment(cls):
+    def supportsAutoDeployment(cls) -> bool:
         return True
 
     @classmethod
-    def supportsWorkerCleanup(cls):
+    def supportsWorkerCleanup(cls) -> bool:
         return True
 
     class ExecutorInfo:
-        def __init__(self, nodeAddress, agentId, nodeInfo, lastSeen):
+        def __init__(
+            self, nodeAddress: str, agentId: str, nodeInfo: str, lastSeen: str
+        ) -> None:
             super().__init__()
             self.nodeAddress = nodeAddress
             self.agentId = agentId
             self.nodeInfo = nodeInfo
             self.lastSeen = lastSeen
 
-    def __init__(self, config, maxCores, maxMemory, maxDisk):
+    def __init__(
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: int
+    ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
 
         # The auto-deployed resource representing the user script. Will be passed along in every
@@ -165,13 +170,13 @@ class MesosBatchSystem(BatchSystemLocalSupport, AbstractScalableBatchSystem, Sch
 
         self._startDriver(config)
 
-    def setUserScript(self, userScript):
+    def setUserScript(self, userScript: str) -> None:
         self.userScript = userScript
 
-    def ignoreNode(self, nodeAddress):
+    def ignoreNode(self, nodeAddress: str) -> None:
         self.ignoredNodes.add(nodeAddress)
 
-    def unignoreNode(self, nodeAddress):
+    def unignoreNode(self, nodeAddress: str) -> None:
         self.ignoredNodes.remove(nodeAddress)
 
     def issueBatchJob(
@@ -179,7 +184,7 @@ class MesosBatchSystem(BatchSystemLocalSupport, AbstractScalableBatchSystem, Sch
         command: str,
         jobNode: JobDescription,
         job_environment: Optional[dict[str, str]] = None,
-    ):
+    ) -> str:
         """
         Issues the following command returning a unique jobID. Command is the string to run, memory
         is an int giving the number of bytes the job needs to run in and cores is the number of cpus

toil/batchSystems/mesos/test/__init__.py

@@ -7,6 +7,7 @@ from abc import ABCMeta, abstractmethod
 from contextlib import closing
 from shutil import which
 from urllib.request import urlopen
+from typing import Optional
 
 from toil.lib.retry import retry
 from toil.lib.threading import ExceptionalThread, cpu_count
@@ -25,7 +26,7 @@ class MesosTestSupport:
         with closing(urlopen("http://127.0.0.1:5050/version")) as content:
             content.read()
 
-    def _startMesos(self, numCores=None):
+    def _startMesos(self, numCores: Optional[int] = None) -> None:
         if numCores is None:
             numCores = cpu_count()
         shutil.rmtree("/tmp/mesos", ignore_errors=True)
@@ -52,7 +53,7 @@ class MesosTestSupport:
             log.warning("Forcibly killing child which ignored SIGTERM")
             process.kill()
 
-    def _stopMesos(self):
+    def _stopMesos(self) -> None:
         self._stopProcess(self.agent.popen)
         self.agent.join()
         self._stopProcess(self.master.popen)

toil/batchSystems/registry.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
 import logging
 import pkgutil
 import warnings
@@ -21,6 +20,7 @@ from typing import TYPE_CHECKING, Callable
 
 from toil.lib.compatibility import deprecated
 from toil.lib.memoize import memoize
+import toil.lib.plugins
 
 if TYPE_CHECKING:
     from toil.batchSystems.abstractBatchSystem import AbstractBatchSystem
@@ -40,17 +40,14 @@ def add_batch_system_factory(
 
     :param class_factory: A function that returns a batch system class (NOT an instance), which implements :class:`toil.batchSystems.abstractBatchSystem.AbstractBatchSystem`.
     """
-    _registry_keys.append(key)
-    _registry[key] = class_factory
+    toil.lib.plugins.register_plugin("batch_system", key, class_factory)
 
 
 def get_batch_systems() -> Sequence[str]:
     """
-    Get the names of all the availsble batch systems.
+    Get the names of all the available batch systems.
     """
-    _load_all_plugins()
-
-    return _registry_keys
+    return toil.lib.plugins.get_plugin_names("batch_system")
 
 
 def get_batch_system(key: str) -> type["AbstractBatchSystem"]:
@@ -60,8 +57,7 @@ def get_batch_system(key: str) -> type["AbstractBatchSystem"]:
     :raises: KeyError if the key is not the name of a batch system, and
              ImportError if the batch system's class cannot be loaded.
     """
-
-    return _registry[key]()
+    return toil.lib.plugins.get_plugin("batch_system", key)()
 
 
 DEFAULT_BATCH_SYSTEM = "single_machine"
@@ -126,114 +122,15 @@ def kubernetes_batch_system_factory():
 
 
 #####
-# Registry implementation
-#####
-
-_registry: dict[str, Callable[[], type["AbstractBatchSystem"]]] = {
-    "aws_batch": aws_batch_batch_system_factory,
-    "single_machine": single_machine_batch_system_factory,
-    "grid_engine": gridengine_batch_system_factory,
-    "lsf": lsf_batch_system_factory,
-    "mesos": mesos_batch_system_factory,
-    "slurm": slurm_batch_system_factory,
-    "torque": torque_batch_system_factory,
-    "htcondor": htcondor_batch_system_factory,
-    "kubernetes": kubernetes_batch_system_factory,
-}
-_registry_keys = list(_registry.keys())
-
-# We will load any packages starting with this prefix and let them call
-# add_batch_system_factory()
-_PLUGIN_NAME_PREFIX = "toil_batch_system_"
-
-
-@memoize
-def _load_all_plugins() -> None:
-    """
-    Load all the batch system plugins that are installed.
-    """
-
-    for finder, name, is_pkg in pkgutil.iter_modules():
-        # For all installed packages
-        if name.startswith(_PLUGIN_NAME_PREFIX):
-            # If it is a Toil batch system plugin, import it
-            importlib.import_module(name)
-
-
-#####
-# Deprecated API
+# Registers all built-in batch system
 #####
 
-# We used to directly access these constants, but now the Right Way to use this
-# module is add_batch_system_factory() to register and get_batch_systems() to
-# get the list/get_batch_system() to get a class by name.
-
-
-def __getattr__(name):
-    """
-    Implement a fallback attribute getter to handle deprecated constants.
-
-    See <https://stackoverflow.com/a/48242860>.
-    """
-    if name == "BATCH_SYSTEM_FACTORY_REGISTRY":
-        warnings.warn(
-            "BATCH_SYSTEM_FACTORY_REGISTRY is deprecated; use get_batch_system() or add_batch_system_factory()",
-            DeprecationWarning,
-        )
-        return _registry
-    elif name == "BATCH_SYSTEMS":
-        warnings.warn(
-            "BATCH_SYSTEMS is deprecated; use get_batch_systems()", DeprecationWarning
-        )
-        return _registry_keys
-    else:
-        raise AttributeError(f"Module {__name__} ahs no attribute {name}")
-
-
-@deprecated(new_function_name="add_batch_system_factory")
-def addBatchSystemFactory(
-    key: str, batchSystemFactory: Callable[[], type["AbstractBatchSystem"]]
-):
-    """
-    Deprecated method to add a batch system.
-    """
-    return add_batch_system_factory(key, batchSystemFactory)
-
-
-#####
-# Testing utilities
-#####
-
-# We need a snapshot save/restore system for testing. We can't just tamper with
-# the globals because module-level globals are their own references, so we
-# can't touch this module's global name bindings from a client module.
-
-
-def save_batch_system_plugin_state() -> (
-    tuple[list[str], dict[str, Callable[[], type["AbstractBatchSystem"]]]]
-):
-    """
-    Return a snapshot of the plugin registry that can be restored to remove
-    added plugins. Useful for testing the plugin system in-process with other
-    tests.
-    """
-
-    snapshot = (list(_registry_keys), dict(_registry))
-    return snapshot
-
-
-def restore_batch_system_plugin_state(
-    snapshot: tuple[list[str], dict[str, Callable[[], type["AbstractBatchSystem"]]]]
-):
-    """
-    Restore the batch system registry state to a snapshot from
-    save_batch_system_plugin_state().
-    """
-
-    # We need to apply the snapshot without rebinding the names, because that
-    # won't affect modules that imported the names.
-    wanted_batch_systems, wanted_registry = snapshot
-    _registry_keys.clear()
-    _registry_keys.extend(wanted_batch_systems)
-    _registry.clear()
-    _registry.update(wanted_registry)
+add_batch_system_factory("aws_batch", aws_batch_batch_system_factory)
+add_batch_system_factory("single_machine", single_machine_batch_system_factory)
+add_batch_system_factory("grid_engine", gridengine_batch_system_factory)
+add_batch_system_factory("lsf", lsf_batch_system_factory)
+add_batch_system_factory("mesos", mesos_batch_system_factory)
+add_batch_system_factory("slurm", slurm_batch_system_factory)
+add_batch_system_factory("torque", torque_batch_system_factory)
+add_batch_system_factory("htcondor", htcondor_batch_system_factory)
+add_batch_system_factory("kubernetes", kubernetes_batch_system_factory)

toil/batchSystems/singleMachine.py

@@ -94,7 +94,7 @@ class SingleMachineBatchSystem(BatchSystemSupport):
         self,
         config: Config,
         maxCores: float,
-        maxMemory: int,
+        maxMemory: float,
         maxDisk: int,
         max_jobs: Optional[int] = None,
     ) -> None:

toil/batchSystems/slurm.py

@@ -101,6 +101,32 @@ def parse_slurm_time(slurm_time: str) -> int:
             total_seconds += multiplier * int(elapsed_split[index])
     return total_seconds
 
+# For parsing user-provided option overrides (or self-generated
+# options) for sbatch, we need a way to recognize long, long-with-equals, and
+# short forms.
+def option_detector(long: str, short: str | None = None) -> Callable[[str], bool]:
+    """
+    Get a function that returns true if it sees the long or short
+    option.
+    """
+    def is_match(option: str) -> bool:
+        return option == f"--{long}" or option.startswith(f"--{long}=") or (short is not None and option == f"-{short}")
+    return is_match
+
+def any_option_detector(options: list[str | tuple[str, str]]) -> Callable[[str], bool]:
+    """
+    Get a function that returns true if it sees any of the long
+    options or long or short option pairs.
+    """
+    detectors = [option_detector(o) if isinstance(o, str) else option_detector(*o) for o in options]
+    def is_match(option: str) -> bool:
+        for detector in detectors:
+            if detector(option):
+                return True
+        return False
+    return is_match
+
+
 
 class SlurmBatchSystem(AbstractGridEngineBatchSystem):
     class PartitionInfo(NamedTuple):
@@ -647,31 +673,6 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             # Also any extra arguments from --slurmArgs or TOIL_SLURM_ARGS
             nativeConfig: str = self.boss.config.slurm_args  # type: ignore[attr-defined]
 
-            # For parsing user-provided option overrides (or self-generated
-            # options) we need a way to recognize long, long-with-equals, and
-            # short forms.
-            def option_detector(long: str, short: str | None = None) -> Callable[[str], bool]:
-                """
-                Get a function that returns true if it sees the long or short
-                option.
-                """
-                def is_match(option: str) -> bool:
-                    return option == f"--{long}" or option.startswith(f"--{long}=") or (short is not None and option == f"-{short}")
-                return is_match
-
-            def any_option_detector(options: list[str | tuple[str, str]]) -> Callable[[str], bool]:
-                """
-                Get a function that returns true if it sees any of the long
-                options or long or short option pairs.
-                """
-                detectors = [option_detector(o) if isinstance(o, str) else option_detector(*o) for o in options]
-                def is_match(option: str) -> bool:
-                    for detector in detectors:
-                        if detector(option):
-                            return True
-                    return False
-                return is_match
-
             is_any_mem_option = any_option_detector(["mem", "mem-per-cpu", "mem-per-gpu"])
             is_any_cpus_option = any_option_detector([("cpus-per-task", "c"), "cpus-per-gpu"])
             is_export_option = option_detector("export")
@@ -842,7 +843,7 @@ class SlurmBatchSystem(AbstractGridEngineBatchSystem):
             return sbatch_line
 
     def __init__(
-        self, config: Config, maxCores: float, maxMemory: int, maxDisk: int
+        self, config: Config, maxCores: float, maxMemory: float, maxDisk: float
     ) -> None:
         super().__init__(config, maxCores, maxMemory, maxDisk)
         self.partitions = SlurmBatchSystem.PartitionSet()

toil/bus.py

@@ -69,13 +69,15 @@ import tempfile
 import threading
 from collections.abc import Iterator
 from dataclasses import dataclass
-from typing import IO, Any, Callable, NamedTuple, Optional, TypeVar, cast
+from typing import IO, Any, Callable, NamedTuple, Optional, TypeVar, TYPE_CHECKING, cast
 
 from pubsub.core import Publisher
 from pubsub.core.listener import Listener
 from pubsub.core.topicobj import Topic
 from pubsub.core.topicutils import ALL_TOPICS
 
+from toil.lib.misc import FileDescriptorOrPath
+
 logger = logging.getLogger(__name__)
 
 # We define some ways to talk about jobs.
@@ -434,7 +436,7 @@ class MessageBus:
         connection._set_bus(self)
         return connection
 
-    def connect_output_file(self, file_path: str) -> Any:
+    def connect_output_file(self, file_path: FileDescriptorOrPath) -> Any:
         """
         Send copies of all messages to the given output file.
 
@@ -736,7 +738,7 @@ class JobStatus:
         )  # if the exit code is -1 and the job id is specified, we assume the job is running
 
 
-def replay_message_bus(path: str) -> dict[str, JobStatus]:
+def replay_message_bus(path: FileDescriptorOrPath) -> dict[str, JobStatus]:
     """
     Replay all the messages and work out what they mean for jobs.