From ed06a83157542d63948cc0cb55ce33e9c818127d Mon Sep 17 00:00:00 2001 From: Thibault Guittet Date: Wed, 4 Dec 2024 17:01:47 -0600 Subject: [PATCH] Add `num_cpu` and `available_memory` steps --- .../arcaflow_plugin_utilities.py | 49 +++++++++++++++++++ poetry.lock | 34 ++++++++++++- pyproject.toml | 1 + 3 files changed, 82 insertions(+), 2 deletions(-) diff --git a/arcaflow_plugin_utilities/arcaflow_plugin_utilities.py b/arcaflow_plugin_utilities/arcaflow_plugin_utilities.py index 6fcc6ab..3fefbf4 100644 --- a/arcaflow_plugin_utilities/arcaflow_plugin_utilities.py +++ b/arcaflow_plugin_utilities/arcaflow_plugin_utilities.py @@ -2,6 +2,8 @@ import sys import uuid +import psutil +import os from datetime import datetime, timezone from time import time, sleep import typing @@ -50,6 +52,24 @@ class SuccessOutputWait: ] +@dataclass +class SuccessOutputNumCpu: + num_cpu: typing.Annotated[ + int, + schema.name("num cpu"), + schema.description("Number of logical CPUs"), + ] + + +@dataclass +class SuccessOutputAvailableMemory: + available_memory: typing.Annotated[ + int, + schema.name("available memory"), + schema.description("Amount of available memory (in bytes)"), + ] + + @dataclass class ErrorOutput: error: typing.Annotated[ @@ -102,6 +122,33 @@ def wait( return "success", SuccessOutputWait((end - start) * 1000) +@plugin.step( + id="num_cpu", + name="Num Cpu", + description=("Number of logical CPUs"), + outputs={"success": SuccessOutputNumCpu, "error": ErrorOutput}, +) +def num_cpu( + params: InputParams, +) -> typing.Tuple[str, typing.Union[SuccessOutputNumCpu, ErrorOutput]]: + cpu_count = os.cpu_count() + if cpu_count is None: + return "error", ErrorOutput("Number of logical CPUs in the system couldn't be determined") + return "success", SuccessOutputNumCpu(cpu_count) + + +@plugin.step( + id="available_memory", + name="Available Memory", + description=("Amount of available memory (in bytes)"), + outputs={"success": SuccessOutputAvailableMemory}, +) +def available_memory( + params: InputParams, +) -> typing.Tuple[str, SuccessOutputAvailableMemory]: + return "success", SuccessOutputAvailableMemory(psutil.virtual_memory().available) + + if __name__ == "__main__": sys.exit( plugin.run( @@ -109,6 +156,8 @@ def wait( generate_uuid, generate_timestamp, wait, + num_cpu, + available_memory, ) ) ) diff --git a/poetry.lock b/poetry.lock index a7f6b42..ab30905 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. [[package]] name = "arcaflow-plugin-sdk" @@ -340,6 +340,36 @@ docs = ["furo (>=2024.8.6)", "proselint (>=0.14)", "sphinx (>=8.0.2)", "sphinx-a test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=8.3.2)", "pytest-cov (>=5)", "pytest-mock (>=3.14)"] type = ["mypy (>=1.11.2)"] +[[package]] +name = "psutil" +version = "6.1.0" +description = "Cross-platform lib for process and system monitoring in Python." +optional = false +python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,>=2.7" +files = [ + {file = "psutil-6.1.0-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:ff34df86226c0227c52f38b919213157588a678d049688eded74c76c8ba4a5d0"}, + {file = "psutil-6.1.0-cp27-cp27m-manylinux2010_i686.whl", hash = "sha256:c0e0c00aa18ca2d3b2b991643b799a15fc8f0563d2ebb6040f64ce8dc027b942"}, + {file = "psutil-6.1.0-cp27-cp27m-manylinux2010_x86_64.whl", hash = "sha256:000d1d1ebd634b4efb383f4034437384e44a6d455260aaee2eca1e9c1b55f047"}, + {file = "psutil-6.1.0-cp27-cp27mu-manylinux2010_i686.whl", hash = "sha256:5cd2bcdc75b452ba2e10f0e8ecc0b57b827dd5d7aaffbc6821b2a9a242823a76"}, + {file = "psutil-6.1.0-cp27-cp27mu-manylinux2010_x86_64.whl", hash = "sha256:045f00a43c737f960d273a83973b2511430d61f283a44c96bf13a6e829ba8fdc"}, + {file = "psutil-6.1.0-cp27-none-win32.whl", hash = "sha256:9118f27452b70bb1d9ab3198c1f626c2499384935aaf55388211ad982611407e"}, + {file = "psutil-6.1.0-cp27-none-win_amd64.whl", hash = "sha256:a8506f6119cff7015678e2bce904a4da21025cc70ad283a53b099e7620061d85"}, + {file = "psutil-6.1.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:6e2dcd475ce8b80522e51d923d10c7871e45f20918e027ab682f94f1c6351688"}, + {file = "psutil-6.1.0-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:0895b8414afafc526712c498bd9de2b063deaac4021a3b3c34566283464aff8e"}, + {file = "psutil-6.1.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9dcbfce5d89f1d1f2546a2090f4fcf87c7f669d1d90aacb7d7582addece9fb38"}, + {file = "psutil-6.1.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:498c6979f9c6637ebc3a73b3f87f9eb1ec24e1ce53a7c5173b8508981614a90b"}, + {file = "psutil-6.1.0-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d905186d647b16755a800e7263d43df08b790d709d575105d419f8b6ef65423a"}, + {file = "psutil-6.1.0-cp36-cp36m-win32.whl", hash = "sha256:6d3fbbc8d23fcdcb500d2c9f94e07b1342df8ed71b948a2649b5cb060a7c94ca"}, + {file = "psutil-6.1.0-cp36-cp36m-win_amd64.whl", hash = "sha256:1209036fbd0421afde505a4879dee3b2fd7b1e14fee81c0069807adcbbcca747"}, + {file = "psutil-6.1.0-cp37-abi3-win32.whl", hash = "sha256:1ad45a1f5d0b608253b11508f80940985d1d0c8f6111b5cb637533a0e6ddc13e"}, + {file = "psutil-6.1.0-cp37-abi3-win_amd64.whl", hash = "sha256:a8fb3752b491d246034fa4d279ff076501588ce8cbcdbb62c32fd7a377d996be"}, + {file = "psutil-6.1.0.tar.gz", hash = "sha256:353815f59a7f64cdaca1c0307ee13558a0512f6db064e92fe833784f08539c7a"}, +] + +[package.extras] +dev = ["black", "check-manifest", "coverage", "packaging", "pylint", "pyperf", "pypinfo", "pytest-cov", "requests", "rstcheck", "ruff", "sphinx", "sphinx_rtd_theme", "toml-sort", "twine", "virtualenv", "wheel"] +test = ["pytest", "pytest-xdist", "setuptools"] + [[package]] name = "pyflakes" version = "3.2.0" @@ -478,4 +508,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = "^3.9" -content-hash = "58966cc9286900d877ec061c4b21c3854f27d40dae7aeaccfe1bd781ca8082e0" +content-hash = "348705b61dc586490f1a992bb0d6108eee8d8da04a63e64b5b90934c8bbdf626" diff --git a/pyproject.toml b/pyproject.toml index 50ee53a..150e76a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ packages = [ [tool.poetry.dependencies] python = "^3.9" arcaflow-plugin-sdk = "^0.14.0" +psutil = "^6.1.0" [tool.poetry.dev-dependencies] docformatter = "^1.5.0"