pangea-flows 0.3.0__py3-none-any.whl

pangea_flows/__init__.py

@@ -0,0 +1,48 @@
+"""Write PANGEA processing workflows in Python.
+
+    from pangea_flows import each_file
+
+    @each_file(inputs="pollyxt/raw", pattern="*.nc", output="pollyxt/quicklooks")
+    def quicklook(path, out_dir):
+        plot(path, out_dir / f"{path.stem}.png")
+
+See README.md.
+"""
+
+__version__ = "0.3.0"
+
+from .client import PangeaClient
+from .errors import ClaimLost, FailPermanently, MissingSecrets, PangeaError
+from .keys import KeyScheme, TimeWindows, Window, daily, hourly, per_datum
+from .models import DataFile, Dataset, Datum
+from .runner import RunReport, connect, dry_run, run_once
+from .simple import each_file, each_window
+from .workflow import Context, Item, Output, Workflow, workflow
+
+__all__ = [
+    "ClaimLost",
+    "Context",
+    "DataFile",
+    "Dataset",
+    "Datum",
+    "FailPermanently",
+    "Item",
+    "KeyScheme",
+    "MissingSecrets",
+    "Output",
+    "PangeaClient",
+    "PangeaError",
+    "RunReport",
+    "TimeWindows",
+    "Window",
+    "Workflow",
+    "connect",
+    "daily",
+    "dry_run",
+    "each_file",
+    "each_window",
+    "hourly",
+    "per_datum",
+    "run_once",
+    "workflow",
+]

pangea_flows/capture.py

@@ -0,0 +1,144 @@
+"""Keeps the end of what a run, and each item, logs and prints, to send to the
+platform with the run's and the attempt's result.
+
+Researchers' code logs in all sorts of ways: ``ctx.log``, its own loggers,
+``print``, warnings. A handler on the root logger catches the loggers, and
+stdout and stderr are teed while an item runs. Output of programs started with
+``subprocess`` goes straight to the file descriptors and isn't caught; the
+runner that started the process keeps that.
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+import threading
+from contextlib import contextmanager
+from typing import Iterable, Iterator, Optional, TextIO
+
+#: The API keeps at most this many characters of a log
+MAX_LOG_CHARS = 64 * 1024
+
+#: Shorter secret values aren't masked: they would blank out whole logs
+MIN_MASKED_LENGTH = 6
+
+
+def mask(text: str, values: Iterable[str]) -> str:
+    """``text`` with each of ``values`` of at least MIN_MASKED_LENGTH characters replaced by ``***``"""
+    for value in sorted((v for v in values if len(v) >= MIN_MASKED_LENGTH), key=len, reverse=True):
+        text = text.replace(value, "***")
+    return text
+
+_FORMAT = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
+
+
+class Tail:
+    """Text that keeps only its last ``limit`` characters"""
+
+    def __init__(self, limit: int = MAX_LOG_CHARS):
+        self.limit = limit
+        self._parts: list[str] = []
+        self._size = 0
+        self._dropped = False
+        self._lock = threading.Lock()
+
+    def write(self, text: str) -> None:
+        if not text:
+            return
+        with self._lock:
+            self._parts.append(text)
+            self._size += len(text)
+            if self._size > 2 * self.limit:
+                joined = "".join(self._parts)[-self.limit :]
+                self._parts, self._size, self._dropped = [joined], len(joined), True
+
+    def text(self) -> str:
+        with self._lock:
+            joined = "".join(self._parts)
+            if len(joined) <= self.limit and not self._dropped:
+                return joined
+        marker = "[...]\n"
+        return marker + joined[-(self.limit - len(marker)) :]
+
+
+class Capture(logging.Handler):
+    """Copies log records into the run's tail, and into the current item's while one runs"""
+
+    def __init__(self, level: int = logging.INFO):
+        super().__init__(level)
+        self.setFormatter(_FORMAT)
+        self.run = Tail()
+        self.item: Optional[Tail] = None
+        #: Values to mask in everything captured
+        self.secrets: list[str] = []
+
+    def emit(self, record: logging.LogRecord) -> None:
+        # The handler sits on the root logger and on the run's logger, whose
+        # records may also reach the root: write each once
+        if getattr(record, "_pangea_captured", None) is self:
+            return
+        record._pangea_captured = self  # type: ignore[attr-defined]
+        try:
+            line = self.format(record) + "\n"
+        except Exception:
+            self.handleError(record)
+            return
+        self.write(line)
+
+    def write(self, text: str) -> None:
+        if self.secrets:
+            text = mask(text, self.secrets)
+        self.run.write(text)
+        item = self.item
+        if item is not None:
+            item.write(text)
+
+    @contextmanager
+    def attached(self, *loggers: logging.Logger) -> Iterator[Capture]:
+        """Catches records of the root logger and ``loggers`` meanwhile"""
+        targets = [logging.getLogger(), *(lg for lg in loggers if lg is not logging.getLogger())]
+        # The run's own messages are kept even where logging isn't set up, e.g.
+        # when run_once is called from a notebook, where loggers default to WARNING
+        levels = {lg: lg.level for lg in loggers if lg.getEffectiveLevel() > self.level}
+        for lg in levels:
+            lg.setLevel(self.level)
+        for target in targets:
+            target.addHandler(self)
+        try:
+            yield self
+        finally:
+            for target in targets:
+                target.removeHandler(self)
+            for lg, level in levels.items():
+                lg.setLevel(level)
+
+    @contextmanager
+    def item_log(self) -> Iterator[Tail]:
+        """Collects an item's log, including what it prints, meanwhile"""
+        tail = Tail()
+        self.item = tail
+        out, err = sys.stdout, sys.stderr
+        sys.stdout, sys.stderr = _Tee(out, self), _Tee(err, self)
+        try:
+            yield tail
+        finally:
+            sys.stdout, sys.stderr = out, err
+            self.item = None
+
+
+class _Tee:
+    """Writes to a stream and to the capture"""
+
+    def __init__(self, stream: TextIO, capture: Capture):
+        self._stream = stream
+        self._capture = capture
+
+    def write(self, text: str) -> int:
+        self._capture.write(text)
+        return self._stream.write(text)
+
+    def flush(self) -> None:
+        self._stream.flush()
+
+    def __getattr__(self, name: str):
+        return getattr(self._stream, name)

pangea_flows/cli.py

@@ -0,0 +1,240 @@
+"""``pangea-flows``: run workflows from the command line.
+
+    pangea-flows new my-flows                 # Start a project
+    pangea-flows list                         # The workflows in flows.py
+    pangea-flows try quicklook                # Run on yesterday's data, into ./out
+    pangea-flows run flows.py:quicklook       # Run for real, against the ledger
+
+A target is a workflow's name in ``flows.py``, ``module``, ``module:name``, or
+a file such as ``flows.py:name``; ``flows.py`` by default. The API comes from
+``$PANGEA_URL``, and tokens from ``$PANGEA_TOKEN_<NAME>`` or ``$PANGEA_TOKEN``,
+which can also be set in a ``.env`` file in the current directory.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import importlib.util
+import json
+import logging
+import os
+import sys
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from types import ModuleType
+from typing import Optional
+
+from . import __version__, template
+from .errors import MissingSecrets, PangeaError
+from .runner import connect, dry_run, run_once
+from .workflow import Workflow
+
+DEFAULT_FILE = "flows.py"
+
+
+def load(target: Optional[str]) -> list[Workflow]:
+    """The workflows a target names"""
+    target = target or DEFAULT_FILE
+    if target.isidentifier() and Path(DEFAULT_FILE).is_file():
+        # A bare name: a workflow in flows.py, if it has one by that name
+        if isinstance(getattr(_load_file(Path(DEFAULT_FILE)), target, None), Workflow):
+            target = f"{DEFAULT_FILE}:{target}"
+    module_name, _, attr = target.partition(":")
+    if module_name.endswith(".py"):
+        if not Path(module_name).is_file():
+            raise SystemExit(f"No file {module_name} here. Run this in your project's folder, or name a module.")
+        module = _load_file(Path(module_name))
+    else:
+        sys.path.insert(0, str(Path.cwd()))
+        module = importlib.import_module(module_name)
+
+    if attr:
+        wf = getattr(module, attr, None)
+        if not isinstance(wf, Workflow):
+            raise SystemExit(f"{target} is not a @workflow")
+        return [wf]
+    found = [v for v in vars(module).values() if isinstance(v, Workflow)]
+    if not found:
+        raise SystemExit(f"No workflows in {module_name}")
+    return found
+
+
+def describe(wf: Workflow) -> dict:
+    """What a runner needs to know about a workflow"""
+    return {
+        "name": wf.name,
+        "keys": wf.keys.name,
+        "inputs": wf.inputs,
+        "outputs": wf.outputs,
+        "first": wf.first,
+        "secrets": wf.secrets,
+        "schedule": wf.schedule,
+        "timezone": wf.timezone,
+        "memory": wf.memory,
+        "cpus": wf.cpus,
+        "timeoutSeconds": int(wf.timeout.total_seconds()) if wf.timeout else None,
+    }
+
+
+def _load_file(path: Path) -> ModuleType:
+    path = path.resolve()
+    name = path.stem
+    loaded = sys.modules.get(name)
+    if loaded is not None and getattr(loaded, "__file__", None) == str(path):
+        return loaded
+    spec = importlib.util.spec_from_file_location(name, path)
+    if spec is None or spec.loader is None:
+        raise SystemExit(f"Can't load {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    # Let flows.py import its neighbours, like a script would
+    sys.path.insert(0, str(path.parent))
+    spec.loader.exec_module(module)
+    return module
+
+
+def _one(target: Optional[str]) -> Workflow:
+    workflows = load(target)
+    if len(workflows) > 1:
+        names = ", ".join(w.name for w in workflows)
+        raise SystemExit(f"{target} has several workflows ({names}); name one as module:name")
+    return workflows[0]
+
+
+def load_dotenv(path: Path = Path(".env")) -> None:
+    """Sets variables from a ``.env`` file, without overriding ones already set"""
+    if not path.is_file():
+        return
+    for line in path.read_text().splitlines():
+        line = line.strip()
+        if not line or line.startswith("#") or "=" not in line:
+            continue
+        name, _, value = line.removeprefix("export ").partition("=")
+        value = value.strip()
+        if len(value) >= 2 and value[0] == value[-1] and value[0] in "'\"":
+            value = value[1:-1]
+        os.environ.setdefault(name.strip(), value)
+
+
+def _since(value: str) -> datetime:
+    parsed = datetime.fromisoformat(value.replace("Z", "+00:00"))
+    return parsed if parsed.tzinfo else parsed.replace(tzinfo=timezone.utc)
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    parser = argparse.ArgumentParser(prog="pangea-flows", description="Run PANGEA workflows")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Log debug messages")
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    new_cmd = commands.add_parser("new", help="Start a project of workflows in a new folder")
+    new_cmd.add_argument("folder", type=Path)
+
+    list_cmd = commands.add_parser("list", help="List the workflows in flows.py, or a module")
+    list_cmd.add_argument("target", nargs="?")
+    list_cmd.add_argument("--json", action="store_true", help="As JSON, for runners")
+
+    run_cmd = commands.add_parser("run", help="Plan new items and process them, once")
+    run_cmd.add_argument("target", nargs="?")
+    run_cmd.add_argument("--max-items", type=int, help="Process at most this many items")
+    run_cmd.add_argument("--no-plan", action="store_true", help="Only process items already in the ledger")
+    run_cmd.add_argument("--run-ref", help="An ID for this run, shown with each attempt")
+    run_cmd.add_argument("--trigger", default="manual", help="What started the run, shown on the platform")
+    run_cmd.add_argument("--run-id", type=int, help="Carry on with a run a runner started, instead of starting one")
+    run_cmd.add_argument(
+        "--accept-changes", action="store_true", help="Accept changed inputs or outputs since the last run"
+    )
+
+    serve_cmd = commands.add_parser("serve", help="Run workflows from their repositories, in containers")
+    serve_cmd.add_argument("--config", type=Path, required=True, help="The runner's TOML settings")
+    serve_cmd.add_argument("--once", action="store_true", help="Check once, wait for the runs, and stop")
+
+    dry_cmd = commands.add_parser(
+        "try",
+        aliases=["dry-run"],
+        help="Try a workflow on real data, writing outputs to ./out. Only reads from the API.",
+    )
+    dry_cmd.add_argument("target", nargs="?")
+    dry_cmd.add_argument(
+        "--since",
+        type=_since,
+        default=None,
+        help="Plan from files uploaded since this time (default: a day ago)",
+    )
+    dry_cmd.add_argument("--key", action="append", help="Run this item; can be repeated. By default the first few.")
+    dry_cmd.add_argument("--limit", type=int, default=5, help="Run at most this many items, without --key")
+    dry_cmd.add_argument("--out", type=Path, default=Path("out"), help="Where to write outputs (default: ./out)")
+
+    args = parser.parse_args(argv)
+    load_dotenv()
+    logging.basicConfig(
+        level=logging.DEBUG if args.verbose else logging.INFO,
+        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
+    )
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    try:
+        return _run(args)
+    except (PangeaError, LookupError, ValueError, MissingSecrets) as err:
+        print(f"pangea-flows: {err}", file=sys.stderr)
+        return 1
+
+
+def _run(args: argparse.Namespace) -> int:
+    if args.command == "new":
+        created = template.create(args.folder)
+        print(f"Created {args.folder}/ with {', '.join(p.name for p in created)}")
+        print(f"Next: cd {args.folder} && uv sync, then see README.md")
+        return 0
+
+    if args.command == "serve":
+        from .serve import Server
+        from .serve.config import Config
+
+        server = Server(Config.load(args.config))
+        if args.once:
+            server.api.report(f"pangea-flows {__version__}")
+            server.tick()
+            server.wait()
+        else:
+            server.serve_forever()
+        return 0
+
+    if args.command == "list":
+        if args.json:
+            print(json.dumps([describe(wf) for wf in load(args.target)]))
+            return 0
+        for wf in load(args.target):
+            print(f"{wf.name}\tkeys={wf.keys.name}\tinputs={wf.inputs}\toutputs={wf.outputs}\ttoken=${wf.token_variable}")
+        return 0
+
+    if args.command == "run":
+        report = run_once(
+            _one(args.target),
+            max_items=args.max_items,
+            plan=not args.no_plan,
+            run_ref=args.run_ref,
+            accept_changes=args.accept_changes,
+            trigger=args.trigger,
+            run_id=args.run_id,
+        )
+        print(report)
+        return 1 if report.failed else 0
+
+    if args.command in ("try", "dry-run"):
+        wf = _one(args.target)
+        since = args.since or datetime.now(timezone.utc) - timedelta(days=1)
+        with connect(wf) as client:
+            written = dry_run(wf, client, since=since, keys=args.key, limit=args.limit, out_dir=args.out)
+        if not written:
+            print(f"Nothing uploaded to the inputs since {since:%Y-%m-%d %H:%M} UTC. Try an earlier --since.")
+        for key, paths in written.items():
+            print(f"{key}: {len(paths)} file(s)")
+            for path in paths:
+                print(f"  {path}")
+        return 0
+
+    return 2
+
+
+if __name__ == "__main__":
+    sys.exit(main())