flux-batch 0.0.1__py3-none-any.whl → 0.0.11__py3-none-any.whl

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/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/service/__init__.py

@@ -2,10 +2,60 @@ import os
 import subprocess
 import sys
 
-from .scribe import SERVICE_TEMPLATE as scribe_template
+import flux_batch.service.scribe as scribe
 
 # Lookup of known services
-services = {"flux-scribe": scribe_template}
+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):

flux_batch/service/scribe/__init__.py

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

flux_batch/service/scribe/template.py

@@ -10,3 +10,44 @@ 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/submit.py

@@ -1,3 +1,4 @@
+import json
 import os
 import stat
 import subprocess
@@ -7,8 +8,56 @@ import flux
 import flux.job
 
 import flux_batch.models as models
+import flux_batch.service as services
 import flux_batch.utils as utils
-from flux_batch.service import ensure_user_service
+
+
+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:
@@ -19,10 +68,7 @@ def submit(handle: flux.Flux, spec: models.BatchJobV1, dry_run=False) -> int:
     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)
+    setup(spec)
 
     with tempfile.TemporaryDirectory() as tmpdir:
         # Write the wrapper script (handling prologs, services, and jobs)
@@ -33,21 +79,8 @@ def submit(handle: flux.Flux, spec: models.BatchJobV1, dry_run=False) -> int:
         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
+        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, result.stdout)
+        return flux.job.submit(handle, jobspec)

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.dist-info → flux_batch-0.0.11.dist-info}/METADATA

@@ -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.dist-info → flux_batch-0.0.11.dist-info}/RECORD

@@ -1,28 +1,28 @@
-flux_batch/__init__.py,sha256=ZAZF-E0SbVVY2U1_WkRGZFB4rge5BGqQZJ2MdjloZhY,230
-flux_batch/jobspec.py,sha256=8eYPPgYmup37AvKQddtQax9RweZnsbPFUBe0eOb7__4,4638
+flux_batch/__init__.py,sha256=35032OVyHkBiQIptg7NBY_6mdhthFu1hpqOCsREm5lQ,280
+flux_batch/jobspec.py,sha256=MUMI4Y_DgjCRssgzUDagTTFw8L3t4L_5az1E3tmT1FI,5464
 flux_batch/models.py,sha256=JjrFqi4Skrop_cyIWfgAHTdY53kotkiGD0JpTnVlByI,2200
-flux_batch/submit.py,sha256=TSDg1Dwa5HKhg7Rj9Due8hTDz5__ihyoUdGGKhtVqWo,1795
-flux_batch/version.py,sha256=EjicWR0QF-5SiGDBpWkctAT6keC05IyZrr61fi4OjQ4,642
+flux_batch/submit.py,sha256=tNJMDvnsxbAuXg_5TkB8HurzQ9I-Dot_gXXpIU5LtTA,2654
+flux_batch/version.py,sha256=h-6jDEwrQuFPdUK7HJbFYJISN7mt1MwBKK0UnnS1-cM,643
 flux_batch/logger/__init__.py,sha256=eDdpw_uppR5mPLHE39qT_haqMxu-2wniLlJZDigRC2k,52
 flux_batch/logger/generate.py,sha256=L9JyMY2oapp0ss7f7LGuihbLomzVJsMq7sByy9NhbZI,4017
 flux_batch/logger/logger.py,sha256=HKymVBNcoPdX87QWy69er5wUzHVeriiKp9p0bIYboUo,5927
 flux_batch/script/__init__.py,sha256=ZR-NCvI42RrAJTKzSmV4AcxlL-aoYh7mlQNPxF8Vzok,400
 flux_batch/script/save_logs.sh,sha256=HeapqvL0iR8aX7LtbwlcFTy19j5WxkQ-1M2J9epu-EU,515
-flux_batch/service/__init__.py,sha256=sFJ3r7XEw4vLcmadTxsxQq715bl_tFREOoUS4erimqQ,1116
+flux_batch/service/__init__.py,sha256=pD7RYpdSLZvYU4qwShYXA6UcaJy191fKtzqJL3uttEc,2724
 flux_batch/service/scribe.py,sha256=dY6geiLvXYIRcIzuP_naZscKgzX4Y5dPzxoWf9Wywg0,253
-flux_batch/service/scribe/__init__.py,sha256=BW4xb-QZvdrtbo73XB26MZozIbvDdnrgLzjB4GRtofU,39
+flux_batch/service/scribe/__init__.py,sha256=773-AzF_WRY6udG5nQSexf7Vga4K1YZLy1sfQAEd1uo,126
 flux_batch/service/scribe/__main__.py,sha256=3S0dyhkHk-bPT_Z0laNUg-HydrCFql4hPOV2ZNF5rO0,3777
 flux_batch/service/scribe/database.py,sha256=EB8OEMfNvfCplGaz-ZNMsfIpd305eP-mfCvSd-fg_k4,5626
 flux_batch/service/scribe/models.py,sha256=7lUrRosnQ2douFL_xD9GMYex4Z4lkN-CcBeWDhzmD8c,2668
-flux_batch/service/scribe/template.py,sha256=YwmAX5qkMwaj9FztftdgY3DjBaF9PVd7hDZebiD50fA,256
+flux_batch/service/scribe/template.py,sha256=yDsheid2FXeYr458loxB3HyZKSWpGQpM8iUQfnfOA-s,1336
 flux_batch/utils/__init__.py,sha256=CqMhw_mBfR0HBcHwv7LtFITq0J7LBV413VQE9xrz8ks,42
 flux_batch/utils/fileio.py,sha256=Elz8WkNkJ9B6x7WmCwiIBW0GgsRSSFCcbuJh7aqu2z4,4879
 flux_batch/utils/text.py,sha256=Ci1BqHs2IbOSn2o60zhLkT4kIA7CSNuGj8mdiGaDIGk,606
 flux_batch/utils/timer.py,sha256=_Weec7Wd5hWQ1r4ZHjownG4YdoIowpVqilXhvYFmIgA,491
-flux_batch-0.0.1.dist-info/LICENSE,sha256=AlyLB1m_z0CENCx1ob0PedLTTohtH2VLZhs2kfygrfc,1108
-flux_batch-0.0.1.dist-info/METADATA,sha256=G4iEbaFF1cRs12NvjbhWF76uj63MPnG4C6zF2T00e7w,5215
-flux_batch-0.0.1.dist-info/NOTICE,sha256=9CR93geVKl_4ZrJORbXN0fzkEM2y4DglWhY1hn9ZwQw,1167
-flux_batch-0.0.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-flux_batch-0.0.1.dist-info/entry_points.txt,sha256=ynoKpD82xn2V2sD-aZIQoq7NnfOu9VEKqW55Y1AoPGI,67
-flux_batch-0.0.1.dist-info/top_level.txt,sha256=jj8zAsZzMmbjiBISJL7lRtA37MSEAQYfObGLUncn9Lw,11
-flux_batch-0.0.1.dist-info/RECORD,,
+flux_batch-0.0.11.dist-info/LICENSE,sha256=AlyLB1m_z0CENCx1ob0PedLTTohtH2VLZhs2kfygrfc,1108
+flux_batch-0.0.11.dist-info/METADATA,sha256=A0LLvVm-RKGxSM26HMfHlOKe5OxF22O7APLNWfvIBj0,5352
+flux_batch-0.0.11.dist-info/NOTICE,sha256=9CR93geVKl_4ZrJORbXN0fzkEM2y4DglWhY1hn9ZwQw,1167
+flux_batch-0.0.11.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+flux_batch-0.0.11.dist-info/entry_points.txt,sha256=ynoKpD82xn2V2sD-aZIQoq7NnfOu9VEKqW55Y1AoPGI,67
+flux_batch-0.0.11.dist-info/top_level.txt,sha256=jj8zAsZzMmbjiBISJL7lRtA37MSEAQYfObGLUncn9Lw,11
+flux_batch-0.0.11.dist-info/RECORD,,