pegasus-wms.worker 5.1.2__py3-none-any.whl

Pegasus/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

Pegasus/cli/__init__.py

@@ -0,0 +1 @@
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)

Pegasus/cli/pegasus-checkpoint.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python3
+import argparse
+import importlib
+import logging
+import logging.handlers
+import os
+import re
+import signal
+import sys
+import tarfile
+import threading
+import time
+from pathlib import Path
+from typing import Iterable, List, Set
+
+# PEGASUS_PYTHONPATH is set by the pegasus-python-wrapper script
+peg_path = os.environ.get("PEGASUS_PYTHONPATH")
+if peg_path:
+    for p in reversed(peg_path.split(":")):
+        if p not in sys.path:
+            sys.path.insert(0, p)
+
+ptransfer_module = importlib.import_module("Pegasus.cli.pegasus-transfer")
+pegasus_transfer = ptransfer_module.pegasus_transfer
+
+log = logging.getLogger("pegasus-checkpoint")
+
+# name of archived & compressed checkpoints
+CHECKPOINT_FILENAME = "pegasus.checkpoint.tar.gz"
+
+# pegasus-transfer url file expected to be generated by pegasus-lite
+PEGASUS_TRANSFER_URL_FILE = "pegasus_checkpoint_transfer_urls.json"
+
+# file to which this process's PID will be written upon startup
+PEGASUS_CHECKPOINT_PID_FILE = "pegasus_checkpoint.pid"
+
+
+def parse_args(args: List[str] = sys.argv[1:]) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Periodically transfer checkpoints back to staging site."
+    )
+
+    parser.add_argument(
+        "-p",
+        "--pattern",
+        action="append",
+        dest="patterns",
+        metavar="PATTERN",
+        help="""regex patterns to match when searching for files to checkpoint""",
+    )
+
+    def validate_interval(value: str) -> int:
+        """
+        Ensure that given interval is an integer and is nonzero.
+
+        :param value: the given interval to validate
+        :type value: str
+        :raises argparse.ArgumentTypeError: interval cannot be casted to an int or is a negative integer
+        :return: the given value
+        :rtype: int
+        """
+        is_valid = True
+        try:
+            value = int(value)
+        except ValueError:
+            is_valid = False
+
+        if not is_valid or value < 0:
+            raise argparse.ArgumentTypeError(
+                f"interval {value} must be a nonnegative integer"
+            )
+
+        return value
+
+    parser.add_argument(
+        "-i",
+        "--interval",
+        type=validate_interval,
+        metavar="SECONDS",
+        help="""interval in seconds at which pegasus-checkpoint will send
+                checkpoint file back to the staging site
+                """,
+    )
+
+    # TODO: add --once to indicate pegasus-checkpoint is to run CheckpointWorker.run once
+
+    parser.add_argument(
+        "-d", "--debug", action="store_true", default=False, help="enable debug logging"
+    )
+
+    parser.add_argument(
+        "-l",
+        "--log-to-file",
+        action="store_true",
+        default=False,
+        help="enable logging to file pegasus-checkpoint.log",
+    )
+
+    return parser.parse_args(args)
+
+
+def configure_logging(level: int, log_to_file: bool):
+    """
+    Setup logging.
+
+    :param level: log level to use; set to :code:`logging.DEBUG` to enable debug logging
+    :type level: int
+    :param log_to_file: whether or not to write logs to the file pegasus-checkpoint.log
+    :type log_to_file: bool
+    """
+    if level == logging.DEBUG:
+        level = logging.DEBUG
+    else:
+        level = logging.INFO
+
+    handlers = [logging.StreamHandler()]
+    if log_to_file:
+        handlers.append(
+            logging.handlers.RotatingFileHandler(
+                filename="pegasus-checkpoint.log",
+                mode="a",
+                maxBytes=(1 << 20),
+                backupCount=2,
+            )
+        )
+
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s %(threadName)s [%(levelname)s] %(message)s",
+        handlers=handlers,
+    )
+
+
+def write_pid():
+    """Write PID to :code:`PEGASUS_CHECKPOINT_PID_FILE`"""
+    with open(PEGASUS_CHECKPOINT_PID_FILE, "w") as f:
+        pid = os.getpid()
+        f.write(str(pid))
+
+        log.debug(f"started with PID: {pid}")
+
+
+class PeriodicCheckpointNotifier(threading.Thread):
+    """
+    Thread to periodically signal to the worker thread to start saving and
+    staging out checkpoint files.
+    """
+
+    log = logging.getLogger("pegasus-checkpoint.PeriodicCheckcpointNotifier")
+
+    def __init__(self, interval: int, notify: threading.Event):
+        """Constructor
+
+        :param interval: interval in seconds to sleep for
+        :type interval: int
+        :param notify: event to set, which will notify checkpoint worker thread to wake up and run 
+        :type notify: threading.Event
+        """
+        super().__init__(group=None, target=None, name="PeriodicCheckpointNotifier")
+        self.interval = interval
+        self.notify = notify
+
+    def run(self):
+        """
+        Repeatedly sleep for the given interval, then notify checkpoint worker
+        thread to start working upon waking up.
+        """
+        self.log.debug(f"{self.name} started with interval: {self.interval} seconds")
+        while True:
+            time.sleep(self.interval)
+            self.notify.set()
+            self.log.debug("notify event set")
+
+
+class CheckpointWorker(threading.Thread):
+    """Thread to handle saving and staging out checkpoint files."""
+
+    log = logging.getLogger("pegasus-checkpoint.CheckpointWorker")
+
+    def __init__(self, notify: threading.Event, patterns: List[str]):
+        """Constructor
+
+        :param notify: event to watch for, which will indicate that this thread needs to run
+        :type notify: threading.Event
+        :param patterns: regex file patterns to match
+        :type patterns: List[str]
+        """
+        super().__init__(group=None, target=None, name="CheckpointWorker")
+
+        self.notify = notify
+        self.patterns = patterns
+
+    def run(self):
+        """
+        Wait for notification to start working. Once notified, create the
+        checkpoint archive file and transfer back to staging site with either
+        pegasus-transfer or condor_chirp.
+        self.log.debug("{} started".format(self.name))
+        """
+        while True:
+            # wait till we are notified
+            self.notify.wait(timeout=None)
+            self.log.debug("got notification, starting work")
+
+            # work
+            matched_files = CheckpointWorker.get_matched_filenames(self.patterns)
+
+            CheckpointWorker.archive_and_compress(matched_files)
+
+            # TODO: add exception handling
+            pegasus_transfer(
+                max_attempts=3,
+                num_threads=8,
+                file=PEGASUS_TRANSFER_URL_FILE,
+                symlink=True,
+            )
+
+            # clear the notification, so we can wait until it is set again
+            self.notify.clear()
+            self.log.debug("work done, cleared notification")
+
+    @staticmethod
+    def get_matched_filenames(patterns: List[str]) -> Set[str]:
+        """
+        Get all files/folders in the current working directory that match the
+        given file patterns.
+
+        :param patterns: regex patterns to match filenames with
+        :type patterns: List[str]
+        :return: all files and or folders which match one of the given patterns
+        :rtype: Set[str]
+        """
+        patterns = [re.compile(p) for p in patterns]
+        files_in_cwd = [str(f) for f in Path(".").iterdir()]
+        matched = set()
+
+        # n^2.. maybe can do a little better by pulling out files from files_in_cwd (using a set
+        # isntead) as we make matches so we don't try to iterate over it again
+        for p in patterns:
+            for f in files_in_cwd:
+                is_match = bool(p.fullmatch(f))
+
+                if is_match:
+                    matched.add(f)
+
+                CheckpointWorker.log.debug(
+                    "pattern: {:<30}, file: {:>20}, match: {}".format(
+                        p.pattern, f, is_match
+                    )
+                )
+
+        CheckpointWorker.log.info(
+            f"given patterns matched the following filenames: {matched}"
+        )
+
+        return matched
+
+    @staticmethod
+    def archive_and_compress(filenames: Iterable[str]):
+        """
+        Given a list of file/folder names in CWD, archive and compress them into a single file
+        :code:`pegasus.checkpoint.tar.gz`
+
+        :param filenames: list of filenames to archive and compress 
+        :type filenames: Iterable[str]
+        """
+        start = time.perf_counter()
+        with tarfile.open(CHECKPOINT_FILENAME, "w|gz") as tar:
+            for f in filenames:
+                tar.add(name=f, recursive=True)
+        end = time.perf_counter()
+
+        checkpoint_size_in_mb = Path(CHECKPOINT_FILENAME).stat().st_size / (1 << 20)
+
+        CheckpointWorker.log.info(
+            "created {} MB checkpoint file: {}".format(
+                checkpoint_size_in_mb, CHECKPOINT_FILENAME
+            )
+        )
+        CheckpointWorker.log.info(f"archive and gzip took {end - start} seconds")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    log_level = logging.INFO
+    if args.debug:
+        log_level = logging.DEBUG
+
+    configure_logging(level=log_level, log_to_file=args.log_to_file)
+
+    write_pid()
+
+    # event to signal to worker thread to start saving checkpoints
+    notify = threading.Event()
+
+    # main thread to handle signal from user application
+    def SIGUSR1_handler(signum, frame):
+        log.info("SIGUSR1 received, setting notification")
+        notify.set()
+
+    signal.signal(signal.SIGUSR1, SIGUSR1_handler)
+
+    # create and start a periodic checkpoint notifier if user specified an interval to run on
+    periodic_notifier = None
+    if args.interval:
+        periodic_notifier = PeriodicCheckpointNotifier(
+            interval=args.interval, notify=notify
+        )
+        periodic_notifier.start()
+
+    worker = CheckpointWorker(notify=notify, patterns=args.patterns)
+    worker.start()
+
+    if periodic_notifier:
+        periodic_notifier.join()
+    worker.join()

Pegasus/cli/pegasus-globus-online-init.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+
+import os
+from argparse import ArgumentParser
+
+import globus_sdk
+from globus_sdk.scopes import (
+    GCSCollectionScopeBuilder,
+    GCSEndpointScopeBuilder,
+    TransferScopes,
+)
+
+try:
+    from configparser import ConfigParser
+except Exception:
+    from ConfigParser import ConfigParser
+
+client_id = "d7382f5a-4ea3-4b69-b094-99c392fc820d"
+config_file = os.path.expanduser("~/.pegasus/globus.conf")
+
+parser = ArgumentParser(
+    description="Initialize Globus OAuth Tokens - SDK Version {}".format(
+        globus_sdk.__version__
+    )
+)
+parser.add_argument(
+    "-p", "--permanent", action="store_true", help="request a refreshable token",
+)
+parser.add_argument(
+    "-e",
+    "--endpoints",
+    nargs="*",
+    default=[],
+    help="list of endpoint uuids to acquire manage_collections consent",
+)
+parser.add_argument(
+    "-c",
+    "--collections",
+    nargs="*",
+    default=[],
+    help="list of collection uuids to acquire data_access consent",
+)
+parser.add_argument(
+    "-d",
+    "--domains",
+    nargs="*",
+    default=[],
+    help="list of domain requirements which must be satisfied by the identities under the globus user account",
+)
+
+
+args = parser.parse_args()
+
+config = ConfigParser()
+config.add_section("oauth")
+config.set("oauth", "client_id", client_id)
+
+pegasus_scopes = [TransferScopes.all]
+for c in args.collections:
+    sb = GCSCollectionScopeBuilder(c, known_scopes=[TransferScopes.all])
+    pegasus_scopes.append(f"{TransferScopes.all}[*{sb.data_access}]")
+    pegasus_scopes.append(sb.data_access)
+
+for e in args.endpoints:
+    sb = GCSEndpointScopeBuilder(e)
+    pegasus_scopes.append(sb.manage_collections)
+
+required_domains = None
+for d in args.domains:
+    if not required_domains:
+        required_domains = [d]
+    else:
+        required_domains.append(d)
+
+client = globus_sdk.NativeAppAuthClient(client_id)
+if args.permanent:
+    client.oauth2_start_flow(refresh_tokens=True, requested_scopes=pegasus_scopes)
+else:
+    client.oauth2_start_flow(requested_scopes=pegasus_scopes)
+
+authorize_url = client.oauth2_get_authorize_url(
+    session_required_single_domain=required_domains, prompt="login"
+)
+print(f"Please go to this URL and login: {authorize_url}")
+
+get_input = getattr(__builtins__, "raw_input", input)
+auth_code = get_input("Please enter the code you get after login here: ").strip()
+token_response = client.oauth2_exchange_code_for_tokens(auth_code)
+
+transfer_data = token_response.by_resource_server["transfer.api.globus.org"]
+
+# get tokens as strings and their expiration timestamp
+transfer_at = str(transfer_data["access_token"])
+transfer_rt = str(transfer_data["refresh_token"])
+transfer_at_exp = str(transfer_data["expires_at_seconds"])
+
+config.set("oauth", "transfer_at", transfer_at)  # transfer access token
+config.set("oauth", "transfer_rt", transfer_rt)  # transfer refresh token
+config.set("oauth", "transfer_at_exp", transfer_at_exp)  # transfer expiration time
+
+cfg = open(config_file, "w+")
+config.write(cfg)
+cfg.close()