Support for new gwf

devtools/conda-recipe/meta.yaml

@@ -1,9 +1,9 @@
 package:
   name: gwf-utilization
-  version: '0.1.7'
+  version: '0.2'
 
 source:
-  git_rev: v0.1.7
+  git_rev: v0.2
   git_url: [email protected]:micknudsen/gwf-utilization.git
 
 build:
@@ -17,7 +17,7 @@ requirements:
     - python >=3.7
   run:
     - click
-    - gwf >=1.7.1
+    - gwf >=2
     - python >=3.7
     - texttable >=1.4.0
 

setup.py

@@ -1,29 +1,15 @@
 from setuptools import setup, find_packages
 
 setup(
-
-    name='gwf-utilization',
-    version='0.1.7',
-
-    packages=find_packages('src'),
-    package_dir={'': 'src'},
-
-    test_suite='tests',
-
-    entry_points={
-        'gwf.plugins': ['utilization = gwf_utilization.main:utilization']
-    },
-
-    python_requires='>=3.7',
-
-    install_requires=[
-        'click',
-        'gwf>=1.7.1',
-        'texttable>=1.4.0'
-    ],
-
-    author='Michael Knudsen',
-    author_email='[email protected]',
-    license='MIT'
-
+    name="gwf-utilization",
+    version="0.1.7",
+    packages=find_packages("src"),
+    package_dir={"": "src"},
+    test_suite="tests",
+    entry_points={"gwf.plugins": ["utilization = gwf_utilization.main:utilization"]},
+    python_requires=">=3.7",
+    install_requires=["click", "gwf>=2", "texttable>=1.4.0"],
+    author="Michael Knudsen",
+    author_email="[email protected]",
+    license="MIT",
 )

src/gwf_utilization/main.py

@@ -1,12 +1,12 @@
 import math
 import click
+import os.path
+import json
 
 from texttable import Texttable
 
 from gwf import Workflow
-from gwf.core import Graph
-from gwf.backends import Backend
-from gwf.backends.slurm import SlurmBackend
+from gwf.core import CachedFilesystem, Graph, pass_context
 from gwf.exceptions import GWFError
 from gwf.filtering import filter_names
 
@@ -17,61 +17,87 @@ def pretty_time(time_in_seconds):
     minutes, seconds = divmod(time_in_seconds, 60)
     hours, minutes = divmod(minutes, 60)
     days, hours = divmod(hours, 24)
-    result = f'{hours:02d}:{minutes:02d}:{seconds:02d}'
+    result = f"{hours:02d}:{minutes:02d}:{seconds:02d}"
     if days:
-        result = f'{days}-{result}'
+        result = f"{days}-{result}"
     return result
 
 
 def pretty_size(size_in_bytes):
     if size_in_bytes == 0:
-        return '0 B'
-    size_name = ('B', 'KB', 'MB', 'GB', 'TB', 'PB')
+        return "0 B"
+    size_name = ("B", "KB", "MB", "GB", "TB", "PB")
     exponent = int(math.floor(math.log(size_in_bytes, 1024)))
     multiplier = math.pow(1024, exponent)
     result = round(size_in_bytes / multiplier, 2)
-    return f'{result} {size_name[exponent]}'
+    return f"{result} {size_name[exponent]}"
+
+
+def load_tracked_jobs(working_dir):
+    try:
+        with open(
+            os.path.join(working_dir, ".gwf", "slurm-backend-tracked.json")
+        ) as state_file:
+            return json.load(state_file)
+    except FileNotFoundError:
+        return {}
 
 
 @click.command()
-@click.argument('targets', nargs=-1)
-@click.pass_obj
-def utilization(obj, targets):
-    # This plugin only works for SlurmBackend
-    backend_cls = Backend.from_config(obj)
-    if not issubclass(backend_cls, SlurmBackend):
-        raise GWFError('Utilization plugin only works for Slurm backend!')
+@click.argument("targets", nargs=-1)
+@pass_context
+def utilization(ctx, targets):
+    # This plugin only works for the slurm backend.
+    if ctx.backend != "slurm":
+        raise GWFError("Utilization plugin only works for Slurm backend!")
 
-    workflow = Workflow.from_config(obj)
-    graph = Graph.from_targets(workflow.targets)
+    workflow = Workflow.from_context(ctx)
+    graph = Graph.from_targets(workflow.targets, fs=CachedFilesystem())
 
     # If user specified list of targets, only report utilization for these.
     # Otherwise, report utilization for all targets.
     matches = graph.targets.values()
     if targets:
         matches = filter_names(matches, targets)
 
-    with backend_cls() as backend:
-        job_ids = []
-        for target in matches:
-            try:
-                job_ids.append(backend.get_job_id(target))
-            except KeyError:
-                pass
-
-    rows = [['Target', 'Cores', 'Walltime Alloc', 'Walltime Used', 'Memory Alloc', 'Memory Used', 'CPU Time Alloc', 'CPU Time Used', 'Walltime %', 'Memory %', 'CPU %']]
+    tracked_jobs = load_tracked_jobs(ctx.working_dir)
+    job_ids = [
+        tracked_jobs[target.name]
+        for target in matches
+        if tracked_jobs.get(target.name) is not None
+    ]
+
+    rows = [
+        [
+            "Target",
+            "Cores",
+            "Walltime Alloc",
+            "Walltime Used",
+            "Memory Alloc",
+            "Memory Used",
+            "CPU Time Alloc",
+            "CPU Time Used",
+            "Walltime %",
+            "Memory %",
+            "CPU %",
+        ]
+    ]
     for job in get_jobs(job_ids):
-        rows.append([job.name,
-                     job.allocated_cores,
-                     pretty_time(job.allocated_time_per_core),
-                     pretty_time(job.used_walltime),
-                     pretty_size(job.allocated_memory),
-                     pretty_size(job.used_memory),
-                     pretty_time(job.allocated_cpu_time),
-                     pretty_time(job.used_cpu_time),
-                     str(format(job.walltime_utilization, '.1f')),
-                     format(job.memory_utilization, '.1f'),
-                     format(job.cpu_utilization, '.1f')])
+        rows.append(
+            (
+                job.name,
+                job.allocated_cores,
+                pretty_time(job.allocated_time_per_core),
+                pretty_time(job.used_walltime),
+                pretty_size(job.allocated_memory),
+                pretty_size(job.used_memory),
+                pretty_time(job.allocated_cpu_time),
+                pretty_time(job.used_cpu_time),
+                str(format(job.walltime_utilization, ".1f")),
+                format(job.memory_utilization, ".1f"),
+                format(job.cpu_utilization, ".1f"),
+            )
+        )
 
     table = Texttable()
 
@@ -80,9 +106,9 @@ def utilization(obj, targets):
     ncols = len(rows[0])
 
     table.set_max_width(0)
-    table.set_header_align('l' * ncols)
-    table.set_cols_align(['l'] + (ncols - 1) * ['r'])
-    table.set_cols_dtype(['t'] * ncols)
+    table.set_header_align("l" * ncols)
+    table.set_cols_align(["l"] + (ncols - 1) * ["r"])
+    table.set_cols_dtype(["t"] * ncols)
 
     table.add_rows(rows)