diff --git a/devtools/conda-recipe/meta.yaml b/devtools/conda-recipe/meta.yaml index efed2b8..84fe5b7 100644 --- a/devtools/conda-recipe/meta.yaml +++ b/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: git@github.com: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 diff --git a/setup.py b/setup.py index 7d346b8..1a58d27 100644 --- a/setup.py +++ b/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='michaelk@clin.au.dk', - 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="michaelk@clin.au.dk", + license="MIT", ) diff --git a/src/gwf_utilization/main.py b/src/gwf_utilization/main.py index dcafd37..366b73e 100644 --- a/src/gwf_utilization/main.py +++ b/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,33 +17,42 @@ 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. @@ -51,27 +60,44 @@ def utilization(obj, targets): 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)