flux-batch 0.0.0__tar.gz → 0.0.11__tar.gz

{flux_batch-0.0.0/flux_batch.egg-info → flux_batch-0.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flux-batch
-Version: 0.0.0
+Version: 0.0.11
 Summary: Python SDK for flux batch jobs and services
 Home-page: https://github.com/converged-computing/flux-batch
 Author: Vanessa Sochat
@@ -26,12 +26,17 @@ Provides-Extra: all
 Requires-Dist: pyyaml; extra == "all"
 Requires-Dist: ply; extra == "all"
 Requires-Dist: pytest>=4.6.2; extra == "all"
+Requires-Dist: sqlalchemy; extra == "all"
+Requires-Dist: rich; extra == "all"
+Provides-Extra: scribe
+Requires-Dist: sqlalchemy; extra == "scribe"
+Requires-Dist: rich; extra == "scribe"
 
 # flux-batch
 
 > Python SDK to generate Flux batch jobs and services
 
-[![PyPI version](https://badge.fury.io/py/flux-batch.svg)](https://badge.fury.io/py/flux-batch)
+![PyPI - Version](https://img.shields.io/pypi/v/flux-batch)
 
 ![https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png](https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png)
 
@@ -46,6 +51,7 @@ Requires-Dist: pytest>=4.6.2; extra == "all"
 
  - **flux-scribe**: Write job events to a local sqlite database via the JournalConsumer (not added yet, written and needs testing)
 
+
 ## Usage
 
 This is a small Flux utility that makes it easy to create Flux batch jobs and services.
@@ -60,9 +66,26 @@ flux start
 pip install -e . --break-system-packages
 ```
 
-### Example
+### Examples
+
+We have a few simple examples:
+
+#### Saving Logs
+
+```bash
+python3 ./examples/save_logs.py
+```
+
+#### Flux Scribe Module
+
+```bash
+export FLUX_SCRIBE_DATABASE=sqlite:///flux-batch-job.db
+python3 ./examples/flux_scribe_module.py
+```
+
+#### General Test
 
-Run the controlled example to see a batch job with prolog and epilog run and complete:
+Or run the controlled example to see a batch job with prolog and epilog run and complete:
 
 ```bash
 python3 ./tests/test_flux_batch.py
@@ -117,7 +140,9 @@ jobspec = flux_batch.BatchJobspecV1.from_jobs(
     nodes=1,
     nslots=1,
     time_limit="10m",
-    job_name="test-batch"
+    job_name="test-batch",
+    # Add saving of logs, info, and metadata
+    logs_dir="./logs",
 )
 
 # Add a prolog and epilog
@@ -125,7 +150,7 @@ jobspec.add_prolog("echo 'Batch Wrapper Starting'")
 jobspec.add_epilog("echo 'Batch Wrapper Finished'")
 
 # Add a service (this assumes user level that exists)
-# jobspec.add_service("my-service'")
+jobspec.add_service("flux-scribe")
 
 # Preview it
 print(flux_batch.submit(handle, jobspec, dry_run=True))

{flux_batch-0.0.0 → flux_batch-0.0.11}/README.md

@@ -2,7 +2,7 @@
 
 > Python SDK to generate Flux batch jobs and services
 
-[![PyPI version](https://badge.fury.io/py/flux-batch.svg)](https://badge.fury.io/py/flux-batch)
+![PyPI - Version](https://img.shields.io/pypi/v/flux-batch)
 
 ![https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png](https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png)
 
@@ -17,6 +17,7 @@
 
  - **flux-scribe**: Write job events to a local sqlite database via the JournalConsumer (not added yet, written and needs testing)
 
+
 ## Usage
 
 This is a small Flux utility that makes it easy to create Flux batch jobs and services.
@@ -31,9 +32,26 @@ flux start
 pip install -e . --break-system-packages
 ```
 
-### Example
+### Examples
+
+We have a few simple examples:
+
+#### Saving Logs
+
+```bash
+python3 ./examples/save_logs.py
+```
+
+#### Flux Scribe Module
+
+```bash
+export FLUX_SCRIBE_DATABASE=sqlite:///flux-batch-job.db
+python3 ./examples/flux_scribe_module.py
+```
+
+#### General Test
 
-Run the controlled example to see a batch job with prolog and epilog run and complete:
+Or run the controlled example to see a batch job with prolog and epilog run and complete:
 
 ```bash
 python3 ./tests/test_flux_batch.py
@@ -88,7 +106,9 @@ jobspec = flux_batch.BatchJobspecV1.from_jobs(
     nodes=1,
     nslots=1,
     time_limit="10m",
-    job_name="test-batch"
+    job_name="test-batch",
+    # Add saving of logs, info, and metadata
+    logs_dir="./logs",
 )
 
 # Add a prolog and epilog
@@ -96,7 +116,7 @@ jobspec.add_prolog("echo 'Batch Wrapper Starting'")
 jobspec.add_epilog("echo 'Batch Wrapper Finished'")
 
 # Add a service (this assumes user level that exists)
-# jobspec.add_service("my-service'")
+jobspec.add_service("flux-scribe")
 
 # Preview it
 print(flux_batch.submit(handle, jobspec, dry_run=True))

{flux_batch-0.0.0 → flux_batch-0.0.11}/flux_batch/__init__.py

@@ -1,7 +1,8 @@
 from .jobspec import BatchJobspecV1
 from .models import BatchAttributesV1, BatchJobV1
+from .submit import preview as jobspec
 from .submit import submit
 
-__all__ = ["BatchJobV1", "BatchAttributesV1", "BatchJobspecV1", "submit"]
+__all__ = ["BatchJobV1", "BatchAttributesV1", "BatchJobspecV1", "submit", "jobspec"]
 
 from .version import __version__  # noqa

{flux_batch-0.0.0 → flux_batch-0.0.11}/flux_batch/jobspec.py

@@ -1,7 +1,9 @@
+import os
 import shlex
 from typing import List
 
 import flux_batch.models as models
+import flux_batch.script as scripts
 
 
 class BatchJobspecV1:
@@ -18,6 +20,7 @@ class BatchJobspecV1:
         self.prologs: List[str] = []
         self.epilogs: List[str] = []
         self.services: List[str] = []
+        self.modules: List[str] = []
 
     @classmethod
     def from_command(cls, command: List[str], **kwargs):
@@ -51,6 +54,9 @@ class BatchJobspecV1:
     def add_epilog(self, cmd: str):
         self.epilogs.append(cmd)
 
+    def add_module(self, service_name: str):
+        self.modules.append(service_name)
+
     def get_cli_flags(self) -> List[str]:
         """
         Converts BatchAttributesV1 into a list of strings for subprocess.
@@ -107,8 +113,25 @@ class BatchJobspecV1:
             for val in getattr(attr, field_name):
                 flags.extend([flag, str(val)])
 
+        # If we have modules, ensure they are added --conf <module>=true
+        if self.modules:
+            # Tell Flux to look in our user home for rc scripts
+            modprobe_path = os.path.expanduser("~/.flux-batch")
+            flags.extend(["--env", f"FLUX_MODPROBE_PATH_APPEND={modprobe_path}"])
+
+            # If modules are used, we need to pass the service names into the Flux config
+            # so the @task 'needs_config' filter allows them to run
+            for mod in self.modules:
+                flags.extend(["--conf", f"{mod}=true"])
+
         return flags
 
+    def render(self) -> str:
+        """
+        Generate the jobspec.
+        """
+        return self.generate_wrapper_script()
+
     def generate_wrapper_script(self) -> str:
         """
         Generate the wrapper script.
@@ -119,14 +142,38 @@ class BatchJobspecV1:
         4. Add jobs/commands
         5. Stop services
         6. And epilogs
-        """
+        7. Custom scripts
 
+        Yes, it's redundant to write them as comments but I like the organization. -v
+        """
+        # hashbang
         lines = ["#!/bin/bash"]
+
+        # prologs
         lines.extend(self.prologs)
         for s in self.services:
             lines.append(f"systemctl --user start {s}")
+
+        # commands that are derived from jobs or command
         lines.extend(self.commands)
+
+        # stop services
         for s in reversed(self.services):
             lines.append(f"systemctl --user stop {s}")
+
+        # epilogs
         lines.extend(self.epilogs)
+
+        # custom user scripts
+        if self.attributes.logs_dir is not None:
+            lines.append(self.script_save_logs())
         return "\n".join(lines)
+
+    def script_save_logs(self):
+        """
+        Custom saving of logs. This is what we wrote for our peformance study!
+        """
+        script_path = scripts.get_script("save_logs.sh")
+
+        # Determine output directory (use home default if not defined)
+        return f"bash {script_path} {self.attributes.logs_dir}"

{flux_batch-0.0.0 → flux_batch-0.0.11}/flux_batch/models.py

@@ -24,6 +24,11 @@ class BatchAttributesV1:
     Explicitly defined arguments allowed by flux batch for V1 spec
     """
 
+    # These are added / custom to our module
+    # If logs directory defined (not None) save output there
+    # We force the user to provide something.
+    logs_dir: Optional[Union[bool, str]] = None
+
     # Resources
     nslots: Optional[int] = None  # -n
     cores_per_slot: Optional[int] = None  # -c

flux_batch-0.0.11/flux_batch/script/__init__.py

@@ -0,0 +1,16 @@
+import os
+
+import flux_batch
+
+
+def get_script(name):
+    """
+    Get a script by name
+    """
+    # Find the path to the installed script
+    base_path = os.path.dirname(os.path.abspath(flux_batch.__file__))
+    script_path = os.path.join(base_path, "script", name)
+    if not os.path.exists(script_path):
+        print(f"Warning: script {name} does not exist")
+        return
+    return script_path

flux_batch-0.0.11/flux_batch/script/save_logs.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+output=$1
+mkdir -p $output
+
+echo "Saving logs and job metadata to ${output}"
+
+# This will save logs, events, and jobspecs
+for jobid in $(flux jobs -a --json | jq -r .jobs[].id)
+  do
+    echo "Parsing jobid ${jobid}"
+    flux job attach $jobid &> $output/${jobid}.out
+    echo "START OF JOBSPEC" >> $output/${jobid}.out
+    flux job info $jobid jobspec >> $output/${jobid}.out
+    echo "START OF EVENTLOG" >> $output/${jobid}.out
+    flux job info $jobid guest.exec.eventlog >> $output/${jobid}.out
+done

flux_batch-0.0.11/flux_batch/service/__init__.py

@@ -0,0 +1,81 @@
+import os
+import subprocess
+import sys
+
+import flux_batch.service.scribe as scribe
+
+# Lookup of known services
+services = {"flux-scribe": scribe.SERVICE_TEMPLATE}
+modules = {
+    "flux-scribe": {
+        "startup": scribe.START_MODULE_TEMPLATE,
+        "shutdown": scribe.STOP_MODULE_TEMPLATE,
+        "module": scribe.MODULE_NAME,
+    }
+}
+
+
+def write_modprobe_script(rc_path, script, args=None):
+    """
+    Shared function to write service file.
+    """
+    args = args or {}
+    if not os.path.exists(rc_path):
+        with open(rc_path, "w") as f:
+            f.write(script.format(**args))
+
+
+def ensure_modprobe_scripts(service_name: str):
+    """
+    Ensures rc1.d (start) and rc3.d (stop) scripts exist for the service.
+    """
+    if service_name not in modules:
+        print("Warning: module {service_name} is not known.")
+        return
+
+    # We will add these to FLUX_MODPROBE_PATH_APPEND
+    base_dir = os.path.expanduser("~/.flux-batch")
+    for subdir in ["rc1.d", "rc3.d"]:
+        os.makedirs(os.path.join(base_dir, subdir), exist_ok=True)
+
+    service_func = service_name.replace("-", "_")
+
+    # Path for rc1.d (startup)
+    args = {
+        "service_name": service_name,
+        "service_func": service_func,
+        "python_bin": sys.executable,
+        "module_name": modules[service_name]["module"],
+    }
+    rc1_path = os.path.join(base_dir, "rc1.d", f"{service_name}.py")
+    script = modules[service_name]["startup"]
+    write_modprobe_script(rc1_path, script, args=args)
+
+    # Path for rc3.d (shutdown)
+    args = {"service_name": service_name, "service_func": service_func}
+    rc3_path = os.path.join(base_dir, "rc3.d", f"{service_name}.py")
+    script = modules[service_name]["shutdown"]
+    write_modprobe_script(rc3_path, script, args=args)
+
+
+def ensure_user_service(service_name: str):
+    """
+    Checks for the existence of a systemd service file in the user's home.
+    If it doesn't exist, it creates it and reloads the daemon.
+    """
+    user_systemd_dir = os.path.expanduser("~/.config/systemd/user")
+    os.makedirs(user_systemd_dir, exist_ok=True)
+    service_path = os.path.join(user_systemd_dir, f"{service_name}.service")
+
+    if not os.path.exists(service_path):
+        if service_name in services:
+            template = services[service_name]
+            print(f"[*] Provisioning {service_name} at {service_path}")
+            with open(service_path, "w") as f:
+                f.write(template.format(python_path=sys.executable))
+
+        else:
+            print(f"[*] Service {service_name} is not known, assuming exists.")
+
+        # Reload the user-session manager to recognize the new unit
+        subprocess.run(["systemctl", "--user", "daemon-reload"], check=True)

flux_batch-0.0.11/flux_batch/service/scribe/__init__.py

@@ -0,0 +1,2 @@
+MODULE_NAME = "flux_batch.service.scribe"
+from .template import SERVICE_TEMPLATE, START_MODULE_TEMPLATE, STOP_MODULE_TEMPLATE

flux_batch-0.0.11/flux_batch/service/scribe/__main__.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+import errno
+import logging
+import os
+import sys
+import time
+
+import flux
+import flux.job
+
+# Not necessary, but it makes it pretty
+from rich import print
+
+# Use the synchronous version of the backend to avoid asyncio-in-thread conflicts
+from flux_batch.service.scribe.database import SQLAlchemyBackend
+
+# Setup logging to stderr (to avoid polluting stdout if run manually)
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s", stream=sys.stderr
+)
+logger = logging.getLogger("flux-scribe")
+
+
+class JournalScribe:
+    def __init__(self, db_url: str):
+        """
+        Initializes the Scribe with a synchronous DB backend and a Flux Journal Consumer.
+        """
+        # Setup Database
+        logger.info(f"Connecting to Database: {db_url}")
+        self.db = SQLAlchemyBackend(db_url)
+        self.db.initialize()
+
+        try:
+            self.handle = flux.Flux()
+            logger.info("Connected to Flux instance.")
+        except Exception as e:
+            logger.critical(f"Failed to connect to Flux: {e}")
+            sys.exit(1)
+
+        # Initialize Journal Consumer
+        # This consumes the global event log for the entire instance
+        self.consumer = flux.job.JournalConsumer(self.handle)
+        self.running = True
+
+    def _normalize_event(self, event) -> dict:
+        """
+        Converts a Flux event object into the dictionary format expected by record_event.
+        Matches the logic provided in your EventsEngine reference.
+        """
+        # Convert the SWIG/CFFI event object to a dictionary
+        payload = dict(event)
+
+        return {
+            "id": str(getattr(event, "jobid", "unknown")),
+            "type": getattr(event, "name", "unknown"),
+            "timestamp": getattr(event, "timestamp", time.time()),
+            "payload": payload,
+            "R": getattr(event, "R", None),
+            "jobspec": getattr(event, "jobspec", None),
+        }
+
+    def run(self):
+        """
+        Main execution loop. Polls the journal and writes to the DB.
+        """
+        try:
+            logger.info("🚀 Flux Scribe (Journal Consumer) started.")
+            self.consumer.start()
+
+            while self.running:
+                try:
+                    # Non-blocking poll (100ms timeout)
+                    # This allows the loop to check for shutdown signals regularly
+                    event = self.consumer.poll(timeout=0.1)
+
+                    if event:
+                        print(event)
+                        # We only care about events associated with a job
+                        if hasattr(event, "jobid"):
+                            clean_event = self._normalize_event(event)
+                            self.db.record_event("local", clean_event)
+                    else:
+                        # If no event, yield a tiny bit of CPU
+                        time.sleep(0.01)
+
+                except EnvironmentError as e:
+                    # Ignore timeouts (no data)
+                    if e.errno == errno.ETIMEDOUT:
+                        continue
+                    logger.error(f"Flux connection error: {e}")
+                    time.sleep(1)
+
+                except Exception as e:
+                    logger.error(f"Unexpected error in event loop: {e}")
+                    time.sleep(1)
+
+        except Exception as e:
+            logger.critical(f"EventsEngine crashed: {e}")
+        finally:
+            self.db.close()
+            logger.info("EventsEngine thread exiting.")
+
+
+def main():
+    # Retrieve DB path from environment or use a default
+    db_path = os.environ.get("FLUX_SCRIBE_DATABASE", "sqlite:///server_state.db")
+    scribe = JournalScribe(db_path)
+    scribe.run()
+
+
+if __name__ == "__main__":
+    main()

flux_batch-0.0.11/flux_batch/service/scribe/database.py

@@ -0,0 +1,150 @@
+import time
+from typing import Any, Dict, List, Optional
+
+from sqlalchemy import and_, create_engine, select, update
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+from sqlalchemy.orm import sessionmaker
+
+from flux_batch.service.scribe.models import Base, EventModel, EventRecord, JobModel, JobRecord
+
+
+def _record_event_internal(session, cluster: str, event: Dict[str, Any]):
+    """
+    Shared synchronous logic for recording events.
+    Used by both Sync and Async backends.
+    """
+    job_id = event.get("id")
+    event_type = event.get("type")
+    data = event.get("payload", {})
+    timestamp = event.get("timestamp", time.time())
+
+    new_event = EventModel(
+        job_id=job_id,
+        cluster=cluster,
+        timestamp=timestamp,
+        event_type=event_type,
+        payload=data,
+    )
+    session.add(new_event)
+
+    if event_type == "submit":
+        stmt = select(JobModel).where(and_(JobModel.job_id == job_id, JobModel.cluster == cluster))
+        job = session.execute(stmt).scalar_one_or_none()
+
+        if not job:
+            job = JobModel(
+                job_id=job_id,
+                cluster=cluster,
+                user=str(data.get("userid", "unknown")),
+                state="submitted",
+                workdir=data.get("cwd", ""),
+                submit_time=timestamp,
+                last_updated=timestamp,
+            )
+            session.add(job)
+        else:
+            job.state = "submitted"
+            job.last_updated = timestamp
+
+    # state transitions
+    elif event_type == "state" or (event_type and event_type.endswith(".finish")):
+        state_name = data.get("state_name", event_type)
+        stmt = select(JobModel).where(and_(JobModel.job_id == job_id, JobModel.cluster == cluster))
+        job = session.execute(stmt).scalar_one_or_none()
+        if job:
+            job.state = state_name
+            job.last_updated = time.time()
+            if "status" in data:
+                job.exit_code = data["status"]
+
+
+class AsyncSQLAlchemyBackend:
+    """
+    Asynchronous backend for the MCP Gateway.
+    """
+
+    def __init__(self, db_url: str):
+        self.engine = create_async_engine(db_url, echo=False)
+        self.SessionLocal = async_sessionmaker(self.engine, expire_on_commit=False)
+
+    async def initialize(self):
+        async with self.engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+
+    async def close(self):
+        await self.engine.dispose()
+
+    async def record_event(self, cluster: str, event: Dict[str, Any]):
+        async with self.SessionLocal() as session:
+            # run_sync bridges our shared logic into the async session
+            await session.run_sync(_record_event_internal, cluster, event)
+            await session.commit()
+
+    async def get_job(self, cluster: str, job_id: int) -> Optional[JobRecord]:
+        async with self.SessionLocal() as session:
+            result = await session.execute(
+                select(JobModel).where(and_(JobModel.job_id == job_id, JobModel.cluster == cluster))
+            )
+            job = result.scalar_one_or_none()
+            return job.to_record() if job else None
+
+    async def get_event_history(self, cluster: str, job_id: int) -> List[EventRecord]:
+        async with self.SessionLocal() as session:
+            result = await session.execute(
+                select(EventModel)
+                .where(and_(EventModel.job_id == job_id, EventModel.cluster == cluster))
+                .order_by(EventModel.timestamp.asc())
+            )
+            return [e.to_record() for e in result.scalars().all()]
+
+    async def search_jobs(
+        self, cluster: str = None, state: str = None, limit: int = 10
+    ) -> List[JobRecord]:
+        async with self.SessionLocal() as session:
+            stmt = select(JobModel)
+            if cluster:
+                stmt = stmt.where(JobModel.cluster == cluster)
+            if state:
+                stmt = stmt.where(JobModel.state == state)
+            result = await session.execute(stmt.limit(limit))
+            return [j.to_record() for j in result.scalars().all()]
+
+
+class SQLAlchemyBackend:
+    """
+    Synchronous backend for the standalone Scribe daemon.
+    """
+
+    def __init__(self, db_url: str):
+        # strip 'aiosqlite+' or similar if passed from shared config
+        url = db_url.replace("+aiosqlite", "").replace("+asyncpg", "")
+        self.engine = create_engine(url, echo=False)
+        self.SessionLocal = sessionmaker(bind=self.engine, expire_on_commit=False)
+
+    def initialize(self):
+        Base.metadata.create_all(self.engine)
+
+    def close(self):
+        self.engine.dispose()
+
+    def record_event(self, cluster: str, event: Dict[str, Any]):
+        with self.SessionLocal() as session:
+            with session.begin():
+                _record_event_internal(session, cluster, event)
+
+    def get_unwatched_job_ids(self, cluster: str) -> List[int]:
+        """Specific for Scribe: find jobs that need a watcher."""
+        with self.SessionLocal() as session:
+            stmt = select(JobModel.job_id).where(
+                and_(JobModel.cluster == cluster, JobModel.state == "submitted")
+            )
+            return list(session.execute(stmt).scalars().all())
+
+    def mark_job_as_watched(self, cluster: str, job_id: int):
+        with self.SessionLocal() as session:
+            with session.begin():
+                session.execute(
+                    update(JobModel)
+                    .where(and_(JobModel.job_id == job_id, JobModel.cluster == cluster))
+                    .values(state="watching")
+                )

flux_batch-0.0.11/flux_batch/service/scribe/models.py

@@ -0,0 +1,94 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from sqlalchemy import JSON, Float, Integer, String
+from sqlalchemy.ext.asyncio import AsyncAttrs
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+# DTOs are "Public Data Transfer Objects" and they are used by
+# our interfaces and tools
+
+
+@dataclass
+class JobRecord:
+    """
+    Represents a snapshot of a job state.
+    Returned by get_job() and search_jobs().
+    """
+
+    job_id: int
+    cluster: str
+    state: str
+    user: str
+    workdir: Optional[str] = None
+    exit_code: Optional[int] = None
+    submit_time: float = 0.0
+    last_updated: float = 0.0
+
+
+@dataclass
+class EventRecord:
+    """
+    Represents a single historical event.
+    Returned by get_event_history().
+    """
+
+    timestamp: float
+    event_type: str
+    payload: Dict[str, Any]
+
+
+# Database models for SQLAlchemy ORM
+
+
+class Base(AsyncAttrs, DeclarativeBase):
+    pass
+
+
+class JobModel(Base):
+    __tablename__ = "jobs"
+
+    # Composite Primary Key
+    job_id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    cluster: Mapped[str] = mapped_column(String(255), primary_key=True)
+
+    state: Mapped[str] = mapped_column(String(50))
+    user: Mapped[str] = mapped_column(String(255), nullable=True)
+    workdir: Mapped[Optional[str]] = mapped_column(String, nullable=True)
+    exit_code: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
+    submit_time: Mapped[float] = mapped_column(Float, default=0.0)
+    last_updated: Mapped[float] = mapped_column(Float, default=0.0)
+
+    def to_record(self) -> JobRecord:
+        """
+        Helper to convert ORM model to public DTO
+        """
+        return JobRecord(
+            job_id=self.job_id,
+            cluster=self.cluster,
+            state=self.state,
+            user=self.user,
+            workdir=self.workdir,
+            exit_code=self.exit_code,
+            submit_time=self.submit_time,
+            last_updated=self.last_updated,
+        )
+
+
+class EventModel(Base):
+    __tablename__ = "events"
+
+    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
+    job_id: Mapped[int] = mapped_column(Integer, index=True)
+    cluster: Mapped[str] = mapped_column(String(255), index=True)
+    timestamp: Mapped[float] = mapped_column(Float)
+    event_type: Mapped[str] = mapped_column(String(50))
+    payload: Mapped[Dict[str, Any]] = mapped_column(JSON)
+
+    def to_record(self) -> EventRecord:
+        """
+        Helper to convert ORM model to public DTO
+        """
+        return EventRecord(
+            timestamp=self.timestamp, event_type=self.event_type, payload=self.payload
+        )

flux_batch-0.0.11/flux_batch/service/scribe/template.py

@@ -0,0 +1,53 @@
+# Template for the Scribe Journal Consumer
+SERVICE_TEMPLATE = """[Unit]
+Description=Flux Scribe Journal Consumer
+After=network.target
+
+[Service]
+ExecStart={python_path} -m flux_batch.service.scribe
+Restart=on-failure
+
+[Install]
+WantedBy=default.target
+"""
+
+START_MODULE_TEMPLATE = """
+from flux.modprobe import task
+import flux.subprocess as subprocess
+
+@task(
+    "start-{service_name}",
+    ranks="0",
+    needs_config=["{service_name}"],
+    after=["resource", "job-list"],
+)
+def start_{service_func}(context):
+    # This triggers the systemd user service provisioned earlier
+    # context.bash("systemctl --user start {service_name}")
+    subprocess.rexec_bg(
+        context.handle,
+        ["{python_bin}", "-m", "{module_name}"],
+        label="{service_name}",
+        nodeid=0
+    )
+"""
+
+STOP_MODULE_TEMPLATE = """
+from flux.modprobe import task
+import flux.subprocess as subprocess
+
+@task(
+    "stop-{service_name}",
+    ranks="0",
+    needs_config=["{service_name}"],
+    before=["resource", "job-list"],
+)
+def stop_{service_func}(context):
+    # context.bash("systemctl --user stop {service_name}")
+    subprocess.kill(context.handle, signum=2, label="{service_name}").get()
+    try:
+        status = subprocess.wait(context.handle, label="{service_name}").get()["status"]
+        print(status)
+    except:
+        pass
+"""

flux_batch-0.0.11/flux_batch/submit.py

@@ -0,0 +1,86 @@
+import json
+import os
+import stat
+import subprocess
+import tempfile
+
+import flux
+import flux.job
+
+import flux_batch.models as models
+import flux_batch.service as services
+import flux_batch.utils as utils
+
+
+def setup(spec):
+    """
+    shared function to generate services / modules from a spec.
+    """
+    # Provision services (like flux-scribe) if requested
+    for service in spec.services:
+        services.ensure_user_service(service)
+    for module in spec.modules:
+        services.ensure_modprobe_scripts(module)
+
+
+def generate_jobspec(spec, script, wrapper_path):
+    """
+    Shared function to write a script to a wrapper path and generate
+    a jobspec for it via flux batch --dry-run.
+    """
+    utils.write_file(script, wrapper_path)
+
+    # Make the script executable so 'flux batch' can analyze it
+    os.chmod(wrapper_path, os.stat(wrapper_path).st_mode | stat.S_IEXEC)
+
+    # Generate the RFC 25 Jobspec JSON via the Flux CLI
+    # This handles all resource mapping (-N, -n, etc.)
+    cmd = ["flux", "batch"] + spec.get_cli_flags() + ["--dry-run", wrapper_path]
+
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Error during flux batch dryrun: {e.stderr}")
+        raise
+    return result.stdout
+
+
+def preview(spec: models.BatchJobV1) -> int:
+    """
+    Preview the jobspec.
+    """
+    setup(spec)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+
+        # Write the wrapper script (handling prologs, services, and jobs)
+        wrapper_path = os.path.join(tmpdir, "wrapper.sh")
+        script = spec.generate_wrapper_script()
+        jobspec = generate_jobspec(spec, script, wrapper_path)
+        return json.loads(jobspec)
+
+
+def submit(handle: flux.Flux, spec: models.BatchJobV1, dry_run=False) -> int:
+    """
+    Orchestrates the submission process:
+    1. Provisions any required user-space services.
+    2. Generates the wrapper shell script.
+    3. Uses 'flux batch --dryrun' to compile the Jobspec JSON.
+    4. Submits the Jobspec to the Flux instance.
+    """
+    setup(spec)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Write the wrapper script (handling prologs, services, and jobs)
+        wrapper_path = os.path.join(tmpdir, "wrapper.sh")
+
+        # dry run here just displays it
+        script = spec.generate_wrapper_script()
+        if dry_run:
+            return script
+
+        jobspec = generate_jobspec(spec, script, wrapper_path)
+
+        # Submit the JSON string to the Flux instance
+        # The result.stdout contains the raw JSON Jobspec
+        return flux.job.submit(handle, jobspec)

{flux_batch-0.0.0 → flux_batch-0.0.11}/flux_batch/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.0"
+__version__ = "0.0.11"
 AUTHOR = "Vanessa Sochat"
 AUTHOR_EMAIL = "vsoch@users.noreply.github.com"
 NAME = "flux-batch"
@@ -13,4 +13,5 @@ INSTALL_REQUIRES = (
 )
 
 TESTS_REQUIRES = (("pytest", {"min_version": "4.6.2"}),)
-INSTALL_REQUIRES_ALL = INSTALL_REQUIRES + TESTS_REQUIRES
+SCRIBE_REQUIRES = (("sqlalchemy", {"min_version": None}), ("rich", {"min_version": None}))
+INSTALL_REQUIRES_ALL = INSTALL_REQUIRES + TESTS_REQUIRES + SCRIBE_REQUIRES

{flux_batch-0.0.0 → flux_batch-0.0.11/flux_batch.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flux-batch
-Version: 0.0.0
+Version: 0.0.11
 Summary: Python SDK for flux batch jobs and services
 Home-page: https://github.com/converged-computing/flux-batch
 Author: Vanessa Sochat
@@ -26,12 +26,17 @@ Provides-Extra: all
 Requires-Dist: pyyaml; extra == "all"
 Requires-Dist: ply; extra == "all"
 Requires-Dist: pytest>=4.6.2; extra == "all"
+Requires-Dist: sqlalchemy; extra == "all"
+Requires-Dist: rich; extra == "all"
+Provides-Extra: scribe
+Requires-Dist: sqlalchemy; extra == "scribe"
+Requires-Dist: rich; extra == "scribe"
 
 # flux-batch
 
 > Python SDK to generate Flux batch jobs and services
 
-[![PyPI version](https://badge.fury.io/py/flux-batch.svg)](https://badge.fury.io/py/flux-batch)
+![PyPI - Version](https://img.shields.io/pypi/v/flux-batch)
 
 ![https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png](https://github.com/converged-computing/flux-batch/raw/main/img/flux-batch-small.png)
 
@@ -46,6 +51,7 @@ Requires-Dist: pytest>=4.6.2; extra == "all"
 
  - **flux-scribe**: Write job events to a local sqlite database via the JournalConsumer (not added yet, written and needs testing)
 
+
 ## Usage
 
 This is a small Flux utility that makes it easy to create Flux batch jobs and services.
@@ -60,9 +66,26 @@ flux start
 pip install -e . --break-system-packages
 ```
 
-### Example
+### Examples
+
+We have a few simple examples:
+
+#### Saving Logs
+
+```bash
+python3 ./examples/save_logs.py
+```
+
+#### Flux Scribe Module
+
+```bash
+export FLUX_SCRIBE_DATABASE=sqlite:///flux-batch-job.db
+python3 ./examples/flux_scribe_module.py
+```
+
+#### General Test
 
-Run the controlled example to see a batch job with prolog and epilog run and complete:
+Or run the controlled example to see a batch job with prolog and epilog run and complete:
 
 ```bash
 python3 ./tests/test_flux_batch.py
@@ -117,7 +140,9 @@ jobspec = flux_batch.BatchJobspecV1.from_jobs(
     nodes=1,
     nslots=1,
     time_limit="10m",
-    job_name="test-batch"
+    job_name="test-batch",
+    # Add saving of logs, info, and metadata
+    logs_dir="./logs",
 )
 
 # Add a prolog and epilog
@@ -125,7 +150,7 @@ jobspec.add_prolog("echo 'Batch Wrapper Starting'")
 jobspec.add_epilog("echo 'Batch Wrapper Finished'")
 
 # Add a service (this assumes user level that exists)
-# jobspec.add_service("my-service'")
+jobspec.add_service("flux-scribe")
 
 # Preview it
 print(flux_batch.submit(handle, jobspec, dry_run=True))

{flux_batch-0.0.0 → flux_batch-0.0.11}/flux_batch.egg-info/SOURCES.txt

@@ -20,8 +20,14 @@ flux_batch.egg-info/top_level.txt
 flux_batch/logger/__init__.py
 flux_batch/logger/generate.py
 flux_batch/logger/logger.py
+flux_batch/script/__init__.py
+flux_batch/script/save_logs.sh
 flux_batch/service/__init__.py
-flux_batch/service/scribe.py
+flux_batch/service/scribe/__init__.py
+flux_batch/service/scribe/__main__.py
+flux_batch/service/scribe/database.py
+flux_batch/service/scribe/models.py
+flux_batch/service/scribe/template.py
 flux_batch/utils/__init__.py
 flux_batch/utils/fileio.py
 flux_batch/utils/text.py

{flux_batch-0.0.0 → flux_batch-0.0.11}/flux_batch.egg-info/requires.txt

@@ -5,3 +5,9 @@ ply
 pyyaml
 ply
 pytest>=4.6.2
+sqlalchemy
+rich
+
+[scribe]
+sqlalchemy
+rich

{flux_batch-0.0.0 → flux_batch-0.0.11}/setup.py

@@ -62,6 +62,7 @@ if __name__ == "__main__":
     INSTALL_REQUIRES = get_reqs(lookup)
     TESTS_REQUIRES = get_reqs(lookup, "TESTS_REQUIRES")
     INSTALL_REQUIRES_ALL = get_reqs(lookup, "INSTALL_REQUIRES_ALL")
+    INSTALL_REQUIRES_SCRIBE = get_reqs(lookup, "SCRIBE_REQUIRES")
 
     setup(
         name=NAME,
@@ -83,6 +84,7 @@ if __name__ == "__main__":
         tests_require=TESTS_REQUIRES,
         extras_require={
             "all": [INSTALL_REQUIRES_ALL],
+            "scribe": [INSTALL_REQUIRES_SCRIBE],
         },
         classifiers=[
             "Intended Audience :: Science/Research",

flux_batch-0.0.0/flux_batch/service/__init__.py

@@ -1,31 +0,0 @@
-import os
-import subprocess
-import sys
-
-from .scribe import SERVICE_TEMPLATE as scribe_template
-
-# Lookup of known services
-services = {"scribe": scribe_template}
-
-
-def ensure_user_service(service_name: str):
-    """
-    Checks for the existence of a systemd service file in the user's home.
-    If it doesn't exist, it creates it and reloads the daemon.
-    """
-    user_systemd_dir = os.path.expanduser("~/.config/systemd/user")
-    os.makedirs(user_systemd_dir, exist_ok=True)
-    service_path = os.path.join(user_systemd_dir, f"{service_name}.service")
-
-    if not os.path.exists(service_path):
-        if service_name in services:
-            template = services[service_name]
-            print(f"[*] Provisioning {service_name} at {service_path}")
-            with open(service_path, "w") as f:
-                f.write(template.format(python_path=sys.executable))
-
-        else:
-            print(f"[*] Service {service_name} is not known, assuming exists.")
-
-        # Reload the user-session manager to recognize the new unit
-        subprocess.run(["systemctl", "--user", "daemon-reload"], check=True)

flux_batch-0.0.0/flux_batch/service/scribe.py

@@ -1,12 +0,0 @@
-# Template for the Scribe Journal Consumer
-SERVICE_TEMPLATE = """[Unit]
-Description=Flux Scribe Journal Consumer
-After=network.target
-
-[Service]
-ExecStart={python_path} -m flux_mcp_server.scribe
-Restart=on-failure
-
-[Install]
-WantedBy=default.target
-"""

flux_batch-0.0.0/flux_batch/submit.py

@@ -1,53 +0,0 @@
-import os
-import stat
-import subprocess
-import tempfile
-
-import flux
-import flux.job
-
-import flux_batch.models as models
-import flux_batch.utils as utils
-from flux_batch.service import ensure_user_service
-
-
-def submit(handle: flux.Flux, spec: models.BatchJobV1, dry_run=False) -> int:
-    """
-    Orchestrates the submission process:
-    1. Provisions any required user-space services.
-    2. Generates the wrapper shell script.
-    3. Uses 'flux batch --dryrun' to compile the Jobspec JSON.
-    4. Submits the Jobspec to the Flux instance.
-    """
-
-    # Provision services (like flux-scribe) if requested
-    for service in spec.services:
-        ensure_user_service(service)
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        # Write the wrapper script (handling prologs, services, and jobs)
-        wrapper_path = os.path.join(tmpdir, "wrapper.sh")
-
-        # dry run here just displays it
-        script = spec.generate_wrapper_script()
-        if dry_run:
-            return script
-
-        utils.write_file(script, wrapper_path)
-
-        # Make the script executable so 'flux batch' can analyze it
-        os.chmod(wrapper_path, os.stat(wrapper_path).st_mode | stat.S_IEXEC)
-
-        # Generate the RFC 25 Jobspec JSON via the Flux CLI
-        # This handles all resource mapping (-N, -n, etc.)
-        cmd = ["flux", "batch"] + spec.get_cli_flags() + ["--dry-run", wrapper_path]
-
-        try:
-            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
-        except subprocess.CalledProcessError as e:
-            print(f"Error during flux batch dryrun: {e.stderr}")
-            raise
-
-        # Submit the JSON string to the Flux instance
-        # The result.stdout contains the raw JSON Jobspec
-        return flux.job.submit(handle, result.stdout)