basebandflow 0.0.1__py3-none-any.whl

basebandflow/api/auth/__init__.py

@@ -0,0 +1,58 @@
+from datetime import timedelta
+
+from basebandflow.api.auth.canfar import is_logged_in
+from basebandflow.api.auth.tokens import (
+    load_tokens,
+    refresh_tokens,
+    token_time_remaining,
+    tokens_need_refresh,
+)
+from basebandflow.api.errors import InvalidTokens, NotLoggedIn, TokenRefreshFailed
+from basebandflow.api.settings import require_settings
+from basebandflow.models import TokenStatus
+from basebandflow.tools import require_tool
+
+
+def tokens_valid() -> bool:
+    settings = require_settings()
+    margin = timedelta(minutes=settings.token_refresh_timer)
+    access, _ = load_tokens()
+    return not tokens_need_refresh(access, margin)
+
+
+def refresh_and_verify() -> None:
+    """Refresh (prompts for a password), then confirm the result. Call from the CLI, not from launch()."""
+    refresh_tokens()
+    if not tokens_valid():
+        raise TokenRefreshFailed("Tokens are still expired after refresh.")
+
+
+def ensure_ready_for_launch() -> None:
+    require_tool("workflow")
+    logged_in = is_logged_in()
+    if logged_in is None:
+        raise NotLoggedIn(
+            "Couldn't reach CANFAR to verify your login. Check your connection and retry."
+        )
+    if not logged_in:
+        raise NotLoggedIn("Not logged in to CANFAR. Run `bbflow login`.")
+    if not tokens_valid():
+        raise InvalidTokens(
+            "CHIME/FRB tokens are expired or about to expire."
+        )  # no refresh here
+
+
+def _token_status(token: str | None, margin: timedelta) -> TokenStatus:
+    return TokenStatus(
+        present=bool(token),
+        remaining=token_time_remaining(token),
+        needs_refresh=tokens_need_refresh(token, margin),
+    )
+
+
+def check_tokens() -> tuple[TokenStatus, TokenStatus]:
+    settings = require_settings()
+    access, refresh = load_tokens()
+    return _token_status(
+        access, timedelta(minutes=settings.token_refresh_timer)
+    ), _token_status(refresh, timedelta(minutes=settings.token_refresh_timer))

basebandflow/api/auth/canfar.py

@@ -0,0 +1,24 @@
+from basebandflow import tools
+from basebandflow.api.errors import ToolNotFound
+
+
+def is_logged_in() -> bool | None:
+    """True/False, or None if it couldn't be verified (timeout)."""
+    try:
+        r = tools.canfar(
+            "ps",
+            stdin=tools.subprocess.DEVNULL,
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+    except ToolNotFound:
+        return False
+    except tools.subprocess.TimeoutExpired:
+        return None
+    return r.returncode == 0
+
+
+def login_interactive() -> int:  # inherits terminal, returns exit code
+    r = tools.canfar("login", "--force")
+    return r.returncode

basebandflow/api/auth/tokens.py

@@ -0,0 +1,102 @@
+import json
+import os
+import re
+import sys
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+
+import jwt
+from platformdirs import user_state_dir
+
+from basebandflow import tools
+from basebandflow.api.errors import TokenRefreshFailed
+from basebandflow.api.settings import require_settings
+
+TOKEN_FILE = Path(user_state_dir("bbflow", appauthor=False)) / "chime_frb_tokens.json"
+
+
+def token_time_remaining(token: str | None) -> timedelta | None:
+    """Time until the token expires (negative if already expired).
+    Returns None if the token is missing, malformed, or has no exp claim."""
+    if not token:
+        return None
+    try:
+        claims = jwt.decode(token, options={"verify_signature": False})
+        exp = claims.get("exp")
+        if exp is None:
+            return None
+        return datetime.fromtimestamp(exp, tz=timezone.utc) - datetime.now(timezone.utc)
+    except (jwt.DecodeError, ValueError, TypeError, OverflowError):
+        return None
+
+
+def tokens_need_refresh(token: str | None, margin: timedelta) -> bool:
+    remaining = token_time_remaining(token)
+    return remaining is None or remaining <= margin
+
+
+def parse_exports(output: str) -> tuple[str, str]:
+    """Parses the output from frb-api to find new tokens."""
+    a = re.search(r"CHIME_FRB_ACCESS_TOKEN=(\S+)", output)
+    r = re.search(r"CHIME_FRB_REFRESH_TOKEN=(\S+)", output)
+    if not (a and r):
+        raise TokenRefreshFailed("Couldn't find tokens in refresh output.")
+    return a.group(1).strip("'\""), r.group(1).strip("'\"")
+
+
+def save_tokens(access: str, refresh: str) -> None:
+    """Save the tokens to local json file."""
+    TOKEN_FILE.parent.mkdir(parents=True, exist_ok=True)
+    tmp = TOKEN_FILE.with_suffix(".tmp")
+    fd = os.open(
+        tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600
+    )  # owner read/write only
+    with os.fdopen(fd, "w") as f:
+        json.dump({"access": access, "refresh": refresh}, f)
+    tmp.replace(TOKEN_FILE)
+
+
+def load_tokens() -> tuple[str | None, str | None]:
+    """Loads the tokens either from the environment (first check), or the stored json file."""
+    # environment wins, so a user who exports them manually still works
+    access = os.environ.get("CHIME_FRB_ACCESS_TOKEN")
+    refresh = os.environ.get("CHIME_FRB_REFRESH_TOKEN")
+    if access and refresh:
+        return access, refresh
+    try:
+        data = json.loads(TOKEN_FILE.read_text())
+        return data.get("access"), data.get("refresh")
+    except (OSError, ValueError):
+        return None, None
+
+
+def refresh_tokens() -> None:
+    # require_tool("frb-api")
+    settings = require_settings()
+    if not sys.stdin.isatty():
+        raise TokenRefreshFailed(
+            "Token refresh needs a password; run `bbflow login` in a terminal."
+        )
+    try:
+        r = tools.frb_api(
+            "generate-tokens",
+            "-u",
+            settings.username,
+            stdout=tools.subprocess.PIPE,
+            text=True,
+            timeout=300,
+        )
+    except tools.subprocess.TimeoutExpired:
+        raise TokenRefreshFailed("Token refresh timed out.") from None
+    if r.returncode != 0:
+        raise TokenRefreshFailed("Token refresh failed (wrong password?).")
+
+    access, refresh = parse_exports(r.stdout)
+    save_tokens(access, refresh)
+    _export_to_env(access, refresh)
+
+
+def _export_to_env(access: str, refresh: str) -> None:
+    """Export the tokens to the environment"""
+    os.environ["CHIME_FRB_ACCESS_TOKEN"] = access
+    os.environ["CHIME_FRB_REFRESH_TOKEN"] = refresh

basebandflow/api/config.py

@@ -0,0 +1,59 @@
+from importlib.resources import files
+from pathlib import Path
+
+import yaml
+from jinja2 import Environment, StrictUndefined
+
+from basebandflow.api.errors import (
+    ConfigRenderError,
+    TemplateError,
+)
+from basebandflow.api.validate import validate_event_id, validate_username
+from basebandflow.models import RenderedConfig
+
+MATRIX_EXPR = "${{ matrix.job_id }}"
+
+
+def load_template() -> str:  # importlib.resources
+    return (files("basebandflow") / "templates" / "pipeline_config.yaml.j2").read_text()
+
+
+def render_config(username: str, event_id: str, directory: Path) -> RenderedConfig:
+    username = validate_username(username)
+    event_id = validate_event_id(event_id)
+
+    env = Environment(undefined=StrictUndefined, keep_trailing_newline=True)
+    try:
+        text = env.from_string(load_template()).render(
+            username=username, event_id=event_id
+        )
+        parsed = yaml.safe_load(text)
+    except (TemplateError, yaml.YAMLError) as e:
+        raise ConfigRenderError(f"Could not render config: {e}") from e
+
+    _check_rendered(parsed, event_id)
+
+    path = directory / "config.yaml"
+    path.write_text(
+        text, encoding="utf-8"
+    )  # exact rendered text: anchors and comments intact
+    return RenderedConfig(text=text, parsed=parsed, path=path, event_id=event_id)
+
+
+def _check_rendered(parsed, event_id: str) -> None:
+    """Sanity checks on the rendered result. Raises ConfigRenderError."""
+    if not isinstance(parsed, dict):
+        raise ConfigRenderError("Rendered config is not a YAML mapping.")
+    if (
+        str(parsed.get("event_number")) != event_id
+        or parsed.get("event_str") != event_id
+    ):
+        raise ConfigRenderError("Event ID did not render correctly.")
+
+    for step in parsed.get("pipeline", {}).get("steps", []):
+        if "matrix" in step:  # workflow expression must survive rendering untouched
+            job_id = step.get("work", {}).get("parameters", {}).get("job_id")
+            if job_id != MATRIX_EXPR:
+                raise ConfigRenderError(
+                    f"Step {step.get('name')!r}: matrix expression was altered."
+                )

basebandflow/api/containers.py

@@ -0,0 +1,199 @@
+from typing import Literal
+
+import nest_asyncio
+from canfar.sessions import Session
+
+from basebandflow.api.auth import ensure_ready_for_launch
+from basebandflow.api.auth.tokens import load_tokens
+from basebandflow.api.settings import require_settings
+from basebandflow.models import Container, Settings
+
+nest_asyncio.apply()
+s = Session()
+
+# _STATUS_RE = re.compile(r"^\s*Status\s{2,}(\S.*?)\s*$", re.MULTILINE)
+baseband_analysis_image = "images.canfar.net/chimefrb-public/baseband-analysis:latest"
+
+
+# Launch the datatrail containers
+def spawn_datatrail_container(settings: Settings) -> tuple[Container]:
+    payload = {
+        "name": "baseband-datatrail",  # name of the session
+        "image": baseband_analysis_image,  # image you're pulling #Try using chimefrb-public/baseband-analysis:latest if this image container does not work
+        "cores": 4,  # data staging now supports multithreaded pulling
+        "ram": 2,
+        "kind": "headless",
+        "cmd": "workflow",
+        "args": f"run --site=canfar baseband-{settings.username}-datatrail",  # the last args is the bucket name!!!!
+        "env": {"SITE": "canfar"},
+        "replicas": 2,
+    }
+    session_ids = s.create(**payload)
+    containers = []
+    for sid in session_ids:
+        container = Container(
+            payload.get("name"),
+            sid,
+            "pending",
+            "datatrail",
+        )
+        containers.append(container)
+    return containers
+
+
+def prune_datatrail_container() -> bool:
+    res = s.destroy_with(prefix="baseband-datatrail", status=None)
+    return bool(res) and all(res.values())
+
+
+def build_container_env(
+    settings: Settings, access_token: str, refresh_token: str
+) -> dict:
+    env = {
+        "SITE": "canfar",
+        "CHIME_FRB_ACCESS_TOKEN": access_token,
+        "CHIME_FRB_REFRESH_TOKEN": refresh_token,
+        # 'PYTHONPATH': '/arc/home/user/baseband-analysis/' # you can set this if you want to run your local branch on this headless session
+    }
+
+    if settings.baseband_path is not None:
+        env.update({"PYTHONPATH": settings.baseband_path})
+    return env
+
+
+# Launch the analysis containers
+def spawn_analysis_container(
+    settings: Settings, access_token: str, refresh_token: str
+) -> tuple[Container]:
+    env = build_container_env(settings, access_token, refresh_token)
+
+    payload = {
+        "name": "baseband-analysis",
+        "image": baseband_analysis_image,
+        "cores": 2,
+        "ram": 12,
+        "kind": "headless",
+        "cmd": "workflow",
+        "args": f"run --site=canfar baseband-{settings.username}",
+        "env": env,
+        "replicas": 10,
+    }
+
+    session_ids = s.create(**payload)
+    containers = []
+    for sid in session_ids:
+        container = Container(
+            payload.get("name"),
+            sid,
+            "pending",
+            "analysis",
+        )
+        containers.append(container)
+    return containers
+
+
+def prune_analysis_container() -> bool:
+    res = s.destroy_with(prefix="baseband-analysis", status=None)
+    return bool(res) and all(res.values())
+
+
+# Launch the refinement containers
+def spawn_refinement_container(
+    settings: Settings, access_token: str, refresh_token: str
+) -> tuple[Container]:
+    env = build_container_env(settings, access_token, refresh_token)
+
+    payload = {
+        "name": "baseband-refinement",
+        "image": baseband_analysis_image,
+        "cores": 8,
+        "ram": 32,
+        "kind": "headless",
+        "cmd": "workflow",
+        "args": f"run --log-level=INFO --site=canfar baseband-{settings.username}-brightest",
+        "env": env,
+        "replicas": 1,
+    }
+
+    session_ids = s.create(**payload)
+    containers = []
+    for sid in session_ids:
+        container = Container(
+            payload.get("name"),
+            sid,
+            "pending",
+            "refinement",
+        )
+        containers.append(container)
+    return containers
+
+
+def prune_refinement_container() -> bool:
+    res = s.destroy_with(prefix="baseband-refinement", status=None)
+    return bool(res) and all(res.values())
+
+
+# Check health
+def check_health(sid: list[str]) -> list[str] | None:
+    session_info = s.info(sid)
+    if session_info is None:
+        return
+
+    return [sinf.get("status") for sinf in session_info]
+
+
+def delete_container(ids: str | list[str]) -> bool:
+    res = s.destroy(ids)
+    return all(res)
+
+
+def prune_containers(session_type: Literal["datatrail", "analysis", "refinement"]):
+    if session_type == "datatrail":
+        return prune_datatrail_container()
+    if session_type == "analysis":
+        return prune_analysis_container()
+    if session_type == "refinement":
+        return prune_refinement_container()
+
+
+def initialize_session_container(session_type) -> list[Container]:
+    settings = require_settings()
+
+    current_containers = list(
+        filter(lambda x: x.type == session_type, list_headless_containers())
+    )
+    if session_type == "datatrail":
+        if 0 < len(current_containers) <= 2:
+            prune_containers(session_type)
+
+        return spawn_datatrail_container(settings)
+
+    ensure_ready_for_launch()
+    access, refresh = load_tokens()
+
+    if session_type == "analysis":
+        if 0 < len(current_containers) <= 10:
+            prune_containers(session_type)
+
+        return spawn_analysis_container(settings, access, refresh)
+
+    if session_type == "refinement":
+        if 0 < len(current_containers) <= 1:
+            prune_containers(session_type)
+
+        return spawn_refinement_container(settings, access, refresh)
+
+
+def list_headless_containers() -> list[Container]:
+    sessions = s.fetch(kind="headless")
+    containers = []
+    for session in sessions:
+        cont_type = session.get("name").split("-")[1]
+        container = Container(
+            session.get("name"), session.get("id"), session.get("status"), cont_type
+        )
+        containers.append(container)
+
+    containers.sort(key=lambda x: x.session_name)
+
+    return containers

basebandflow/api/errors.py

@@ -0,0 +1,50 @@
+class BBFlowError(Exception):
+    """Baseband Flow error."""
+
+
+class SettingsMissing(BBFlowError):
+    """Configuration file or required settings are missing."""
+
+
+class InvalidEventId(BBFlowError):
+    """Event ID must be a non-empty numeric value."""
+
+
+class InvalidUsername(BBFlowError):
+    """Username is invalid or not recognized."""
+
+
+class TemplateError(BBFlowError):
+    """Failed to load the pipeline template."""
+
+
+class ConfigRenderError(BBFlowError):
+    """Failed to render the pipeline configuration."""
+
+
+class NotLoggedIn(BBFlowError):
+    """You must log in before performing this action."""
+
+
+class TokenRefreshFailed(BBFlowError):
+    """Unable to refresh the authentication token."""
+
+
+class InvalidTokens(BBFlowError):
+    """Tokens are invalid. Refresh them."""
+
+
+class ToolNotFound(BBFlowError):
+    """Required command-line tool was not found on PATH"""
+
+
+class DeployFailed(BBFlowError):
+    """Pipeline deployment failed."""
+
+
+class SessionError(BBFlowError):
+    """The user session is invalid or could not be restored."""
+
+
+class SettingsCorrupt(SettingsMissing):
+    """Invalid configuration file. Reset it."""

basebandflow/api/history.py

@@ -0,0 +1,73 @@
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+
+import yaml
+from pydantic import ValidationError
+
+from basebandflow.api.settings import RUNS_DIR
+from basebandflow.models import LaunchResult
+
+
+@dataclass
+class RunRecord:
+    path: Path  # the run folder
+    event_id: str
+    started: (
+        datetime | None
+    )  # parsed from folder names; None if they don't match the format
+    result: LaunchResult | None  # None = launch_results.yaml missing or unreadable
+
+
+def load_run(path: Path) -> LaunchResult | None:
+    try:
+        data = yaml.safe_load(
+            (path / "launch_results.yaml").read_text(encoding="utf-8")
+        )
+        return LaunchResult(**data)
+    except (OSError, yaml.YAMLError, ValidationError, TypeError):
+        return None
+
+
+def _parse_run_dir(path: Path) -> tuple[str, datetime | None]:
+    """Folder layout: <YYYY-MM-DD>/<HHMMSS>_<event_id>"""
+    time_part, _, event_id = path.name.partition("_")
+    try:
+        started = datetime.strptime(
+            f"{path.parent.name} {time_part}", "%Y-%m-%d %H%M%S"
+        ).astimezone()
+    except ValueError:
+        started = None
+    return event_id, started
+
+
+def list_runs(
+    event_id: str | None = None,
+    *,
+    limit: int | None = None,
+    include_dry_runs: bool = False,
+) -> list[RunRecord]:
+    """Newest first. Dry runs (status 'rendered') are hidden unless include_dry_runs."""
+    if not RUNS_DIR.exists():
+        return []
+
+    records: list[RunRecord] = []
+    # date folders and HHMMSS_ names both sort chronologically as plain strings
+    for path in sorted((p for p in RUNS_DIR.glob("*/*") if p.is_dir()), reverse=True):
+        eid, started = _parse_run_dir(path)
+        if event_id and eid != event_id:
+            continue
+        result = load_run(path)
+        if result is not None and result.status == "rendered" and not include_dry_runs:
+            continue
+        records.append(
+            RunRecord(path=path, event_id=eid, started=started, result=result)
+        )
+        if limit and len(records) >= limit:
+            break
+    return records
+
+
+def latest_run(event_id: str) -> RunRecord | None:
+    runs = list_runs(event_id, limit=1)
+    return runs[0] if runs else None