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

{flux_batch-0.0.1/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.1
+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
@@ -36,7 +36,7 @@ Requires-Dist: rich; extra == "scribe"
 
 > 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)
 
@@ -70,10 +70,21 @@ pip install -e . --break-system-packages
 
 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
+
 Or run the controlled example to see a batch job with prolog and epilog run and complete:
 
 ```bash

{flux_batch-0.0.1 → 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)
 
@@ -36,10 +36,21 @@ pip install -e . --break-system-packages
 
 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
+
 Or run the controlled example to see a batch job with prolog and epilog run and complete:
 
 ```bash

{flux_batch-0.0.1 → 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.1 → flux_batch-0.0.11}/flux_batch/jobspec.py

@@ -1,3 +1,4 @@
+import os
 import shlex
 from typing import List
 
@@ -19,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):
@@ -52,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.
@@ -108,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.

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/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.1 → flux_batch-0.0.11}/flux_batch/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.0.1"
+__version__ = "0.0.11"
 AUTHOR = "Vanessa Sochat"
 AUTHOR_EMAIL = "vsoch@users.noreply.github.com"
 NAME = "flux-batch"

{flux_batch-0.0.1 → 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.1
+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
@@ -36,7 +36,7 @@ Requires-Dist: rich; extra == "scribe"
 
 > 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)
 
@@ -70,10 +70,21 @@ pip install -e . --break-system-packages
 
 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
+
 Or run the controlled example to see a batch job with prolog and epilog run and complete:
 
 ```bash

flux_batch-0.0.1/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 = {"flux-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.1/flux_batch/service/scribe/__init__.py

@@ -1 +0,0 @@
-from .template import SERVICE_TEMPLATE

flux_batch-0.0.1/flux_batch/service/scribe/template.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_batch.service.scribe
-Restart=on-failure
-
-[Install]
-WantedBy=default.target
-"""

flux_batch-0.0.1/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)