feat: graph solver backend

Signed-off-by: vsoch <[email protected]>

Author: vsoch

README.md

@@ -18,6 +18,18 @@ We want to:
 
 See [examples/fractale](examples/fractale) for a detailed walk-through of the above.
 
+For graph tool:
+
+```bash
+conda install -c conda-forge graph-tool
+```
+
+## Questions
+
+- Should other subsystem types have edges? How used?
+- Should we try to map them to nodes in the graph or use another means (or assume global across cluster nodes as we do now)?
+- Can we simplify spack subsystem graph (it's really big...)
+
 <!-- ⭐️ [Documentation](https://compspec.github.io/fractale) ⭐️ -->
 
 ## License

examples/fractale/README.md

@@ -28,7 +28,7 @@ fractale generate --cluster A spack /home/vanessa/Desktop/Code/spack
 Satisfy asks two questions:
 
 1. Which clusters have the subsystem resources that I need?
-2. Which clusters have the job resources that I need?
+2. Which clusters have the job resources that I need (containment subsystem)?
 
 This is the step where we want to say "Run gromacs on 2-4 nodes with these requirements." Since we haven't formalized a way to do that, I'm going to start with a flux jobspec, and then add attributes that can be used to search our subsystems. For example, I generated [software-gromacs.json](software-gromacs.json) with:
 
@@ -52,6 +52,40 @@ fractale satisfy ./examples/fractale/software-curl.yaml
 fractale satisfy ./examples/fractale/software-curl.json
 ```
 
+Try the graph backend, install [graph-tool](https://graph-tool.skewed.de/installation.html) then:
+
+```bash
+export PYTHONPATH=/home/vanessa/anaconda3/lib/python3.12/site-packages
+fractale satisfy --solver graph ./examples/fractale/software-curl.yaml
+fractale satisfy --solver graph ./examples/fractale/software-curl.json
+```
+```console
+=> 🍇 Loading cluster "a" subsystem "containment"
+=> 🍇 Loading cluster "a" subsystem "modules"
+=> 🍇 Loading cluster "a" subsystem "spack"
+    => Exploring cluster "a" containment subsystem
+          (1/1) satisfied resource core 
+Cluster "a" is a match
+```
+
+Here is for a nested slot:
+
+```bash
+fractale satisfy --solver graph ./examples/fractale/software-nested-slot.yaml 
+```
+```console
+=> 🍇 Loading cluster "a" subsystem "containment"
+=> 🍇 Loading cluster "a" subsystem "modules"
+=> 🍇 Loading cluster "a" subsystem "spack"
+    => Exploring cluster "a" containment subsystem
+          (1/1) satisfied resource socket 
+          (1/4) found resource core 
+          (2/4) found resource core 
+          (3/4) found resource core 
+          (4/4) satisfied resource core 
+Cluster "a" is a match
+```
+
 By default, the above assumes subsystems located in the fractale home. If you want to adjust that, set `fractale --config-dir=<path> satisfy...` to adjust that (and note you will need to have generated the tree here. What we basically do with satisfy is build a database with tables for:
 
 - clusters
@@ -77,5 +111,27 @@ Here is a jobspec that can't be satisfied because we ask for too many resources
 ```bash
 fractale satisfy ./examples/fractale/jobspec-containment-unsatisfied.yaml 
 ```
-
+```console
+=> 🍇 Loading cluster "a" subsystem "containment"
+=> 🍇 Loading cluster "a" subsystem "modules"
+=> 🍇 Loading cluster "a" subsystem "spack"
+       (2) SELECT * from subsystems WHERE type = 'software';
+=> No Matches due to containment
+```
 We likely want to have a more structured query syntax that can handle AND, OR, and other specifics. The actual search should remain general to support any generic key/value pair of attributes. My database structure and queries are also bad. 
+
+## Save
+
+We can save an image of our subystem for a cluster. E.g.,
+
+```bash
+$ fractale save a --out ./examples/fractale/cluster-a-containment.png
+```
+```console
+=> 🍇 Loading cluster "a" subsystem "containment"
+=> 🍇 Loading cluster "a" subsystem "modules"
+=> 🍇 Loading cluster "a" subsystem "spack"
+Saving to "./examples/fractale/cluster-a-containment.png"
+```
+
+<img src="./cluster-a-containment.svg">

examples/fractale/cluster-a-containment-graph.svg

@@ -0,0 +1,24 @@
+version: 1
+resources:
+- type: slot
+  count: 2
+  with:
+  - type: socket
+    count: 1
+    with:
+      - type: core
+        count: 4
+  label: task
+tasks:
+- command:
+  - gmx
+  slot: task
+  count:
+    per_slot: 1
+attributes:
+  system:
+    duration: 0
+    requires:
+      software:
+        - name: curl
+          type: binary

examples/fractale/cluster-a-containment.png

@@ -72,17 +72,27 @@ def get_parser():
     generate.add_argument("-c", "--cluster", help="cluster name")
 
     # run.add_argument("-t", "--transform", help="transformer to use", default="flux")
+    save = subparsers.add_parser(
+        "save",
+        formatter_class=argparse.RawTextHelpFormatter,
+        description="save a picture of a subsystem graph",
+    )
+    save.add_argument("cluster", help="cluster to save")
+    save.add_argument(
+        "--subsystem", help="cluster to save (defaults to containment)", default="containment"
+    )
+    save.add_argument("--out", help="output file name")
 
     # This does just the user space subsystem match
     satisfy = subparsers.add_parser(
         "satisfy",
         formatter_class=argparse.RawTextHelpFormatter,
         description="determine clusters that satisfy a jobspec based on user subsystems",
     )
-    for cmd in [satisfy]:
-        cmd.add_argument("jobspec", help="jobspec yaml or json file")
+    satisfy.add_argument("jobspec", help="jobspec yaml or json file")
+    for cmd in [satisfy, save]:
         cmd.add_argument(
-            "--backend",
+            "--solver",
             help="subsystem solved backend",
             default=defaults.solver_backend_default,
             choices=defaults.solver_backends,
@@ -138,6 +148,8 @@ def help(return_code=0):
         from .generate_subsystem import main
     elif args.command == "satisfy":
         from .satisfy import main
+    elif args.command == "save":
+        from .save import main
     else:
         help(1)
     global registry

examples/fractale/software-nested-slot.yaml

@@ -12,6 +12,6 @@ def main(args, extra, **kwargs):
     This is a fairly simple (flat) check.
     """
     store = FractaleStore(args.config_dir)
-    solver = get_subsystem_solver(store.clusters_root, args.backend)
+    solver = get_subsystem_solver(store.clusters_root, args.solver)
     is_satisfied = solver.satisfied(args.jobspec)
     sys.exit(0 if is_satisfied else -1)

fractale/cli/__init__.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+
+import sys
+
+from fractale.store import FractaleStore
+from fractale.subsystem import get_subsystem_solver
+
+
+def main(args, extra, **kwargs):
+    """
+    Save a cluster and subsystem graph.
+    """
+    store = FractaleStore(args.config_dir)
+    solver = get_subsystem_solver(store.clusters_root, args.solver)
+    outfile = args.out
+    if not outfile:
+        outfile = f"cluster-{args.cluster}-{args.subsystem}-{args.solver}.svg"
+    solver.save(args.cluster, args.subsystem, outfile)

fractale/cli/satisfy.py

@@ -2,4 +2,4 @@
 valid_settings = {"sharedfs", "stage"}
 sharedfs = True
 solver_backends = ["database", "graph"]
-solver_backend_default = "database"
+solver_backend_default = "graph"

fractale/cli/save.py

@@ -1,17 +1,27 @@
 import copy
+from contextlib import contextmanager
+
+import fractale.utils as utils
 
 
 def flatten_jobspec_resources(jobspec):
     """
     Given a jobspec, turn the required resources into a flattened version.
+
+    This is intended for subsystem matching that doesn't assume any structure.
+    E.g., the database solved backend, which was first just for prototyping.
     """
     resources = {}
     resource_list = copy.deepcopy(jobspec["resources"])
     multiplier = 1
+
+    # Resource lists are nested, under "with"
     while resource_list:
         requires = resource_list.pop(0)
         resource_type = requires["type"]
         resource_count = requires.get("count")
+
+        # The slot is a marker for the set we care about matching
         if resource_type == "slot":
             multiplier = resource_count or 1
         else:
@@ -20,3 +30,146 @@ def flatten_jobspec_resources(jobspec):
             resources[resource_type] += resource_count * multiplier
         resource_list += requires.get("with") or []
     return resources
+
+
+def extract_slot(jobspec):
+    """
+    Given a jobspec, parse into slot and return counts and requirements.
+    This is intended for the containment subsystem.
+    """
+    js = utils.load_jobspec(jobspec)
+    slots = []
+
+    # Tasks will have per slot (this is a list of one item, required)
+    # In the future Flux could theoretically support more than one task
+    # (and slot) in which case we would need to select the label
+    task = js["tasks"][0]
+    slot_name = task["slot"]
+
+    # I'm not sure I've ever seen anything other than "per_slot"
+    # Might as well raise an error and alert if that happens :)
+    slot_count = task.get("count", {})
+    if "per_slot" not in slot_count:
+        raise ValueError(f"Unexpected value for tasks->slot->count: {slot_count}")
+
+    # Assume a default of one per slot
+    slot_count = slot_count.get("per_slot") or 1
+
+    def check_resource(resource, is_slot=False):
+        """
+        Recursive function to dig into resources under "with.
+        Once we find the slot, we save the entire thing (with nesting)
+        as a Slot class, which better exposes the counts, etc.
+        """
+        if "with" in resource:
+            for item in resource["with"]:
+                if is_slot:
+                    new_slot = copy.deepcopy(item)
+                    total = item["count"] * slot_count
+                    slots.append(Slot(new_slot, total=total))
+
+                # Again, we can eventually support multiple slots by
+                # checking the label. Flux right now only expects one.
+                check_resource(item, item["type"] == "slot")
+
+    # Kick it off, y'all
+    for resource in js["resources"]:
+        check_resource(resource, resource["type"] == "slot")
+    return slots
+
+
+class Slot:
+    """
+    A slot is like a wedge we stick in the containment graph, and say
+    "We are defining our needs for the resources below this wedge, and we need
+    this many of this wedge." Since we need to evaluate the resource counts
+    dynamically, we provide a context-based function to do that. It will restore
+    back to the original state when the context is left. The expected usage is:
+
+    # Temporary view of the slot to evaluate with (make changes to)
+    with slot.evaluate() as evaluator:
+        # This is the current spot we are at in the slot
+        requires = evaluator.next_requirement()
+        v_type, count = next(requires)
+
+    Check for a StopIteration exception to know when we are done.
+    """
+
+    def __init__(self, spec, total=1):
+        # This is a nested structure, e.g.,
+        # {'type': 'socket', 'count': 1, 'with': [{'type': 'core', 'count': 4}]}
+        self.spec = spec
+        self.total = total
+        # This holds state for what we find during an evaluation
+        self._found = {}
+
+    @property
+    def start_type(self):
+        """
+        The start type identifies the top of the slot
+        """
+        return self.spec["type"]
+
+    @contextmanager
+    def evaluate(self):
+        """
+        Yield a temporary copy of the spec in case it is mutated.
+        Also init and reset found, so we can evaluate multiple spots
+        for satisfying the same slot instance.
+        """
+        spec = copy.deepcopy(self.spec)
+        self._found = {}
+        try:
+            yield self
+        # Restore the original state, you mangy animal
+        finally:
+            self.spec = spec
+            self._found = {}
+
+    def next_requirement(self):
+        """
+        Yield requirements. We are strict for now, requiring that we see all levels.
+        """
+        resources = [copy.deepcopy(self.spec)]
+        while resources:
+            resource = resources.pop(0)
+            yield resource["type"], resource.get("count") or 1
+            if "with" in resource:
+                resources += resource["with"]
+
+    def count(self, v_type):
+        """
+        Get the count for a resource type that has been found.
+        """
+        return self._found.get(v_type) or 0
+
+    def found(self, v_type, count=1, needed=None):
+        """
+        Indicate that a type was found.
+        """
+        if v_type not in self._found:
+            self._found[v_type] = 0
+        self._found[v_type] += 1
+
+        # This tells the caller if we are done with the type.
+        if needed is not None and self._found[v_type] == needed:
+            return True
+        return False
+
+    def satisfied(self):
+        """
+        Determine if the slot is satisifed, meaning all needed counts
+        are <= 0. This should be run in the context of evaluate.
+        """
+        # These are all the requirements
+        needed = {}
+        for requires_type, requires_count in self.next_requirement():
+            needed[requires_type] = requires_count
+
+        updated = copy.deepcopy(needed)
+        for found_type, found_count in self._found.items():
+            if found_count >= needed[found_type]:
+                del updated[found_type]
+
+        # If updated is empty, we got all needed
+        return not updated