Source code for gain.task_graph.logging

import hashlib
import os
import re
from typing import Any

from gain import logging
from gain.utils import fs_utils
from gain.utils.verbosity_configuration import VerbosityConfiguration




class FsspecHandler(logging.StreamHandler[Any]):
    """Class to create fsspec based logging handler."""

    def __init__(self, logfile: str):
        self.stream: Any = None
        fs, logpath = fs_utils.url_to_fs(logfile)
        stream = fs.open(logpath, "w")
        super().__init__(stream=stream)



    def close(self) -> None:
        """Close the stream.

        Copied from logging.FileHandler.close().
        """
        self.acquire()
        try:
            try:
                if self.stream:
                    try:
                        self.flush()
                    finally:
                        stream = self.stream
                        self.stream = None
                        stream.close()
            finally:
                # Issue #19523: call unconditionally to
                # prevent a handler leak when delay is set
                # Also see Issue #42378: we also rely on
                # self._closed being set to True there
                logging.StreamHandler.close(self)
        finally:
            self.release()






def ensure_log_dir(**kwargs: Any) -> str:
    """Ensure logging directory exists."""
    log_dir = kwargs.get("task_log_dir")
    if log_dir is None:
        log_dir = os.path.join(os.getcwd(), ".task-log")

    log_dir = fs_utils.abspath(log_dir)
    fs, path = fs_utils.url_to_fs(log_dir)
    fs.makedirs(path, exist_ok=True)
    return log_dir





def configure_task_logging(
    log_dir: str | None, task_id: str, verbosity: int,
) -> logging.Handler:
    """Configure and return task logging hadnler."""
    if log_dir is None:
        return logging.NullHandler()
    loglevel = VerbosityConfiguration.verbosity(verbosity)

    logfile = fs_utils.join(log_dir, f"log_{task_id}.log")

    handler = FsspecHandler(logfile)
    formatter = logging.Formatter(
        f"{task_id}: %(asctime)s %(name)s %(levelname)s %(message)s")
    handler.setFormatter(formatter)
    handler.setLevel(loglevel)

    return handler



_RE_TASK_ID = re.compile(r"[\. /,()\-:;]")




def safe_task_id(task_id: str) -> str:
    """Sanitize a task id into a string usable as a file name.

    Must be a pure function of ``task_id``: the forked executor computes
    the ``.result`` file name in the child and recomputes it in the parent,
    so a non-deterministic result makes the parent read a path that does
    not exist and silently yield ``None`` (gain#573).
    """
    result = _RE_TASK_ID.sub("_", task_id)
    # Coupled to the 200 char truncation in ``TaskGraph.make_task``, which
    # runs before a Task is ever built. Sanitization above is length
    # preserving, so ids coming from there always return here -- the branch
    # below is defence in depth for callers that bypass ``make_task``.
    # Raising or removing make_task's limit arms it.
    if len(result) <= 200:
        return result
    # Digest the RAW id, never the truncated one: ids derived from GRR
    # resource paths share long prefixes and differ only in their tail,
    # so a digest of ``result[:150]`` would collide for exactly those.
    # 150 + 1 + 40 == 191 keeps the result under the 200 char budget.
    digest = hashlib.sha256(task_id.encode()).hexdigest()[:40]
    return f"{result[:150]}_{digest}"