Merge pull request #4 from compspec/add-k8s-transformer

feat: simple transform utility

Author: vsoch

README.md

@@ -9,6 +9,14 @@ This library is primarily being used for development for the descriptive thrust
 
 ## Design
 
+### Simple
+
+We provide a simple translation layer between job specifications. We take the assumption that although each manager has many options, the actual options a user would use is a much smaller set, and it's relatively straight forward to translate (and have better accuracy).
+
+See [examples/transform](examples/transform) for an example.
+
+### Complex
+
 We want to:
 
 1. Generate software graphs for some cluster (fluxion JGF) (this is done with [compspec](https://github.com/compspec/compspec)

examples/fractale/README.md

@@ -174,3 +174,8 @@ Saving to "./examples/fractale/cluster-a-containment.png"
 ```
 
 <img src="./cluster-a-containment.svg">
+
+
+## Transform
+
+This mirrors the previous transform functionality for the jobspec library. It's a simple conversion, and no LLMs needed.

examples/transform/flux-to-kubernetes/README.md

@@ -0,0 +1,50 @@
+# Transform
+
+This is an example of doing a transformation between types. We do a simple mapping of parameters.
+To start testing, we will assume one node runs, and of the equivalent container. This way we can create a Job in Kubernetes without considering MPI networking.
+
+```bash
+# Print pretty
+fractale transform --to kubernetes --from flux ./flux_batch.sh --pretty
+
+# Print as raw yaml (to pipe to file)
+fractale transform --to kubernetes --from flux ./flux_batch.sh
+```
+```console
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: lammps
+spec:
+  activeDeadlineSeconds: 100
+  backoffLimit: 4
+  completions: 1
+  parallelism: 1
+  template:
+    metadata:
+      labels:
+        job-name: lammps
+    spec:
+      apiVersion: batch/v1
+      kind: Job
+      metadata:
+        name: lammps
+      spec:
+        backoffLimit: 0
+        template:
+          spec:
+            containers:
+            - args:
+              - lmp -v x 8 -v y 8 -v z 8 -in in.reaxc.hns -nocite
+              command:
+              - /bin/bash
+              - -c
+              image: ghcr.io/converged-computing/lammps-reax:ubuntu22.04
+              name: lammps
+              resources:
+                limits:
+                  cpu: '64'
+                requests:
+                  cpu: '64'
+            restartPolicy: Never
+```

examples/transform/flux-to-kubernetes/flux-batch.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+#FLUX: -N 1
+#FLUX: -n 64
+#FLUX: -t 100s
+#FLUX: -o cpu-affinity=per-task
+#FLUX: --queue=pbatch
+#FLUX: --setattr=container_image=ghcr.io/converged-computing/lammps-reax:ubuntu22.04
+#FLUX: --output=job.{id}.out
+#FLUX: --error=job.{id}.err
+#FLUX: --job-name=lammps
+
+lmp -v x 8 -v y 8 -v z 8 -in in.reaxc.hns -nocite

fractale/cli/__init__.py

@@ -71,6 +71,29 @@ def get_parser():
     )
     generate.add_argument("-c", "--cluster", help="cluster name")
 
+    # Transform jobspecs from flux to Kubernetes (starting specific)
+    transform = subparsers.add_parser(
+        "transform",
+        formatter_class=argparse.RawTextHelpFormatter,
+        description="transform a jobspec",
+    )
+    transform.add_argument(
+        "-t",
+        "--to",
+        dest="to_transformer",
+        help="transform into this jobspec",
+        default="kubernetes",
+    )
+    transform.add_argument(
+        "-f", "--from", dest="from_transformer", help="transform from this jobspec", default="flux"
+    )
+    transform.add_argument(
+        "--pretty",
+        help="pretty print in the terminal",
+        default=False,
+        action="store_true",
+    )
+
     # run.add_argument("-t", "--transform", help="transformer to use", default="flux")
     save = subparsers.add_parser(
         "save",
@@ -101,7 +124,7 @@ def get_parser():
     script.add_argument(
         "--transformer", help="transformer to use", default="flux", choices=["flux"]
     )
-    for cmd in [satisfy, script]:
+    for cmd in [satisfy, script, transform]:
         cmd.add_argument("jobspec", help="jobspec yaml or json file")
 
     for cmd in [satisfy, save, script]:
@@ -166,6 +189,8 @@ def help(return_code=0):
         from .script import main
     elif args.command == "save":
         from .save import main
+    elif args.command == "transform":
+        from .transform import main
     else:
         help(1)
     global registry

fractale/cli/script.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+
+import json
+import sys
+
+from rich import print, print_json
+from rich.json import JSON
+from rich.pretty import pprint
+
+from fractale.store import FractaleStore
+from fractale.subsystem import get_subsystem_solver
+from fractale.transformer import get_transformer
+
+
+def main(args, extra, **kwargs):
+    """
+    Save a cluster and subsystem graph.
+    """
+    store = FractaleStore(args.config_dir)
+
+    # Prepare selector and transformer
+    solver = get_subsystem_solver(store.clusters_root, args.solver)
+    # This is probably overloaded, but we need to be able to look up
+    # the templating logic for a subsystem
+    transformer = get_transformer(args.transformer, args.selector, solver)
+    matches = solver.satisfied(args.jobspec, return_results=True)
+    if matches.count == 0:
+        sys.exit(-1)
+
+    # Select the results to generate. This consolidates matches (which might include different nodes)
+    # into clusters and groups of subsystems
+    for script in transformer.render(matches, args.jobspec):
+        # Generate batch script or jobspec and print
+        pprint(script)

fractale/cli/transform.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+
+import os
+import sys
+
+import yaml
+from rich.pretty import pprint
+
+from fractale.transformer import get_transformer
+
+
+def main(args, extra, **kwargs):
+    """
+    Transform between jobspecs.
+    """
+    # The jobspec needs to exist as a file here
+    if not os.path.exists(args.jobspec):
+        sys.exit(f"JobSpec {args.jobspec} does not exist.")
+
+    # No selector or solver, just manual transform
+    from_transformer = get_transformer(args.from_transformer)
+    to_transformer = get_transformer(args.to_transformer)
+
+    # Parse the jobspec to transform from
+    normalized_jobspec = from_transformer.parse(args.jobspec)
+    final_jobspec = to_transformer.convert(normalized_jobspec)
+
+    if args.pretty:
+        pprint(final_jobspec, indent_guides=True)
+    elif args.to_transformer in ["kubernetes"]:
+        yaml.dump(final_jobspec, sys.stdout, sort_keys=True, default_flow_style=False)
+    else:
+        print(final_jobspec)

fractale/transformer/__init__.py

@@ -1,11 +1,13 @@
 from .flux import Transformer as FluxTransformer
+from .kubernetes import Transformer as KubernetesTransformer
 
 plugins = {
+    "kubernetes": KubernetesTransformer,
     "flux": FluxTransformer,
 }
 
 
-def get_transformer(name, selector, solver):
+def get_transformer(name, selector="random", solver=None):
     if name not in plugins:
         raise ValueError(f"{name} is not a valid transformer.")
     return plugins[name](selector, solver)

fractale/transformer/base.py

@@ -21,6 +21,18 @@ def __init__(self, selector, solver):
         self.selector = get_selector(selector)
         self.solver = solver
 
+    def parse(self, *args, **kwargs):
+        """
+        Parse converts the native jobspec to the standard JobSpec
+        """
+        raise NotImplementedError
+
+    def convert(self, *args, **kwargs):
+        """
+        Convert a normalized jobspec to the format here.
+        """
+        raise NotImplementedError
+
     def render(self, matches, jobspec):
         """
         Run the transformer:

fractale/transformer/common.py

@@ -0,0 +1,51 @@
+import sys
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+# Requires Python 3.8+ for dataclass
+
+
+@dataclass
+class JobSpec:
+    """
+    A scheduler-agnostic data structure for defining a computational job.
+    Version 2: Now includes accounting, priority, environment, and more constraints.
+    """
+
+    # Job Identity & Accounting
+    job_name: Optional[str] = None
+    account: Optional[str] = None
+
+    # What to Run
+    executable: Optional[str] = None
+    arguments: List[str] = field(default_factory=list)
+    container_image: Optional[str] = None
+    working_directory: Optional[str] = None
+
+    # Custom attributes or options
+    attrs: Optional[dict] = field(default_factory=dict)
+    options: Optional[dict] = field(default_factory=dict)
+
+    # Resource Requests ---
+    num_tasks: int = 1
+    num_nodes: int = 1
+    cpus_per_task: int = 1
+    mem_per_task: Optional[str] = None
+    gpus_per_task: int = 0
+
+    # Scheduling and Constraints
+    wall_time: Optional[str] = None
+    queue: Optional[str] = None
+    priority: Optional[int] = None
+    exclusive_access: bool = False
+    constraints: List[str] = field(default_factory=list)
+    begin_time: Optional[str] = None
+
+    # Environment and I/O
+    environment: Dict[str, str] = field(default_factory=dict)
+    input_file: Optional[str] = None
+    output_file: Optional[str] = None
+    error_file: Optional[str] = None
+
+    # Dependencies
+    depends_on: Optional[Union[str, List[str]]] = None