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

toil/fileStores/abstractFileStore.py

@@ -14,34 +14,31 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from collections.abc import Generator, Iterator
 from contextlib import contextmanager
 from tempfile import mkstemp
 from threading import Event, Semaphore
-from typing import (IO,
-                    TYPE_CHECKING,
-                    Any,
-                    Callable,
-                    ContextManager,
-                    Dict,
-                    Generator,
-                    Iterator,
-                    List,
-                    Literal,
-                    Optional,
-                    Set,
-                    Tuple,
-                    Type,
-                    Union,
-                    cast,
-                    overload)
+from typing import (
+    IO,
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    ContextManager,
+    Literal,
+    Optional,
+    Union,
+    cast,
+    overload,
+)
 
 import dill
 
-from toil.common import Toil, cacheDirName
+from toil.common import Toil, cacheDirName, getDirSizeRecursively
 from toil.fileStores import FileID
-from toil.job import Job, JobDescription
+from toil.job import DebugStoppingPointReached, Job, JobDescription
 from toil.jobStores.abstractJobStore import AbstractJobStore
 from toil.lib.compatibility import deprecated
+from toil.lib.conversions import bytes2human
 from toil.lib.io import WriteWatchingStream, mkdtemp
 
 logger = logging.getLogger(__name__)
@@ -75,9 +72,10 @@ class AbstractFileStore(ABC):
     Also responsible for committing completed jobs back to the job store with
     an update operation, and allowing that commit operation to be waited for.
     """
+
     # Variables used for syncing reads/writes
     _pendingFileWritesLock = Semaphore()
-    _pendingFileWrites: Set[str] = set()
+    _pendingFileWrites: set[str] = set()
     _terminateEvent = Event()  # Used to signify crashes in threads
 
     def __init__(
@@ -110,21 +108,28 @@ class AbstractFileStore(ABC):
         # This gets replaced with a subdirectory of itself on open()
         self.localTempDir: str = os.path.abspath(file_store_dir)
         assert self.jobStore.config.workflowID is not None
-        self.workflow_dir: str = Toil.getLocalWorkflowDir(self.jobStore.config.workflowID, self.jobStore.config.workDir)
-        self.coordination_dir: str =Toil.get_local_workflow_coordination_dir(self.jobStore.config.workflowID, self.jobStore.config.workDir, self.jobStore.config.coordination_dir)
-        self.jobName: str = (
-            self.jobDesc.command.split()[1] if self.jobDesc.command else ""
+        self.workflow_dir: str = Toil.getLocalWorkflowDir(
+            self.jobStore.config.workflowID, self.jobStore.config.workDir
+        )
+        self.coordination_dir: str = Toil.get_local_workflow_coordination_dir(
+            self.jobStore.config.workflowID,
+            self.jobStore.config.workDir,
+            self.jobStore.config.coordination_dir,
         )
+        self.jobName: str = str(self.jobDesc)
         self.waitForPreviousCommit = waitForPreviousCommit
-        self.loggingMessages: List[Dict[str, Union[int, str]]] = []
+        self.logging_messages: list[dict[str, Union[int, str]]] = []
+        self.logging_user_streams: list[dict[str, str]] = []
         # Records file IDs of files deleted during the current job. Doesn't get
         # committed back until the job is completely successful, because if the
         # job is re-run it will need to be able to re-delete these files.
         # This is a set of str objects, not FileIDs.
-        self.filesToDelete: Set[str] = set()
+        self.filesToDelete: set[str] = set()
         # Holds records of file ID, or file ID and local path, for reporting
         # the accessed files of failed jobs.
-        self._accessLog: List[Tuple[str, ...]] = []
+        self._accessLog: list[tuple[str, ...]] = []
+        # Holds total bytes of observed disk usage for the last job run under open()
+        self._job_disk_used: Optional[int] = None
 
     @staticmethod
     def createFileStore(
@@ -139,13 +144,17 @@ class AbstractFileStore(ABC):
         from toil.fileStores.cachingFileStore import CachingFileStore
         from toil.fileStores.nonCachingFileStore import NonCachingFileStore
 
-        fileStoreCls: Union[Type["CachingFileStore"], Type["NonCachingFileStore"]] = (
+        fileStoreCls: Union[type["CachingFileStore"], type["NonCachingFileStore"]] = (
             CachingFileStore if caching else NonCachingFileStore
         )
         return fileStoreCls(jobStore, jobDesc, file_store_dir, waitForPreviousCommit)
 
     @staticmethod
-    def shutdownFileStore(workflowID: str, config_work_dir: Optional[str], config_coordination_dir: Optional[str]) -> None:
+    def shutdownFileStore(
+        workflowID: str,
+        config_work_dir: Optional[str],
+        config_coordination_dir: Optional[str],
+    ) -> None:
         """
         Carry out any necessary filestore-specific cleanup.
 
@@ -165,7 +174,9 @@ class AbstractFileStore(ABC):
         from toil.fileStores.nonCachingFileStore import NonCachingFileStore
 
         workflowDir = Toil.getLocalWorkflowDir(workflowID, config_work_dir)
-        coordination_dir = Toil.get_local_workflow_coordination_dir(workflowID, config_work_dir, config_coordination_dir)
+        coordination_dir = Toil.get_local_workflow_coordination_dir(
+            workflowID, config_work_dir, config_coordination_dir
+        )
         cacheDir = os.path.join(workflowDir, cacheDirName(workflowID))
         if os.path.exists(cacheDir):
             # The presence of the cacheDir suggests this was a cached run. We don't need
@@ -187,15 +198,46 @@ class AbstractFileStore(ABC):
 
         :param job: The job instance of the toil job to run.
         """
-        failed = True
+        job_requested_disk = job.disk
         try:
             yield
             failed = False
-        finally:
-            # Do a finally instead of an except/raise because we don't want
-            # to appear as "another exception occurred" in the stack trace.
-            if failed:
+        except BaseException as e:
+            if isinstance(e, DebugStoppingPointReached):
+                self._dumpAccessLogs(job_type="Debugged", log_level=logging.INFO)
+            else:
                 self._dumpAccessLogs()
+            raise
+        finally:
+            # See how much disk space is used at the end of the job.
+            # Not a real peak disk usage, but close enough to be useful for warning the user.
+            self._job_disk_used = getDirSizeRecursively(self.localTempDir)
+
+            # Report disk usage
+            percent: float = 0.0
+            if job_requested_disk and job_requested_disk > 0:
+                percent = float(self._job_disk_used) / job_requested_disk * 100
+            disk_usage: str = (
+                f"Job {self.jobName} used {percent:.2f}% disk ({bytes2human(self._job_disk_used)}B [{self._job_disk_used}B] used, "
+                f"{bytes2human(job_requested_disk)}B [{job_requested_disk}B] requested)."
+            )
+            if self._job_disk_used > job_requested_disk:
+                self.log_to_leader(
+                    "Job used more disk than requested. For CWL, consider increasing the outdirMin "
+                    f"requirement, otherwise, consider increasing the disk requirement. {disk_usage}",
+                    level=logging.WARNING,
+                )
+            else:
+                self.log_to_leader(disk_usage, level=logging.DEBUG)
+
+    def get_disk_usage(self) -> Optional[int]:
+        """
+        Get the number of bytes of disk used by the last job run under open().
+
+        Disk usage is measured at the end of the job.
+        TODO: Sample periodically and record peak usage.
+        """
+        return self._job_disk_used
 
     # Functions related to temp files and directories
     def getLocalTempDir(self) -> str:
@@ -211,7 +253,9 @@ class AbstractFileStore(ABC):
         """
         return os.path.abspath(mkdtemp(dir=self.localTempDir))
 
-    def getLocalTempFile(self, suffix: Optional[str] = None, prefix: Optional[str] = None) -> str:
+    def getLocalTempFile(
+        self, suffix: Optional[str] = None, prefix: Optional[str] = None
+    ) -> str:
         """
         Get a new local temporary file that will persist for the duration of the job.
 
@@ -228,12 +272,14 @@ class AbstractFileStore(ABC):
         handle, tmpFile = mkstemp(
             suffix=".tmp" if suffix is None else suffix,
             prefix="tmp" if prefix is None else prefix,
-            dir=self.localTempDir
+            dir=self.localTempDir,
         )
         os.close(handle)
         return os.path.abspath(tmpFile)
 
-    def getLocalTempFileName(self, suffix: Optional[str] = None, prefix: Optional[str] = None) -> str:
+    def getLocalTempFileName(
+        self, suffix: Optional[str] = None, prefix: Optional[str] = None
+    ) -> str:
         """
         Get a valid name for a new local file. Don't actually create a file at the path.
 
@@ -287,7 +333,7 @@ class AbstractFileStore(ABC):
         basename: Optional[str] = None,
         encoding: Optional[str] = None,
         errors: Optional[str] = None,
-    ) -> Iterator[Tuple[WriteWatchingStream, FileID]]:
+    ) -> Iterator[tuple[WriteWatchingStream, FileID]]:
         """
         Similar to writeGlobalFile, but allows the writing of a stream to the job store.
         The yielded file handle does not need to and should not be closed explicitly.
@@ -327,18 +373,23 @@ class AbstractFileStore(ABC):
             def handle(numBytes: int) -> None:
                 # No scope problem here, because we don't assign to a fileID local
                 fileID.size += numBytes
+
             wrappedStream.onWrite(handle)
 
             yield wrappedStream, fileID
 
-    def _dumpAccessLogs(self) -> None:
+    def _dumpAccessLogs(
+        self, job_type: str = "Failed", log_level: int = logging.WARNING
+    ) -> None:
         """
-        When something goes wrong, log a report.
+        Log a report of the files accessed.
 
         Includes the files that were accessed while the file store was open.
+
+        :param job_type: Adjective to describe the job in the report.
         """
         if len(self._accessLog) > 0:
-            logger.warning('Failed job accessed files:')
+            logger.log(log_level, "%s job accessed files:", job_type)
 
             for item in self._accessLog:
                 # For each access record
@@ -347,14 +398,29 @@ class AbstractFileStore(ABC):
                     file_id, dest_path = item
                     if os.path.exists(dest_path):
                         if os.path.islink(dest_path):
-                            logger.warning('Symlinked file \'%s\' to path \'%s\'', file_id, dest_path)
+                            logger.log(
+                                log_level,
+                                "Symlinked file '%s' to path '%s'",
+                                file_id,
+                                dest_path,
+                            )
                         else:
-                            logger.warning('Downloaded file \'%s\' to path \'%s\'', file_id, dest_path)
+                            logger.log(
+                                log_level,
+                                "Downloaded file '%s' to path '%s'",
+                                file_id,
+                                dest_path,
+                            )
                     else:
-                        logger.warning('Downloaded file \'%s\' to path \'%s\' (gone!)', file_id, dest_path)
+                        logger.log(
+                            log_level,
+                            "Downloaded file '%s' to path '%s' (gone!)",
+                            file_id,
+                            dest_path,
+                        )
                 else:
                     # Otherwise dump without the name
-                    logger.warning('Streamed file \'%s\'', *item)
+                    logger.log(log_level, "Streamed file '%s'", *item)
 
     def logAccess(
         self, fileStoreID: Union[FileID, str], destination: Union[str, None] = None
@@ -421,14 +487,12 @@ class AbstractFileStore(ABC):
         fileStoreID: str,
         encoding: Literal[None] = None,
         errors: Optional[str] = None,
-    ) -> ContextManager[IO[bytes]]:
-        ...
+    ) -> ContextManager[IO[bytes]]: ...
 
     @overload
     def readGlobalFileStream(
         self, fileStoreID: str, encoding: str, errors: Optional[str] = None
-    ) -> ContextManager[IO[str]]:
-        ...
+    ) -> ContextManager[IO[str]]: ...
 
     @abstractmethod
     def readGlobalFileStream(
@@ -472,7 +536,7 @@ class AbstractFileStore(ABC):
         :return: File's size in bytes, as stored in the job store
         """
         # First try and see if the size is still attached
-        size = getattr(fileStoreID, 'size', None)
+        size = getattr(fileStoreID, "size", None)
 
         if size is None:
             # It fell off
@@ -525,7 +589,7 @@ class AbstractFileStore(ABC):
     ) -> Optional[FileID]:
         return self.jobStore.import_file(src_uri, shared_file_name=shared_file_name)
 
-    @deprecated(new_function_name='export_file')
+    @deprecated(new_function_name="export_file")
     def exportFile(self, jobStoreFileID: FileID, dstUrl: str) -> None:
         return self.export_file(jobStoreFileID, dstUrl)
 
@@ -554,7 +618,7 @@ class AbstractFileStore(ABC):
     class _StateFile:
         """Read and write dill-ed state dictionaries from/to a file into a namespace."""
 
-        def __init__(self, stateDict: Dict[str, Any]):
+        def __init__(self, stateDict: dict[str, Any]):
             assert isinstance(stateDict, dict)
             self.__dict__.update(stateDict)
 
@@ -582,7 +646,7 @@ class AbstractFileStore(ABC):
             """
             # Read the value from the cache state file then initialize and instance of
             # _CacheState with it.
-            with open(fileName, 'rb') as fH:
+            with open(fileName, "rb") as fH:
                 infoDict = dill.load(fH)
             return cls(infoDict)
 
@@ -592,14 +656,14 @@ class AbstractFileStore(ABC):
 
             :param fileName: Path to the state file.
             """
-            with open(fileName + '.tmp', 'wb') as fH:
+            with open(fileName + ".tmp", "wb") as fH:
                 # Based on answer by user "Mark" at:
                 # http://stackoverflow.com/questions/2709800/how-to-pickle-yourself
                 # We can't pickle nested classes. So we have to pickle the variables
                 # of the class.
                 # If we ever change this, we need to ensure it doesn't break FileID
                 dill.dump(self.__dict__, fH)
-            os.rename(fileName + '.tmp', fileName)
+            os.rename(fileName + ".tmp", fileName)
 
     # Functions related to logging
     def log_to_leader(self, text: str, level: int = logging.INFO) -> None:
@@ -611,13 +675,29 @@ class AbstractFileStore(ABC):
         :param level: The logging level.
         """
         logger.log(level=level, msg=("LOG-TO-MASTER: " + text))
-        self.loggingMessages.append(dict(text=text, level=level))
-
+        self.logging_messages.append(dict(text=text, level=level))
 
-    @deprecated(new_function_name='export_file')
+    @deprecated(new_function_name="export_file")
     def logToMaster(self, text: str, level: int = logging.INFO) -> None:
         self.log_to_leader(text, level)
-        
+
+    def log_user_stream(self, name: str, stream: IO[bytes]) -> None:
+        """
+        Send a stream of UTF-8 text to the leader as a named log stream.
+
+        Useful for things like the error logs of Docker containers. The leader
+        will show it to the user or organize it appropriately for user-level
+        log information.
+
+        :param name: A hierarchical, .-delimited string.
+        :param stream: A stream of encoded text. Encoding errors will be
+            tolerated.
+        """
+
+        # Read the whole stream into memory
+        steam_data = stream.read().decode("utf-8", errors="replace")
+        # And remember it for the worker to fish out
+        self.logging_user_streams.append(dict(name=name, text=steam_data))
 
     # Functions run after the completion of the job.
     @abstractmethod