From 0b9c5c4d8e97e0bf68d7508b069d360b94b5c641 Mon Sep 17 00:00:00 2001 From: Oleg Smirnov Date: Sun, 10 Sep 2023 21:54:04 +0200 Subject: [PATCH] moved backend.close() into end of run in prefernce to atexit --- stimela/backends/__init__.py | 39 +++++++++++-------------- stimela/backends/kube/__init__.py | 4 +++ stimela/backends/kube/infrastructure.py | 15 +++++----- stimela/commands/run.py | 4 +++ 4 files changed, 32 insertions(+), 30 deletions(-) diff --git a/stimela/backends/__init__.py b/stimela/backends/__init__.py index 799f9360..369cb7e3 100644 --- a/stimela/backends/__init__.py +++ b/stimela/backends/__init__.py @@ -4,6 +4,7 @@ from enum import Enum from omegaconf import ListConfig, OmegaConf from stimela.exceptions import BackendSpecificationError, BackendError +from stimela.stimelogging import log_exception from scabha.basetypes import EmptyDictDefault from .singularity import SingularityBackendOptions @@ -96,7 +97,7 @@ def resolve_image_name(backend: StimelaBackendOptions, image: 'stimela.kitchen.C return f"{image_name}:{version}" -def init_backends(backend_opts: StimelaBackendOptions, log: logging.Logger): +def _call_backends(backend_opts: StimelaBackendOptions, log: logging.Logger, method: str, desc: str, raise_exc: bool=True): selected = backend_opts.select or ['native'] if type(selected) is str: selected = [selected] @@ -106,32 +107,26 @@ def init_backends(backend_opts: StimelaBackendOptions, log: logging.Logger): opts = getattr(backend_opts, engine, None) if not opts or opts.enable: backend = get_backend(engine) - if backend: + func = backend and getattr(backend, method) + if func: try: - backend.init(backend_opts, log) + func(backend_opts, log) except BackendError as exc: - raise BackendError(f"error initializing {engine} backend", exc) - + exc1 = BackendError(f"error {desc} {engine} backend", exc) + if raise_exc: + raise exc1 from None + else: + log_exception(exc1, log=log) -def cleanup_backends(backend_opts: StimelaBackendOptions, log: logging.Logger): - selected = backend_opts.select or ['native'] - if type(selected) is str: - selected = [selected] - for engine in selected: - # check that backend has not been disabled - opts = getattr(backend_opts, engine, None) - if not opts or opts.enable: - backend = get_backend(engine) - if backend: - if hasattr(backend, 'cleanup'): - try: - backend.cleanup(backend_opts, log) - except BackendError as exc: - raise BackendError(f"error cleaning up {engine} backend", exc) from None - else: - log.info(f"nothing to clean up for {engine} backend") +def init_backends(backend_opts: StimelaBackendOptions, log: logging.Logger): + return _call_backends(backend_opts, log, "init", "initializing") + +def close_backends(backend_opts: StimelaBackendOptions, log: logging.Logger): + return _call_backends(backend_opts, log, "close", "closing") +def cleanup_backends(backend_opts: StimelaBackendOptions, log: logging.Logger): + return _call_backends(backend_opts, log, "cleanup", "cleaning up") ## commenting out for now -- will need to fix when we reactive the kube backend (and have tests for it) diff --git a/stimela/backends/kube/__init__.py b/stimela/backends/kube/__init__.py index fb53df54..ff95afc9 100644 --- a/stimela/backends/kube/__init__.py +++ b/stimela/backends/kube/__init__.py @@ -45,6 +45,10 @@ def init(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger): from . import infrastructure infrastructure.init(backend, log) +def close(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger): + from . import infrastructure + infrastructure.close(backend, log) + def cleanup(backend: 'stimela.backend.StimelaBackendOptions', log: logging.Logger): from . import infrastructure infrastructure.cleanup(backend, log) diff --git a/stimela/backends/kube/infrastructure.py b/stimela/backends/kube/infrastructure.py index 4b3da6bb..f095a505 100644 --- a/stimela/backends/kube/infrastructure.py +++ b/stimela/backends/kube/infrastructure.py @@ -29,8 +29,6 @@ # logger used for global kube messages klog: Optional[logging.Logger] = None -initialized = True - def _delete_pod(kube_api, podname, namespace, log, warn_not_found=True): log.info(f"deleting pod {podname}") try: @@ -55,7 +53,7 @@ def init(backend: StimelaBackendOptions, log: logging.Logger, cleanup: bool = Fa if cleanup: klog.info("cleaning up backend") else: - atexit.register(teardown, kube) + atexit.register(close, kube, klog) klog.info("initializing kube backend") if cleanup or kube.infrastructure.on_startup.report_pods or kube.infrastructure.on_startup.cleanup_pods: @@ -295,10 +293,10 @@ def delete_pvcs(kube: KubeBackendOptions, pvc_names, log: logging.Logger, force= terminating_pvcs[pvc.name] = name -def teardown(kube: KubeBackendOptions): - global initialized - if not initialized: - return +def close(backend: StimelaBackendOptions, log: logging.Logger): + kube = backend.kube + klog.info("closing kube backend") + # release PVCs delete_pvcs(kube, list(active_pvcs.keys()), log=klog, session=True, step=True, refresh=False) @@ -325,4 +323,5 @@ def teardown(kube: KubeBackendOptions): for podname in running_pods: _delete_pod(kube_api, podname, kube.namespace, klog) - initialized = False + atexit.unregister(close) + diff --git a/stimela/commands/run.py b/stimela/commands/run.py index 1ed4ae01..da69f7e7 100644 --- a/stimela/commands/run.py +++ b/stimela/commands/run.py @@ -299,6 +299,8 @@ def elapsed(): try: outputs = outer_step.run(backend=stimela.CONFIG.opts.backend) except Exception as exc: + stimela.backends.close_backends(backend, log) + task_stats.save_profiling_stats(outer_step.log, print_depth=profile if profile is not None else stimela.CONFIG.opts.profile.print_depth, unroll_loops=stimela.CONFIG.opts.profile.unroll_loops) @@ -321,6 +323,8 @@ def elapsed(): else: outer_step.log.info(f"run successful after {elapsed()}") + stimela.backends.close_backends(backend, log) + task_stats.save_profiling_stats(outer_step.log, print_depth=profile if profile is not None else stimela.CONFIG.opts.profile.print_depth, unroll_loops=stimela.CONFIG.opts.profile.unroll_loops)