phaserEM 0.1__py3-none-any.whl

phaser/web/slurm.py

@@ -0,0 +1,180 @@
+import asyncio
+import logging
+import weakref
+import json
+import os
+import re
+import importlib.resources
+import typing as t
+
+from .server import WorkerID, Worker
+
+
+SlurmID: t.TypeAlias = int
+SlurmState: t.TypeAlias = t.Literal[
+    'BOOT_FAIL', 'CANCELLED', 'COMPLETED', 'CONFIGURING', 'COMPLETING',
+    'DEADLINE', 'FAILED', 'NODE_FAIL', 'OUT_OF_MEMORY', 'PENDING',
+    'PREEMPTED', 'RUNNING', 'RESV_DEL_HOLD', 'REQUEUE_FED', 'REQUEUE_HOLD',
+    'REQUEUED', 'RESIZING', 'REVOKED', 'SIGNALING', 'SPECIAL_EXIT',
+    'STAGE_OUT', 'STOPPED', 'SUSPENDED', 'TIMEOUT',
+]
+
+
+class SlurmJobInfo(t.TypedDict):
+    job_id: SlurmID
+    job_state: SlurmState
+    name: str
+    nodes: str
+    user_name: str
+
+    # path to stdout and stderr
+    standard_output: str
+    standard_error: str
+
+    submit_time: int
+    start_time: int
+    eligible_time: int
+    end_time: int
+
+
+class SqueueResult(t.TypedDict):
+    meta: t.Dict[str, t.Any]
+    jobs: t.List[SlurmJobInfo]
+    warnings: t.List[t.Any]
+    errors: t.List[t.Any]
+
+
+class SlurmWorker(Worker):
+    def __init__(self, worker_id: WorkerID, slurm_job_id: SlurmID):
+        self.slurm_job_id: SlurmID = slurm_job_id
+        super().__init__(worker_id)
+
+    async def cancel(self):
+        from .server import server
+        if self.status == 'queued':
+            await server.slurm_manager.cancel_queued_worker(self.slurm_job_id)
+
+        await self.set_status('stopping')
+
+
+class SlurmManager:
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+        self._slurm_exists: t.Optional[bool] = None
+        self._slurm_workers: weakref.WeakValueDictionary[SlurmID, SlurmWorker] = weakref.WeakValueDictionary()
+
+        self._poll_task: t.Optional[asyncio.Task[None]] = None
+        self._new_worker_event: asyncio.Event = asyncio.Event()
+
+    async def make_worker(self, worker_id: WorkerID, url: str) -> SlurmWorker:
+        await self.check_slurm_exists()
+
+        job_name = f"phaser_{worker_id}"
+
+        from .. import web
+        with importlib.resources.path(web, 'slurm_worker.sh') as script_path:
+            slurm_args = "--qos=high --time=4-0 --partition=xeon-g6-volta --cpus-per-task=20 --gres=gpu:volta:1 --signal=SIGINT@120"
+            #slurm_args = "--qos=high --time=01:00:00 --partition=debug-cpu --cpus-per-task=20"
+            shell = os.environ.get('SHELL', '/bin/bash')
+            proc = await asyncio.create_subprocess_exec(
+                shell, "-c", f"sbatch --job-name={job_name} {slurm_args} \"{script_path.absolute()}\" \"{url}\"",
+                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
+            )
+            stdout, stderr = await proc.communicate()
+
+            if proc.returncode != 0:
+                raise RuntimeError(f"Failed to start slurm worker, stderr: '{stderr}'")
+
+            try:
+                match = re.fullmatch(r'Submitted batch job (\d+)', stdout.decode().strip())
+                job_id: SlurmID = int(match[1])  # type: ignore
+            except (ValueError, TypeError):
+                raise ValueError(f"Failed to parse slurm return: '{stdout.decode()}'")
+
+            worker = SlurmWorker(worker_id, job_id)
+            self._slurm_workers[job_id] = worker
+
+            if self._poll_task is None:
+                self._poll_task = asyncio.create_task(self._poll_slurm())
+            else:
+                self._new_worker_event.set()
+
+            return worker
+
+    async def cancel_queued_worker(self, slurm_job_id: SlurmID):
+        proc = await asyncio.create_subprocess_shell(
+            f"scancel --state=PENDING {slurm_job_id}",
+            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await proc.communicate()
+
+        if proc.returncode != 0:
+            self.logger.warning(f"Failed to cancel slurm job: '{stderr.decode()}'")
+
+    async def _poll_slurm(self):
+        while True:
+            while True:
+                try:
+                    await self._poll_slurm_status()
+                except Exception as e:
+                    self.logger.warning(f"Failed to poll slurm worker statuses: {e}")
+
+                if any(worker.state == 'queued' for worker in self._slurm_workers.values()):
+                    # fast wait cycle
+                    await asyncio.sleep(5.0)
+                elif any(worker.state != 'stopped' for worker in self._slurm_workers.values()):
+                    # if there's any workers to wait for, slow wait cycle
+                    await asyncio.sleep(30.0)
+                else:
+                    # no workers waiting, wait until one is started
+                    self._new_worker_event.clear()
+                    break
+            await self._new_worker_event.wait()
+
+    async def _poll_slurm_status(self):
+        jobs = ",".join(str(worker.slurm_job_id) for worker in self._slurm_workers.values())
+
+        proc = await asyncio.create_subprocess_shell(
+            f"squeue --job={jobs} --states=all --json",
+            stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+        )
+        stdout, stderr = await proc.communicate()
+
+        if proc.returncode != 0:
+            self.logger.warning(f"Failed to poll slurm worker statuses, stderr: '{stderr}'")
+            return
+
+        d = t.cast(SqueueResult, json.loads(stdout))
+
+        for job in d["jobs"]:
+            try:
+                worker = self._slurm_workers[job["job_id"]]
+            except KeyError:
+                continue
+
+            if job['job_state'] in ('CONFIGURING', 'RUNNING'):
+                if worker.status == 'queued':
+                    await worker.set_status('starting')
+            elif job['job_state'] != 'PENDING':
+                if worker.status != 'stopped':
+                    # TODO grab some exit information here
+                    await worker.set_status('stopped')
+
+    async def check_slurm_exists(self):
+        if self._slurm_exists is None:
+            proc = await asyncio.create_subprocess_shell("sbatch --version", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
+            stdout, stderr = await proc.communicate()
+
+            if proc.returncode == 0:
+                self._slurm_exists = True
+                self.logger.info(f"Slurm found, version '{stdout.decode().strip()}'")
+            else:
+                self._slurm_exists = False
+                self.logger.warning(f"Slurm not found, stderr '{stderr.decode()}'")
+
+        if not self._slurm_exists:
+            raise RuntimeError("Slurm not found, is it installed and on PATH?")
+
+    async def finalize(self):
+        if self._poll_task is not None:
+            self._poll_task.cancel()

phaser/web/templates/base.html

@@ -0,0 +1,13 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+    <title>{% block title %}Phaser{% endblock %}</title>
+    <link href="{{ url_for('static', filename='style.css') }}" rel='stylesheet'>
+    {% block head %}{% endblock %}
+  </head>
+  <body>
+    {% block body %}{% endblock %}
+    <noscript>This page contains webassembly and javascript content, please enable javascript in your browser.</noscript>
+  </body>
+</html>

phaser/web/templates/dashboard.html

@@ -0,0 +1,14 @@
+{% extends "base.html" %}
+{% block title %}Phaser frontend{% endblock %}
+
+{% block head %}
+  <script type="text/javascript" src="{{ url_for('static', filename='bundle-dashboard.js') }}" defer></script>
+{% endblock %}
+
+{% block body %}
+  <div id="app"></div>
+  <!--<div class="section-header">Probe</div>
+  <div class="section">
+    <div style="width: 100%; margin: auto;" class="plot-cont" id="probe-plot"></div>
+  </div>-->
+{% endblock %}

phaser/web/templates/manager.html

@@ -0,0 +1,19 @@
+{% extends "base.html" %}
+{% block title %}Phaser frontend{% endblock %}
+
+{% block head %}
+  <script type="text/javascript" src="{{ url_for('static', filename='bundle-manager.js') }}" defer></script>
+{% endblock %}
+
+{% block body %}
+  <div id="app"></div>
+  <!--<div class="section-header">Start new reconstruction</div>
+  <div class="section">
+    <button id="start">Start reconstruction</button>
+    <button id="start-fake">Start fake reconstruction</button>
+  </div>
+  <div class="section-header">Current reconstructions</div>
+  <div class="section">
+    <table id="recons-table"></table>
+  </div>-->
+{% endblock %}

phaser/web/types.py

@@ -0,0 +1,255 @@
+import datetime
+import logging
+import typing as t
+
+import pane
+from pane.annotations import Tagged
+from typing_extensions import Self
+
+from phaser.types import Cancelled
+from .util import ReconsStateConverter
+
+
+JobID: t.TypeAlias = str
+WorkerID: t.TypeAlias = str
+
+WorkerStatus: t.TypeAlias = t.Literal['queued', 'starting', 'reloading', 'idle', 'running', 'stopping', 'stopped', 'unknown']
+"""
+Enum indicating the state of the connection, as described by either the server or client.
+
+- 'queued': Worker is queued/scheduled to start
+- 'starting': Worker is actively being started
+- 'reloading': Worker is being reloaded
+- 'idle': No job is running
+- 'running': Job is running
+- 'stopping': Worker is being stopped
+- 'stopped': Worker has been stopped (normally or abnormally)
+
+- 'unknown': Job hasn't talked in a while, unknown state
+"""
+
+JobStatus: t.TypeAlias = t.Literal['queued', 'starting', 'running', 'stopping', 'stopped']
+"""
+Enum indicating the state of a reconstruction job
+"""
+
+Result: t.TypeAlias = t.Literal['finished', 'errored', 'cancelled', 'interrupted']
+"""
+Enum indicating a job or worker result.
+
+- 'finished': Finished normally
+- 'errored': Encounted an error running
+- 'cancelled': Cancelled/shutdown by the user (from server side)
+- 'interrupted': Cancelled/shutdown by the client/OS
+"""
+
+Signal: t.TypeAlias = t.Literal['shutdown', 'cancel', 'reload']
+"""
+Signal sent to a job or worker.
+"""
+
+class ValidationError(Exception):
+    def __init__(self, msg: str):
+        self.msg: str = msg
+
+class SignalException(Cancelled):
+    def __init__(self, signal: Signal, urgent: bool = False):
+        self.signal: Signal = signal
+        self.urgent: bool = urgent
+
+# server -> client messages
+
+class WorkerState(pane.PaneBase):
+    worker_id: WorkerID
+    status: WorkerStatus
+    links: t.Dict[str, str] = pane.field(default_factory=dict)
+
+class WorkerUpdate(pane.PaneBase):
+    worker_id: WorkerID
+    status: WorkerStatus
+
+    msg: t.Literal['status_change'] = 'status_change'
+
+class JobState(pane.PaneBase):
+    job_id: JobID
+    status: JobStatus
+    links: t.Dict[str, str] = pane.field(default_factory=dict)
+    state: t.Dict[str, t.Any] = pane.field(converter=ReconsStateConverter(), default_factory=dict)
+    worker_id: t.Optional[WorkerID] = None
+
+    def strip_state(self) -> Self:
+        return type(self)(
+            self.job_id, self.status, self.links
+        )
+
+class LogRecord(pane.PaneBase):
+    i: int
+
+    timestamp: datetime.datetime
+
+    log: str
+
+    logger_name: str
+    log_level: int
+
+    line_number: int
+    func_name: t.Optional[str] = None
+    stack_info: t.Optional[str] = None
+
+class JobStatusChange(pane.PaneBase):
+    status: JobStatus
+    job_id: JobID
+
+    msg: t.Literal['status_change'] = 'status_change'
+
+class JobUpdate(pane.PaneBase):
+    state: t.Dict[str, t.Any] = pane.field(converter=ReconsStateConverter())
+    job_id: JobID
+
+    msg: t.Literal['job_update'] = 'job_update'
+
+class LogUpdate(pane.PaneBase):
+    new_logs: t.List[LogRecord]
+
+    msg: t.Literal['log'] = 'log'
+
+class JobStopped(pane.PaneBase):
+    result: Result
+    error: t.Optional[str] = None
+
+    msg: t.Literal['job_stopped'] = 'job_stopped'
+
+JobMessage: t.TypeAlias = t.Annotated[t.Union[
+    JobStatusChange, JobUpdate, LogUpdate, JobStopped
+], Tagged('msg')]
+
+class DashboardConnected(pane.PaneBase):
+    state: JobState
+    msg: t.Literal['connected'] = 'connected'
+
+DashboardMessage: t.TypeAlias = t.Annotated[t.Union[
+    JobStatusChange, JobUpdate, JobStopped, DashboardConnected
+], Tagged('msg')]
+
+class WorkersUpdate(pane.PaneBase):
+    event: t.Optional[WorkerUpdate]
+    state: t.List[WorkerState]
+
+    msg: t.Literal['workers_update'] = 'workers_update'
+
+class JobsUpdate(pane.PaneBase):
+    event: JobMessage
+    state: t.List[JobState]
+
+    msg: t.Literal['jobs_update'] = 'jobs_update'
+
+class ManagerConnected(pane.PaneBase):
+    workers: t.List[WorkerState]
+    jobs: t.List[JobState]
+
+    msg: t.Literal['connected'] = 'connected'
+
+ManagerMessage: t.TypeAlias = t.Annotated[t.Union[
+    JobsUpdate, WorkersUpdate, ManagerConnected
+], Tagged('msg')]
+
+# worker -> server messages
+
+class ConnectMessage(pane.PaneBase):
+    """Message sent when a worker starts up"""
+    msg: t.Literal['connect'] = 'connect'
+
+class PollMessage(pane.PaneBase):
+    """Message polling the server for a job"""
+    msg: t.Literal['poll'] = 'poll'
+
+class PingMessage(pane.PaneBase):
+    """Message pinging the server, potentially receiving a cancellation"""
+    msg: t.Literal['ping'] = 'ping'
+
+class UpdateMessage(pane.PaneBase):
+    """Message containing some state update"""
+    state: t.Dict[str, t.Any] = pane.field(converter=ReconsStateConverter())
+    job_id: JobID
+    msg: t.Literal['job_update'] = 'job_update'
+
+class LogMessage(pane.PaneBase):
+    """Message corresponding to a log entry"""
+    job_id: t.Optional[JobID]
+
+    timestamp: datetime.datetime
+
+    log: str
+
+    logger_name: str
+    log_level: int
+
+    line_number: int
+    func_name: t.Optional[str] = None
+    stack_info: t.Optional[str] = None
+
+    msg: t.Literal['log'] = 'log'
+
+    @classmethod
+    def from_logrecord(cls, job_id: t.Optional[JobID], record: logging.LogRecord) -> Self:
+        timestamp = datetime.datetime.fromtimestamp(record.created)
+        return cls(
+            timestamp=timestamp,
+            log=record.getMessage(),
+            job_id=job_id,
+            logger_name=record.name,
+            log_level=record.levelno,
+            line_number=record.lineno,
+            func_name=record.funcName,
+            stack_info=record.stack_info,
+        )
+
+    def into_record(self, i: int) -> LogRecord:
+        return LogRecord.make_unchecked(
+            i, self.timestamp, self.log, self.logger_name, self.log_level,
+            self.line_number, self.func_name, self.stack_info
+        )
+
+class JobResultMessage(pane.PaneBase):
+    """Message indicating job has stopped (normally or abnormally)"""
+    job_id: JobID
+    result: Result
+    error: t.Optional[str] = None
+
+    msg: t.Literal['job_result'] = 'job_result'
+
+class WorkerShutdownMessage(pane.PaneBase):
+    """Message indicating a worker has been stopped"""
+    result: Result
+    error: t.Optional[str] = None
+    detail: t.Optional[str] = None
+
+    msg: t.Literal['shutdown'] = 'shutdown'
+
+
+# server -> worker responses
+
+class JobResponse(pane.PaneBase):
+    """Response containing a job to run"""
+    job_id: JobID
+    plan: str
+    msg: t.Literal['job'] = 'job'
+
+class OkResponse(pane.PaneBase):
+    """Response from server indicating no action should be taken"""
+    msg: t.Literal['ok'] = 'ok'
+
+class SignalResponse(pane.PaneBase):
+    """Response from server signalling some action (cancellation, shutdown, or reload)"""
+    signal: t.Literal['shutdown', 'cancel', 'reload'] = 'cancel'
+    urgent: bool = False
+    msg: t.Literal['signal'] = 'signal'
+
+WorkerMessage: t.TypeAlias = t.Annotated[t.Union[
+    PollMessage, ConnectMessage, PingMessage, UpdateMessage,
+    LogMessage, JobResultMessage, WorkerShutdownMessage,
+], Tagged('msg')]
+
+ServerResponse: t.TypeAlias = t.Annotated[t.Union[
+    JobResponse, SignalResponse, OkResponse
+], Tagged('msg')]

phaser/web/util.py

@@ -0,0 +1,116 @@
+import asyncio
+import base64
+import dataclasses
+import typing as t
+
+import numpy
+from pane.converters import Converter
+from pane.errors import ProductErrorNode, WrongTypeError, ParseInterrupt
+
+T = t.TypeVar('T')
+
+class _array_dummy():
+    def __init__(self, array_interface: t.Any):
+        self.__array_interface__ = array_interface
+
+
+async def merge_streams(
+    *its: t.AsyncIterable[T]
+) -> t.AsyncIterator[T]:
+    # TODO: better error handling here
+    queue = asyncio.Queue(1)
+
+    async def task(it: t.AsyncIterable[T]):
+        async for item in it:
+            await queue.put(item)
+
+    async with asyncio.TaskGroup() as tg:
+        tasks = [tg.create_task(task(it)) for it in its]
+
+        while not all(task.done() for task in tasks):
+            yield await queue.get()
+
+
+def encode_obj(obj: t.Any, to_numpy: bool = True) -> t.Any:
+    if isinstance(obj, numpy.ndarray):
+        if not to_numpy:
+            return obj.tolist()
+        d = obj.__array_interface__
+        d['data'] = base64.urlsafe_b64encode(obj.tobytes()).decode('ascii')
+        d['_ty'] = 'numpy'
+        return d
+
+    if isinstance(obj, bytes):
+        return {
+            '_ty': 'bytes',
+            'data': base64.urlsafe_b64encode(obj).decode('ascii')
+        }
+
+    if isinstance(obj, dict):
+        d = obj
+    elif dataclasses.is_dataclass(obj):
+        d = obj.asdict()  # type: ignore
+    else:
+        return obj
+
+    return {k: encode_obj(v, to_numpy and k not in ('sampling',)) for (k, v) in d.items()}
+
+
+def decode_obj(obj: t.Any) -> t.Any:
+    if isinstance(obj, dict):
+        d = obj
+    elif dataclasses.is_dataclass(obj):
+        d = obj.asdict()  # type: ignore
+    else:
+        return obj
+
+    if '_ty' not in d:
+        return {k: decode_obj(v) for (k, v) in d.items()}
+    ty = d.pop('_ty')
+
+    if ty == 'numpy':
+        d['data'] = base64.urlsafe_b64decode(d['data'].encode('utf-8'))
+        d['shape'] = tuple(d['shape'])
+        return numpy.array(_array_dummy(d))
+
+    if ty == 'bytes':
+        return base64.urlsafe_b64decode(d['data'].encode('utf-8'))
+
+    raise ValueError(f"Unknown custom type '{ty}', while parsing reconstruction state update")
+
+
+class ReconsStateConverter(Converter[t.Dict[str, t.Any]]):
+    def __init__(self):
+        ...
+
+    def expected(self, plural: bool = False) -> str:
+        return "reconstruction state"
+
+    def into_data(self, val: t.Any) -> t.Any:
+        return encode_obj(val)
+
+    def try_convert(self, val: t.Any) -> t.Dict[str, t.Any]:
+        try:
+            val = decode_obj(val)
+        except Exception:
+            raise ParseInterrupt()
+        #if not isinstance(val, t.Mapping) or ({'iter', 'wavelength'} - val.keys()):
+        #    raise ParseInterrupt()
+
+        return val  # type: ignore
+
+    def collect_errors(self, val: t.Any) -> t.Union[ProductErrorNode, WrongTypeError, None]:
+        try:
+            val = decode_obj(val)
+        except Exception as e:
+            import traceback
+
+            tb = e.__traceback__.tb_next  # type: ignore
+            tb = traceback.TracebackException(type(e), e, tb)
+            return WrongTypeError(self.expected(), val, tb)
+        if not isinstance(val, t.Mapping):
+            return WrongTypeError(self.expected(), val)
+        #if (missing := {'iter', 'wavelength'} - val.keys()):
+        #    return ProductErrorNode(self.expected(), {}, val, missing)
+
+        return None