From e734a052f6b95abec1f2e3a0ba14ac948074cf9d Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 3 Mar 2025 01:16:02 +0100 Subject: [PATCH 01/44] snap --- src/radical/utils/flux.py | 79 +++++++++++++++++++++------------------ 1 file changed, 43 insertions(+), 36 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 92980e0b3..4eafc496f 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -7,7 +7,8 @@ import json import shlex -from typing import Optional, List, Dict, Any, Callable +from typing import Optional, List, Dict, Any, Callable +from functools import partial import threading as mt import subprocess as sp @@ -312,8 +313,8 @@ def __init__(self, uid:str = None) -> None: if uid: self._uid = uid else : self._uid = generate_id('flux.%(item_counter)04d', ID_CUSTOM) - self._log = Logger(self._uid, ns='radical.utils') - self._prof = Profiler(self._uid, ns='radical.utils') + self._log = Logger ('radical.utils', ns='radical.utils', level='DEBUG') + self._prof = Profiler('radical.utils', ns='radical.utils') self._lock = mt.RLock() @@ -505,41 +506,47 @@ def submit_jobs(self, if not self._uri: raise RuntimeError('FluxHelper is not connected') - assert self._exe, 'no executor' - - futures = list() - for spec in specs: - jobspec = json.dumps(spec) - fut = self._flux_job.submit_async(self._handle, jobspec) - futures.append(fut) + assert self._exe - ids = list() - for fut in futures: - flux_id = fut.get_id() - ids.append(flux_id) - self._log.debug('submit: %s', flux_id) - - if cb: - def app_cb(fut, event): - try: - cb(flux_id, event) - except: - self._log.exception('app cb failed') + def app_cb(flux_id, exe_fut, event): + try : cb(flux_id, event) + except: self._log.exception('app cb failed for %s [%s]', + flux_id, event) - for ev in [ - 'submit', - 'alloc', - 'start', - 'finish', - 'release', - # 'free', - # 'clean', - 'exception', - ]: - fut.add_event_callback(ev, app_cb) - - self._log.debug('submitted: %s', ids) - return ids + futures = list() + def id_cb(fut): + flux_id = fut.jobid() + idx = fut.ru_idx + for ev in ['alloc', 'start', 'finish', 'release', 'exception']: + # 'submit', 'free', 'clean', + tmp_cb = partial(app_cb, flux_id) + fut.add_event_callback(ev, tmp_cb) + futures.append([flux_id, idx, fut]) + self._log.debug('got flux id: %s: %s', idx, flux_id) + + for idx, spec in enumerate(specs): + jobspec = json.dumps(spec) + fut = self._exe.submit(jobspec) + fut.ru_idx = idx + self._log.debug('submitted : %s', idx) + fut.add_jobid_callback(id_cb) + + # wait until we saw all jobid callbacks (assume 10 tasks/sec) + timeout = len(specs) / 10 + timeout = max(10, timeout) + start = time.time() + self._log.debug('wait %.2fsec for %d flux IDs', timeout, len(specs)) + while len(futures) < len(specs): + time.sleep(0.1) + self._log.debug('wait %s / %s', len(futures), len(specs)) + if time.time() - start > timeout: + raise RuntimeError('timeout on job submission') + self._log.info('got %d flux IDs', len(futures)) + + # get flux_ids sorted by submission order (idx) + flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] + + return flux_ids # -------------------------------------------------------------------------- From ac30f57073e02ec4843042b7cbd757c2de52fb59 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 3 Mar 2025 23:43:29 +0100 Subject: [PATCH 02/44] add simple yappi resource manager --- src/radical/utils/__init__.py | 1 + src/radical/utils/profile.py | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/src/radical/utils/__init__.py b/src/radical/utils/__init__.py index 00ac1910c..1f8c861fc 100644 --- a/src/radical/utils/__init__.py +++ b/src/radical/utils/__init__.py @@ -75,6 +75,7 @@ from .profile import read_profiles, combine_profiles, clean_profile from .profile import TIME, EVENT, COMP, TID, UID, STATE, MSG, ENTITY from .profile import PROF_KEY_MAX +from .profile import Yappi from .json_io import read_json, read_json_str, write_json from .json_io import parse_json, parse_json_str, dumps_json diff --git a/src/radical/utils/profile.py b/src/radical/utils/profile.py index e18e759bb..a0f987321 100644 --- a/src/radical/utils/profile.py +++ b/src/radical/utils/profile.py @@ -12,6 +12,8 @@ from .threads import get_thread_name as ru_get_thread_name from .config import DefaultConfig from .atfork import atfork +from .shell import sh_callout +from .modules import import_module # ------------------------------------------------------------------------------ @@ -843,5 +845,35 @@ def event_to_label(event): return event[EVENT] +# ------------------------------------------------------------------------------ +# +class Yappi(object): + + def __init__(self, name, method='wall'): + self._yappi = import_module('yappi') + self._yappi.set_clock_type(method) + self._name = name + + def __enter__(self): + if self._yappi: + self._yappi.start(builtins=True) + + def __exit__(self, etype, value, traceback): + + if self._yappi: + + # self._yappi.get_thread_stats().print_all() + + fstats = self._yappi.get_func_stats() + pstats = self._yappi.convert2pstats(fstats) + pstats.dump_stats('%s.pstats' % self._name) + + cmd = 'gprof2dot -e 0.00 -n 0.12 --skew=0.3' + cmd += ' --node-label="total-time" -f pstats %s.pstats' % self._name + cmd += ' | dot -Tpng -o %s.png' % self._name + + sh_callout(cmd, shell=True) + + # ------------------------------------------------------------------------------ From 0b4a88ab2fd65c929cc649dcf957e248934bd512 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 19 Mar 2025 11:46:15 +0100 Subject: [PATCH 03/44] snap --- src/radical/utils/flux.py | 3 ++- src/radical/utils/ids.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 4eafc496f..ac45a0b05 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -105,7 +105,8 @@ def _watch(self) -> None: time.sleep(1) - _, err, ret = sh_callout('flux ping -c 1 kvs') + out, err, ret = sh_callout('flux ping -c 1 kvs') + self._log.debug('flux ping:%s', out) if ret: self._log.error('flux watcher err: %s', err) break diff --git a/src/radical/utils/ids.py b/src/radical/utils/ids.py index ecec34c76..15f730563 100644 --- a/src/radical/utils/ids.py +++ b/src/radical/utils/ids.py @@ -198,6 +198,10 @@ def generate_id(prefix: str, mode=ID_SIMPLE, ns=None): elif mode == ID_PRIVATE: template = TEMPLATE_PRIVATE else: raise ValueError("unsupported mode '%s'", mode) + print('template:', template) + print('prefix :', prefix) + print('mode :', mode) + return _generate_id(template, prefix, ns) @@ -205,6 +209,8 @@ def generate_id(prefix: str, mode=ID_SIMPLE, ns=None): # def _generate_id(template, prefix, ns=None): + print('template:', template) + # FIXME: several of the vars below are constants, and many of them are # rarely used in IDs. They should be created only once per module instance, # and/or only if needed. @@ -272,8 +278,12 @@ def _read_file_counter(name): fname = os.path.join(state_dir, 'ru_%s_%s.cnt' % (user, days)) info['day_counter'] = _read_file_counter(fname) + print('template 2:', template) if '%(item_counter)' in template: + print('template:', template) + print('prefix :', prefix) + # clean up "prefix" to use in file name # FIXME: extend same procedure for other cases (with regex?) if '%(item_counter)' in prefix: @@ -284,8 +294,11 @@ def _read_file_counter(name): break prefix = '.'.join(prefix_parts) + print('prefix :', prefix) + fname = os.path.join(state_dir, 'ru_%s_%s.cnt' % (user, prefix)) info['item_counter'] = _read_file_counter(fname) + print('item_counter:', info['item_counter']) if '%(counter)' in template: info['counter'] = _id_registry.get_counter(prefix.replace('%', '')) From 048e3568fd1878a3e1c3faabfe819b58bb3c3874 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 19 Mar 2025 06:47:01 -0400 Subject: [PATCH 04/44] snap --- src/radical/utils/flux.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 4eafc496f..a962644ff 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -517,8 +517,8 @@ def app_cb(flux_id, exe_fut, event): def id_cb(fut): flux_id = fut.jobid() idx = fut.ru_idx - for ev in ['alloc', 'start', 'finish', 'release', 'exception']: - # 'submit', 'free', 'clean', + for ev in ['submit', 'free', 'clean', + 'alloc', 'start', 'finish', 'release', 'exception']: tmp_cb = partial(app_cb, flux_id) fut.add_event_callback(ev, tmp_cb) futures.append([flux_id, idx, fut]) @@ -532,8 +532,8 @@ def id_cb(fut): fut.add_jobid_callback(id_cb) # wait until we saw all jobid callbacks (assume 10 tasks/sec) - timeout = len(specs) / 10 - timeout = max(10, timeout) + timeout = len(specs) + # timeout = max(10, timeout) start = time.time() self._log.debug('wait %.2fsec for %d flux IDs', timeout, len(specs)) while len(futures) < len(specs): From 08256815a8da0f7d7b4a19abe4e0598328b5d10d Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 19 Mar 2025 11:49:36 +0100 Subject: [PATCH 05/44] snap --- src/radical/utils/flux.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index ef8447b30..98844db88 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -533,8 +533,8 @@ def id_cb(fut): fut.add_jobid_callback(id_cb) # wait until we saw all jobid callbacks (assume 10 tasks/sec) - timeout = len(specs) - # timeout = max(10, timeout) + timeout = len(specs) + timeout = max(100, timeout) start = time.time() self._log.debug('wait %.2fsec for %d flux IDs', timeout, len(specs)) while len(futures) < len(specs): From 3da53e714cbf7088c7916ff0cdfa0f5748820a28 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 19 Mar 2025 06:54:14 -0400 Subject: [PATCH 06/44] snap --- src/radical/utils/ids.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/src/radical/utils/ids.py b/src/radical/utils/ids.py index 15f730563..ecec34c76 100644 --- a/src/radical/utils/ids.py +++ b/src/radical/utils/ids.py @@ -198,10 +198,6 @@ def generate_id(prefix: str, mode=ID_SIMPLE, ns=None): elif mode == ID_PRIVATE: template = TEMPLATE_PRIVATE else: raise ValueError("unsupported mode '%s'", mode) - print('template:', template) - print('prefix :', prefix) - print('mode :', mode) - return _generate_id(template, prefix, ns) @@ -209,8 +205,6 @@ def generate_id(prefix: str, mode=ID_SIMPLE, ns=None): # def _generate_id(template, prefix, ns=None): - print('template:', template) - # FIXME: several of the vars below are constants, and many of them are # rarely used in IDs. They should be created only once per module instance, # and/or only if needed. @@ -278,12 +272,8 @@ def _read_file_counter(name): fname = os.path.join(state_dir, 'ru_%s_%s.cnt' % (user, days)) info['day_counter'] = _read_file_counter(fname) - print('template 2:', template) if '%(item_counter)' in template: - print('template:', template) - print('prefix :', prefix) - # clean up "prefix" to use in file name # FIXME: extend same procedure for other cases (with regex?) if '%(item_counter)' in prefix: @@ -294,11 +284,8 @@ def _read_file_counter(name): break prefix = '.'.join(prefix_parts) - print('prefix :', prefix) - fname = os.path.join(state_dir, 'ru_%s_%s.cnt' % (user, prefix)) info['item_counter'] = _read_file_counter(fname) - print('item_counter:', info['item_counter']) if '%(counter)' in template: info['counter'] = _id_registry.get_counter(prefix.replace('%', '')) From cc127cbe61625e3221c18bb54fd74ed2c0be3b8d Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Thu, 20 Mar 2025 14:17:05 +0100 Subject: [PATCH 07/44] more details, faster typed-dict --- src/radical/utils/profile.py | 20 ++++++++++++++------ src/radical/utils/typeddict.py | 20 +++++++++++++------- 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/src/radical/utils/profile.py b/src/radical/utils/profile.py index b51df8fe3..2f83874a5 100644 --- a/src/radical/utils/profile.py +++ b/src/radical/utils/profile.py @@ -13,6 +13,7 @@ from .config import DefaultConfig from .atfork import atfork from .shell import sh_callout +from .which import which from .modules import import_module @@ -847,10 +848,11 @@ def event_to_label(event): # class Yappi(object): - def __init__(self, name, method='wall'): + def __init__(self, name, method='wall', verbose=False): self._yappi = import_module('yappi') self._yappi.set_clock_type(method) self._name = name + self._verb = verbose def __enter__(self): if self._yappi: @@ -860,17 +862,23 @@ def __exit__(self, etype, value, traceback): if self._yappi: - # self._yappi.get_thread_stats().print_all() + if self._verb: + self._yappi.get_func_stats().print_all() + self._yappi.get_thread_stats().print_all() fstats = self._yappi.get_func_stats() pstats = self._yappi.convert2pstats(fstats) pstats.dump_stats('%s.pstats' % self._name) - cmd = 'gprof2dot -e 0.00 -n 0.12 --skew=0.3' - cmd += ' --node-label="total-time" -f pstats %s.pstats' % self._name - cmd += ' | dot -Tpng -o %s.png' % self._name + if which('gprof2dot'): - sh_callout(cmd, shell=True) + cmd = 'gprof2dot -e 0.00 -n 0.12 --skew=0.3' + cmd += ' --node-label="total-time" -f pstats %s.pstats' % self._name + cmd += ' | dot -Tpng -o %s.png' % self._name + + out, err, ret = sh_callout(cmd, shell=True) + if ret: + print('gprof2dot failed: %s' % err) # ------------------------------------------------------------------------------ diff --git a/src/radical/utils/typeddict.py b/src/radical/utils/typeddict.py index a160f5829..252686956 100644 --- a/src/radical/utils/typeddict.py +++ b/src/radical/utils/typeddict.py @@ -158,11 +158,13 @@ def __init__(self, from_dict=None, **kwargs): if self._deep: - self.update(copy.deepcopy(self._defaults)) + self.__dict__['_data'] = copy.deepcopy(self._defaults) else: - self.update(self._defaults) + self.__dict__['_data'] = dict() + self.__dict__['_data'].update(self._defaults) - self.update(from_dict) + if from_dict: + self.update(from_dict) if kwargs: self.update(kwargs) @@ -228,7 +230,10 @@ def __getitem__(self, k): return self._data[k] def __setitem__(self, k, v): - self._data[k] = self._verify_setter(k, v) + if self._check : + self._data[k] = self._verify_setter(k, v) + else: + self._data[k] = v def __delitem__(self, k): del self._data[k] @@ -295,8 +300,6 @@ def popitem(self): def __getattr__(self, k): if k == '_data': - if '_data' not in self.__dict__: - self.__dict__['_data'] = dict() return self.__dict__['_data'] if k.startswith('__'): @@ -314,7 +317,10 @@ def __setattr__(self, k, v): if k.startswith('__'): return object.__setattr__(self, k, v) - self._data[k] = self._verify_setter(k, v) + if self._check : + self._data[k] = self._verify_setter(k, v) + else: + self._data[k] = v def __delattr__(self, k): From e78cdb24b6b45f6c340a8694c42e9c69053974e4 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Thu, 20 Mar 2025 17:15:25 +0100 Subject: [PATCH 08/44] capture flux stderr --- src/radical/utils/flux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 98844db88..c1b0515a0 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -151,7 +151,7 @@ def _locked_start_service(self, self._log.debug('flux command: %s', ' '.join(cmd)) flux_proc = sp.Popen(cmd, encoding="utf-8", - stdin=sp.DEVNULL, stdout=sp.PIPE, stderr=sp.PIPE) + stdin=sp.DEVNULL, stdout=sp.PIPE, stderr=sp.STDOUT) flux_env = dict() while flux_proc.poll() is None: From 1eb90c246c34108d73ab8dbcc26b8a91988a4812 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Fri, 21 Mar 2025 12:26:38 +0100 Subject: [PATCH 09/44] snap --- src/radical/utils/flux.py | 39 +++++++++++++++++++--------------- src/radical/utils/typeddict.py | 16 +++++++------- 2 files changed, 30 insertions(+), 25 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index c1b0515a0..827f7facb 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -92,14 +92,14 @@ def _watch(self) -> None: # we want to call `flux ping` via `sh_callout`. We should # instead use the Flux Python API to run the pings and pass the # URI explicitly. - self._log.info('starting flux watcher') + self._log.info('%s: starting flux watcher', self._uid) if self._env: for k,v in self._env.items(): os.environ[k] = v out, err, ret = sh_callout('flux resource list') - self._log.info('flux resources [ %d %s]:\n%s', ret, err, out) + self._log.info('%s: flux resources [ %d %s]:\n%s', self._uid, ret, err, out) while not self._term.is_set(): @@ -108,12 +108,12 @@ def _watch(self) -> None: out, err, ret = sh_callout('flux ping -c 1 kvs') self._log.debug('flux ping:%s', out) if ret: - self._log.error('flux watcher err: %s', err) + self._log.error('%s: flux watcher err: %s', self._uid, err) break # we only get here when the ping failed - set the event self._term.set() - self._log.warn('flux stopped') + self._log.warn('%s: flux stopped', self._uid) # -------------------------------------------------------------------------- @@ -148,7 +148,7 @@ def _locked_start_service(self, cmd += ['flux', 'start', 'bash', '-c', 'echo "HOST:$(hostname) URI:$FLUX_URI" && sleep inf'] - self._log.debug('flux command: %s', ' '.join(cmd)) + self._log.debug('%s: flux command: %s', self._uid, ' '.join(cmd)) flux_proc = sp.Popen(cmd, encoding="utf-8", stdin=sp.DEVNULL, stdout=sp.PIPE, stderr=sp.STDOUT) @@ -160,13 +160,14 @@ def _locked_start_service(self, line = flux_proc.stdout.readline() except Exception as e: - self._log.exception('flux service failed to start') + self._log.exception('%s: flux service failed to start', + self._uid) raise RuntimeError('could not start flux') from e if not line: continue - self._log.debug('flux output: %s', line) + self._log.debug('%s: flux output: %s', self._uid, line) if line.startswith('HOST:'): @@ -209,7 +210,8 @@ def _locked_start_service(self, self._watcher.daemon = True self._watcher.start() - self._log.info("flux startup successful: [%s]", flux_env['FLUX_URI']) + self._log.info("%s: flux startup successful: [%s]", + self._uid, flux_env['FLUX_URI']) return self._uri @@ -511,8 +513,8 @@ def submit_jobs(self, def app_cb(flux_id, exe_fut, event): try : cb(flux_id, event) - except: self._log.exception('app cb failed for %s [%s]', - flux_id, event) + except: self._log.exception('%s: app cb failed for %s [%s]', + self._uid, flux_id, event) futures = list() def id_cb(fut): @@ -529,24 +531,27 @@ def id_cb(fut): jobspec = json.dumps(spec) fut = self._exe.submit(jobspec) fut.ru_idx = idx - self._log.debug('submitted : %s', idx) + self._log.debug('%s: submitted : %s', self._uid, idx) fut.add_jobid_callback(id_cb) # wait until we saw all jobid callbacks (assume 10 tasks/sec) timeout = len(specs) timeout = max(100, timeout) start = time.time() - self._log.debug('wait %.2fsec for %d flux IDs', timeout, len(specs)) + self._log.debug('%s: wait %.2fsec for %d flux IDs', + self._uid, timeout, len(specs)) while len(futures) < len(specs): time.sleep(0.1) - self._log.debug('wait %s / %s', len(futures), len(specs)) + self._log.debug('%s: wait %s / %s', self._uid, + len(futures), len(specs)) if time.time() - start > timeout: - raise RuntimeError('timeout on job submission') + raise RuntimeError('%s: timeout on submission', self._uid) self._log.info('got %d flux IDs', len(futures)) # get flux_ids sorted by submission order (idx) flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] + self._log.debug('%s: submitted: %s', self._uid, ids) return flux_ids @@ -567,14 +572,14 @@ def attach_jobs(self, for flux_id in ids: fut = self._exe.attach(flux_id) - self._log.debug('attach %s : %s', flux_id, fut) + self._log.debug('%s: attach %s : %s', self._uid, flux_id, fut) if cb: def app_cb(fut, event): try: cb(flux_id, event) except: - self._log.exception('app cb failed') + self._log.exception('%s: app cb failed', self._uid) for ev in [ 'submit', @@ -599,7 +604,7 @@ def cancel_jobs(self, flux_ids: List[int]) -> None: for flux_id in flux_ids: fut = self._exe.attach(flux_id) - self._log.debug('cancel %s : %s', flux_id, fut) + self._log.debug('%s: cancel %s : %s', self._uid, flux_id, fut) fut.cancel() diff --git a/src/radical/utils/typeddict.py b/src/radical/utils/typeddict.py index 252686956..bb32539df 100644 --- a/src/radical/utils/typeddict.py +++ b/src/radical/utils/typeddict.py @@ -200,14 +200,14 @@ def update(self, other): for k, v in other.items(): if isinstance(v, dict): - t = self._schema.get(k) or \ - (type(self) if self._self_default else TypedDict) - if isinstance(t, type) and issubclass(t, TypedDict): - # cast to expected TypedDict type - if not self.get(k): - self[k] = t() - self[k].update(v) - continue + t = type(self._schema.get(k) or self) + # (type(self) if self._self_default else TypedDict) + # if isinstance(t, type) and issubclass(t, TypedDict): + # cast to expected TypedDict type + if not self.get(k): + self[k] = t() + self[k].update(v) + continue self[k] = v From 130dfc50393fbcf47af5c212245782190cc57364 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Fri, 21 Mar 2025 13:20:42 +0100 Subject: [PATCH 10/44] snap --- src/radical/utils/typeddict.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/radical/utils/typeddict.py b/src/radical/utils/typeddict.py index bb32539df..252686956 100644 --- a/src/radical/utils/typeddict.py +++ b/src/radical/utils/typeddict.py @@ -200,14 +200,14 @@ def update(self, other): for k, v in other.items(): if isinstance(v, dict): - t = type(self._schema.get(k) or self) - # (type(self) if self._self_default else TypedDict) - # if isinstance(t, type) and issubclass(t, TypedDict): - # cast to expected TypedDict type - if not self.get(k): - self[k] = t() - self[k].update(v) - continue + t = self._schema.get(k) or \ + (type(self) if self._self_default else TypedDict) + if isinstance(t, type) and issubclass(t, TypedDict): + # cast to expected TypedDict type + if not self.get(k): + self[k] = t() + self[k].update(v) + continue self[k] = v From f99a1899c0e735d6354129d2b0842528f1d37948 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 24 Mar 2025 15:50:06 +0100 Subject: [PATCH 11/44] snap --- src/radical/utils/flux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 827f7facb..93f1f870a 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -551,7 +551,7 @@ def id_cb(fut): # get flux_ids sorted by submission order (idx) flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] - self._log.debug('%s: submitted: %s', self._uid, ids) + self._log.debug('%s: submitted: %s', self._uid, flux_ids) return flux_ids From b0b45230ba3216d441c92728d0156d81c5069524 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Thu, 27 Mar 2025 10:56:53 +0100 Subject: [PATCH 12/44] snap --- requirements.txt | 1 + src/radical/utils/flux.py | 754 +++++++++------------------ src/radical/utils/ids.py | 32 +- tests/integration_tests/test_flux.py | 14 +- 4 files changed, 262 insertions(+), 539 deletions(-) diff --git a/requirements.txt b/requirements.txt index b55e00ed8..4339c51f4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,4 +5,5 @@ ntplib pyzmq regex setproctitle +rc.process diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 93f1f870a..837ef02bd 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -1,654 +1,372 @@ -# pylint: disable=cell-var-from-loop - -import os -import sys import time -import json import shlex -from typing import Optional, List, Dict, Any, Callable -from functools import partial +import threading as mt + +from rc.process import Process +from functools import partial +from collections import defaultdict +from typing import List, Dict, Any + +from .misc import as_list +from .which import which +from .ids import generate_id +from .logger import Logger +from .modules import import_module -import threading as mt -import subprocess as sp +try: + _flux = import_module('flux') + _flux_job = import_module('flux.job') + _flux_exc = None -from .url import Url -from .ids import generate_id, ID_CUSTOM -from .shell import sh_callout -from .logger import Logger -from .profile import Profiler -from .modules import import_module +except Exception as e: + _flux = None + _flux_job = None + _flux_exc = e -# -------------------------------------------------------------------------- +# ------------------------------------------------------------------------------ # class _FluxService(object): - ''' - Helper class to handle a private Flux instance, including configuration, - start, monitoring and termination. - ''' # -------------------------------------------------------------------------- # - def __init__(self, uid : str, - log : Logger, - prof : Profiler) -> None: + def __init__(self, uid : str, + log : Logger) -> None: - self._uid = uid - self._log = log - self._prof = prof + self._uid = uid + self._log = log - self._lock = mt.RLock() - self._term = mt.Event() + self._fexe = which('flux') + self._uri = None + self._tout = 60 - self._uri = None - self._env = None - self._proc = None - self._watcher = None + if not _flux: + raise RuntimeError('flux module not found') from self._exception - try: - cmd = 'flux python -c "import flux; print(flux.__file__)"' - out, err, ret = sh_callout(cmd) + if not _flux_job: + raise RuntimeError('flux.job module not found') from self._exception - if ret: - raise RuntimeError('flux not found: %s' % err) + if not self._fexe: + raise RuntimeError('flux executable not found') - flux_path = os.path.dirname(out.strip()) - mod_path = os.path.dirname(flux_path) - sys.path.append(mod_path) - - self._flux = import_module('flux') - self._flux_job = import_module('flux.job') - - except Exception: - self._log.exception('flux import failed') - raise + self._start() # -------------------------------------------------------------------------- # @property - def uid(self): - return self._uid - - - @property - def uri(self): + def uri(self) -> str: return self._uri - @property - def env(self): - return self._env - + def timeout(self) -> int: + return self._tout - # -------------------------------------------------------------------------- - # - def _watch(self) -> None: - - # FIXME: this thread will change `os.environ` for this *process* because - # we want to call `flux ping` via `sh_callout`. We should - # instead use the Flux Python API to run the pings and pass the - # URI explicitly. - self._log.info('%s: starting flux watcher', self._uid) - - if self._env: - for k,v in self._env.items(): - os.environ[k] = v - - out, err, ret = sh_callout('flux resource list') - self._log.info('%s: flux resources [ %d %s]:\n%s', self._uid, ret, err, out) - - while not self._term.is_set(): - - time.sleep(1) - - out, err, ret = sh_callout('flux ping -c 1 kvs') - self._log.debug('flux ping:%s', out) - if ret: - self._log.error('%s: flux watcher err: %s', self._uid, err) - break - - # we only get here when the ping failed - set the event - self._term.set() - self._log.warn('%s: flux stopped', self._uid) + @timeout.setter + def timeout(self, tout) -> None: + self._tout = tout # -------------------------------------------------------------------------- # - def start_service(self, - launcher: Optional[str] = None, - env : Optional[Dict[str,str]] = None - ) -> Optional[str]: - - with self._lock: - - if self._proc is not None: - raise RuntimeError('already started Flux: %s' % self._uri) - - self._term.clear() - - return self._locked_start_service(launcher, env) + def _proc_line_cb(self, prefix: str, + proc : Process, + lines : List[str] + ) -> None: + for line in lines: + if line.startswith('FLUX_URI:'): + self._uri = line.strip().split(':', 1)[1] # -------------------------------------------------------------------------- # - def _locked_start_service(self, - launcher: Optional[str] = None, - env : Optional[Dict[str,str]] = None - ) -> Optional[str]: - - cmd = list() - - if launcher: - cmd += shlex.split(launcher) - - cmd += ['flux', 'start', 'bash', '-c', - 'echo "HOST:$(hostname) URI:$FLUX_URI" && sleep inf'] - - self._log.debug('%s: flux command: %s', self._uid, ' '.join(cmd)) - - flux_proc = sp.Popen(cmd, encoding="utf-8", - stdin=sp.DEVNULL, stdout=sp.PIPE, stderr=sp.STDOUT) - - flux_env = dict() - while flux_proc.poll() is None: - - try: - line = flux_proc.stdout.readline() - - except Exception as e: - self._log.exception('%s: flux service failed to start', - self._uid) - raise RuntimeError('could not start flux') from e - - if not line: - continue - - self._log.debug('%s: flux output: %s', self._uid, line) - - if line.startswith('HOST:'): - - flux_host, flux_uri = line.split(' ', 1) - - flux_host = flux_host.split(':', 1)[1].strip() - flux_uri = flux_uri.split(':', 1)[1].strip() - - flux_env['FLUX_HOST'] = flux_host - flux_env['FLUX_URI'] = flux_uri - break - - if flux_proc.poll() is not None: - raise RuntimeError('could not execute `flux start`') - - # fr = self._flux.uri.uri.FluxURIResolver() - # ret = fr.resolve('pid:%d' % flux_proc.pid) - # flux_env = {'FLUX_URI': ret} - - assert 'FLUX_URI' in flux_env, 'no FLUX_URI in env' - - # make sure that the flux url can be reached from other hosts - # FIXME: this also routes local access via ssh which may slow comm - flux_url = Url(flux_env['FLUX_URI']) - flux_url.host = flux_env['FLUX_HOST'] - flux_url.schema = 'ssh' - flux_uri = str(flux_url) - flux_env['FLUX_URI'] = flux_uri - - self._uri = flux_uri - self._env = flux_env - self._proc = flux_proc - - self._log.debug('flux uri: %s', flux_uri) - - self._prof.prof('flux_started', msg=self._uid) - - # start watcher thread to monitor the instance - self._watcher = mt.Thread(target=self._watch) - self._watcher.daemon = True - self._watcher.start() - - self._log.info("%s: flux startup successful: [%s]", - self._uid, flux_env['FLUX_URI']) - - return self._uri - - - # -------------------------------------------------------------------------- - # - def check_service(self) -> Optional[str]: - - with self._lock: - - if not self._proc: - raise RuntimeError('flux service was not yet started') - - if self._term.is_set(): - raise RuntimeError('flux service was terminated') - - return self._uri + def _proc_state_cb(self, proc: Process, state: str) -> None: + self._log.info('flux instance state update: %s', state) # -------------------------------------------------------------------------- # - def close_service(self) -> None: - - with self._lock: + def _start(self) -> None: - self.check_service() + fcmd = 'echo FLUX_URI:\\$FLUX_URI && sleep inf' + cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) - if not self._proc: - raise RuntimeError('cannot kill flux from this process') + self._log.info('%s: start flux instance: %s', self._uid, cmd) - if self._watcher: - self._watcher.join() + p = Process(cmd) + p.register_cb(p.CB_OUT_LINE, partial(self._proc_line_cb, 'out')) + p.register_cb(p.CB_ERR_LINE, partial(self._proc_line_cb, 'err')) + p.register_cb(p.CB_STATE, self._proc_state_cb) + p.polldelay = 0.1 + p.start() - # terminate the service process - # FIXME: send termination signal to flux for cleanup - self._proc.kill() + start = time.time() + while time.time() - start < self._tout: time.sleep(0.1) - self._proc.terminate() - self._proc.wait() + if self._uri is not None: + break + + if self.uri is None: + self._log.error('%s: flux instance did not start', self._uid) + raise RuntimeError('%s: flux instance did not start', self._uid) - self._uri = None - self._env = None + self._log.info('%s: found flux uri: %s', self._uid, self.uri) # ------------------------------------------------------------------------------ # class FluxHelper(object): - ''' - Helper CLASS to programnatically handle flux instances and to obtain state - update events for flux jobs known in that instance. - ''' - # -------------------------------------------------------------------------- # - def __init__(self, uid:str = None) -> None: - ''' - The Flux Helper c'tor takes no arguments and will initially not be - connected to a Flux instance. After construction, the application can - call either one of the following methods: - - FluxHelper.connect_flux(uri=None) - FluxHelper.start_flux() - - The first will attempt to connect to the Flux instance referenced by - that URI - a `ValueError` exception will be raised if that instance - cannot be reached. If no URI is provided, the environment variable - `FLUX_URI` will be used. - - The second method will instantiate a new flux instance in the current - process environment. - - In both cases, the properties - - FluxHelper.uri - FluxHelper.env + def __init__(self, uri : str = None, + log : Logger = None) -> None: - will provide information about the connected Flux instance. The `uri` - is provided as a string, the `env` as a dictionary of environment - settings (including `FLUX_URI` again). + self._uri = uri + self._log = log or Logger('radical.utils.flux') + self._uid = generate_id('ru.flux') + self._handle = None + self._journal = None + self._service = None + self._jthread = None + self._started = False - The method + self._tasks = dict() # task ID -> task + self._task_ids = dict() # flux ID -> task ID + self._flux_ids = dict() # task ID -> flux ID - FluxHelper.reset() + self._elock = mt.Lock() # lock event dict + self._events = defaultdict(list) # flux ID -> event list + self._cbacks = list() # list of callbacks - will disconnect from the Flux instance, and in the case where - `start_flux` created a private instance, that instance will be killed. - The `uri` and `env` properties will be reset to `None`. + if not _flux: + raise RuntimeError('flux module not found') from _flux_exc + if not _flux_job: + raise RuntimeError('flux.job module not found') from _flux_exc - While connected to a Flux instance, the following methods can be used to - interact with the instance: - FluxHelper.get_executor() - return a flux.job.Executor instance - FluxHelper.get_handle() - return a flux.job.Flux instance - - All provided executors and handles will be invalidated upon `reset()`. - ''' - - self._service : Optional[_FluxService] = None - - if uid: self._uid = uid - else : self._uid = generate_id('flux.%(item_counter)04d', ID_CUSTOM) - - self._log = Logger ('radical.utils', ns='radical.utils', level='DEBUG') - self._prof = Profiler('radical.utils', ns='radical.utils') - - self._lock = mt.RLock() - - self._uri = None - self._env = None - self._exe = None - self._handle = None - self._handles = list() # TODO - self._executors = list() # TODO + # -------------------------------------------------------------------------- + # + def start(self) -> None: - try: - cmd = 'flux python -c "import flux; print(flux.__file__)"' - out, err, ret = sh_callout(cmd) + if self._started: + return - if ret: - raise RuntimeError('flux not found: %s' % err) + if not self._uri: + self._flux_service = _FluxService(uid=self._uid, log=self._log) + self._uri = self._flux_service.uri - flux_path = os.path.dirname(out.strip()) - mod_path = os.path.dirname(flux_path) - sys.path.append(mod_path) + self._handle = _flux.Flux(self._uri) - self._flux = import_module('flux') - self._flux_job = import_module('flux.job') + self._jthread = mt.Thread(target=self._watcher) + self._jthread.daemon = True + self._jthread.start() - except Exception: - self._log.exception('flux import failed') - raise + self._started = True # -------------------------------------------------------------------------- # - def __del__(self): - - # FIXME: are handles / executors correctly garbage collected? - self.reset() + @property + def uid(self) -> str: + return self._uid # -------------------------------------------------------------------------- # - def reset(self): - ''' - Close the connection to the FLux instance (if it exists), and terminate - the Flux service if it was started by this instance. All handles and - executors created for this service will be invalidated. - ''' - - with self._lock: - - for idx in range(len(self._handles)): - del self._handles[idx] + def _watcher(self): - for idx in range(len(self._executors)): - del self._executors[idx] + # NOTE: *never* used self._handle in this thread, as it is not thread + # safe. Instead, use the private handle created here + handle = _flux.Flux(self._uri) - self._exe = None - self._handle = None + # start watching the event journal + self._journal = _flux_job.JournalConsumer(handle) + self._journal.start() - if self._uri: - try: - self._service.close_service() - except: - pass - self._uri = None - self._env = None + while True: + try: + event = self._journal.poll(timeout=1.0) + self._handle_events(event.jobid, event) - # -------------------------------------------------------------------------- - # - @property - def uid(self): - ''' - unique ID for this FluxHelper instance - ''' - - with self._lock: - return self._uid + except TimeoutError: + pass # -------------------------------------------------------------------------- # - @property - def uri(self): - ''' - uri for the connected Flux instance. Returns `None` if no instance is - connected. - ''' + def register_cb(self, cb: callable) -> None: - with self._lock: - return self._uri + with self._elock: + self._cbacks.append(cb) # -------------------------------------------------------------------------- # - @property - def env(self): - ''' - environment dict for the connected Flux instance. Returns `None` if no - instance is connected. - ''' + def unregister_cb(self, cb: callable) -> None: - with self._lock: - return self._env + with self._elock: + self._cbacks.remove(cb) # -------------------------------------------------------------------------- # - def start_flux(self, launcher: Optional[str] = None) -> None: - ''' - Start a private Flux instance - - FIXME: forward env - ''' + def _handle_events(self, flux_id: 'flux.job.JobID', + event : 'flux.job.journal.JournalEvent' = None + ) -> None: - with self._lock: + with self._elock: - if self._uri: - raise RuntimeError('service already connected: %s' % self._uri) + # if triggered by submit, check if we have anything to do + if not event: + if flux_id not in self._events: + return - self._service = _FluxService(self._uid, self._log, self._prof) - self._service.start_service(launcher=launcher) + # check if we can handle the event - otherwise store it + if not self._cbacks: + self._events[flux_id].append(event) + return - self._uri = self._service.check_service() - self._env = self._service.env + # check if we already know the task - otherwise store the event + if flux_id not in self._task_ids: + self._events[flux_id].append(event) + return - # with ru_open(self._uid + '.dump', 'a') as fout: - # fout.write('start flux pid %d: %s\n' % (os.getpid(), self._uri)) - # for l in get_stacktrace()[:-1]: - # fout.write(l) + # task is known, process stored events + for ev in self._events[flux_id]: + for cb in self._cbacks: + cb(self._task_ids[flux_id], ev) - self._setup() + # process the current event + if event: + for cb in self._cbacks: + cb(self._task_ids[flux_id], event) # -------------------------------------------------------------------------- # - def connect_flux(self, uri : Optional[str] = None) -> None: - ''' - Connect to an existing Flux instance - ''' + def spec_from_command(self, cmd: str) -> 'flux.job.JobspecV1': - with self._lock: - - # with ru_open(self._uid + '.dump', 'a') as fout: - # fout.write('connect flux %d: %s\n' % (os.getpid(), uri)) - # for l in get_stacktrace(): - # fout.write(l + '\n') - - if self._uri: - raise RuntimeError('service already connected: %s' % self._uri) - - if not uri: - uri = os.environ.get('FLUX_URI') - - if not uri: - raise RuntimeError('no Flux instance found via FLUX_URI') - - self._uri = uri - self._env = {'FLUX_URI': uri} - - # FIXME: run a ping test to ensure the service is up - - self._setup() - - - # ---------------------------------------------------------------------- - # - def _setup(self): - ''' - Once a service is connected, create a handle and executor - ''' - - with self._lock: - - assert self._uri, 'not initialized' - - # create a executor and handle for job management - self._exe = self.get_executor() - self._handle = self.get_handle() + return _flux_job.JobspecV1.from_command(shlex.split(cmd)) # -------------------------------------------------------------------------- # - def submit_jobs(self, - specs: List[Dict[str, Any]], - cb : Optional[Callable[[str, Any], None]] = None - ) -> Any: - - with self._lock: - - if not self._uri: - raise RuntimeError('FluxHelper is not connected') - - assert self._exe - - def app_cb(flux_id, exe_fut, event): - try : cb(flux_id, event) - except: self._log.exception('%s: app cb failed for %s [%s]', - self._uid, flux_id, event) - - futures = list() - def id_cb(fut): - flux_id = fut.jobid() - idx = fut.ru_idx - for ev in ['submit', 'free', 'clean', - 'alloc', 'start', 'finish', 'release', 'exception']: - tmp_cb = partial(app_cb, flux_id) - fut.add_event_callback(ev, tmp_cb) - futures.append([flux_id, idx, fut]) - self._log.debug('got flux id: %s: %s', idx, flux_id) - - for idx, spec in enumerate(specs): - jobspec = json.dumps(spec) - fut = self._exe.submit(jobspec) - fut.ru_idx = idx - self._log.debug('%s: submitted : %s', self._uid, idx) - fut.add_jobid_callback(id_cb) - - # wait until we saw all jobid callbacks (assume 10 tasks/sec) - timeout = len(specs) - timeout = max(100, timeout) - start = time.time() - self._log.debug('%s: wait %.2fsec for %d flux IDs', - self._uid, timeout, len(specs)) - while len(futures) < len(specs): - time.sleep(0.1) - self._log.debug('%s: wait %s / %s', self._uid, - len(futures), len(specs)) - if time.time() - start > timeout: - raise RuntimeError('%s: timeout on submission', self._uid) - self._log.info('got %d flux IDs', len(futures)) - - # get flux_ids sorted by submission order (idx) - flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] - - self._log.debug('%s: submitted: %s', self._uid, flux_ids) - return flux_ids + def spec_from_dict(self, td: dict) -> 'flux.job.JobspecV1': + + version = 1 + tasks = [{'command': [td['executable']] + td.get('arguments', []), + 'slot' : 'task', + 'count' : {'per_slot': 1}}] + + system = {'duration': td.get('duration', 0.0)} + + if 'environment' in td: system['environment'] = td['environment'] + if 'sandbox' in td: system['cwd'] = td['sandbox'] + if 'shell' in td: system['shell'] = td['shell'] + if 'stdin' in td: system['stdin'] = td['stdin'] + if 'stdout' in td: system['stdout'] = td['stdout'] + if 'stderr' in td: system['stderr'] = td['stderr'] + if 'uid' in td: system['job'] = {'name': td['uid']} + + attributes = {'system' : system} + resources = [{'count': td.get('ranks', 1), + 'type' : 'slot', + 'label': 'task', + 'with' : [{ + 'count': int(td.get('cores_per_rank', 1)), + 'type' : 'core'}]}] + # 'count': int(td.get('gpus_per_rank', 0)) or None, + # 'type' : 'gpu' + + if 'gpus_per_rank' in td: + resources[0]['with'].append({ + # flux likes integer GPU counts + 'count': math.ceil(td['gpus_per_rank']), + 'type' : 'gpu'}) + + # import json + # return json.dumps({ + # 'version' : version, + # 'resources' : resources, + # 'attributes': attributes, + # 'tasks' : tasks}) + + spec = _flux_job.JobspecV1(resources=resources, + attributes=attributes, + tasks=tasks, + version=version) + + return spec # -------------------------------------------------------------------------- # - def attach_jobs(self, - ids: List[int], - cb : Optional[Callable[[int, Any], None]] = None - ) -> Any: - - with self._lock: + def submit(self, descriptions: List[Dict[str, Any]]) -> List[str]: - if not self._uri: - raise RuntimeError('FluxHelper is not connected') + jobs = list() - assert self._exe, 'no executor' + def _submit_cb(tid: str, f: _flux.future.Future) -> None: + # care must be taken on the order of some of the operations: the + # `journal_cb` will check if a flux_id is known, and if so will + # assume that the task id is known also and callbacks can be issued. + # So *first* register the task ID, *then* the flux ID, to avoid the + # need for additional locking. - for flux_id in ids: + flux_id = _flux_job.submit_get_id(f) + self._flux_ids[tid] = flux_id + self._task_ids[flux_id] = tid + jobs.append(flux_id) - fut = self._exe.attach(flux_id) - self._log.debug('%s: attach %s : %s', self._uid, flux_id, fut) + # if we already got events we'll invoke the callbacks now + self._handle_events(flux_id) - if cb: - def app_cb(fut, event): - try: - cb(flux_id, event) - except: - self._log.exception('%s: app cb failed', self._uid) - - for ev in [ - 'submit', - 'alloc', - 'start', - 'finish', - 'release', - # 'free', - # 'clean', - 'exception', - ]: - fut.add_event_callback(ev, app_cb) + try: + # asynchronously submit jobspec files from a directory + for i, descr in enumerate(descriptions): + cb = partial(_submit_cb, 'task.%04d' % i) + spec = descr + fut = _flux_job.submit_async(self._handle, spec, waitable=True) + fut.then(cb) + # make sure we get jobid events + if self._handle.reactor_run() < 0: + self._log.error("reactor run failed") + self._handle.fatal_error("reactor start failed") - # -------------------------------------------------------------------------- - # - def cancel_jobs(self, flux_ids: List[int]) -> None: + # wait for all jobs to get IDs + # FIXME: Do we need a timeout? Also, the jobid cb can signal on + # completion, so we could wait for that instead of polling. + while len(jobs) < len(descriptions): + time.sleep(0.001) - with self._lock: + return jobs - assert self._exe, 'no executor' - for flux_id in flux_ids: - fut = self._exe.attach(flux_id) - self._log.debug('%s: cancel %s : %s', self._uid, flux_id, fut) - fut.cancel() + except Exception: + self._log.exception("exception") + raise # -------------------------------------------------------------------------- # - def get_handle(self) -> Any: - - with self._lock: - - if not self._uri: - raise RuntimeError('FluxHelper is not connected') - - try: - handle = self._flux.Flux(url=self._uri) - assert handle, 'no handle' + def cancel(self, flux_ids: [str|List[str]]) -> None: - except Exception as e: - raise RuntimeError('failed to connect at %s' % self._uri) from e - - self._handles.append(handle) - - return handle + for flux_id in as_list(flux_ids): + _flux_job.cancel_async(self._handle, flux_id, reason='user cancel') # -------------------------------------------------------------------------- # - def get_executor(self) -> Any: - - with self._lock: - - if not self._uri: - raise RuntimeError('FluxHelper is not connected') - - try: - args = {'url': self._uri} - exe = self._flux_job.executor.FluxExecutor(handle_kwargs=args) - assert exe, 'no executor' - - except Exception as e: - raise RuntimeError('failed to connect at %s' % self._uri) from e - - self._executors.append(exe) + def wait(self, flux_ids: [str|List[str]]) -> None: - return exe + for flux_id in as_list(flux_ids): + _flux_job.wait(self._handle, flux_id) # ------------------------------------------------------------------------------ diff --git a/src/radical/utils/ids.py b/src/radical/utils/ids.py index ecec34c76..c974530c7 100644 --- a/src/radical/utils/ids.py +++ b/src/radical/utils/ids.py @@ -252,20 +252,26 @@ def _generate_id(template, prefix, ns=None): if '%(uuid)' in template: info['uuid'] = uuid.uuid1() # plain uuid def _read_file_counter(name): - fd = os.open(name, os.O_RDWR | os.O_CREAT) + output = 0 try: - fcntl.flock(fd, fcntl.LOCK_EX) - except OSError: - # fcntl.flock might cause OSError: [Errno 524] Unknown error 524 - # (the case for Theta@ALCF) - fcntl.lockf(fd, fcntl.LOCK_EX) - os.lseek(fd, 0, os.SEEK_SET) - data = os.read(fd, 256) - if not data: output = 0 - else : output = int(data) - os.lseek(fd, 0, os.SEEK_SET) - os.write(fd, str.encode("%d\n" % (output + 1))) - os.close(fd) + fd = os.open(name, os.O_RDWR | os.O_CREAT) + try: + fcntl.flock(fd, fcntl.LOCK_EX) + except OSError: + # fcntl.flock might cause OSError: [Errno 524] Unknown error 524 + # (the case for Theta@ALCF) + fcntl.lockf(fd, fcntl.LOCK_EX) + os.lseek(fd, 0, os.SEEK_SET) + data = os.read(fd, 256) + if data: output = int(data) + os.lseek(fd, 0, os.SEEK_SET) + os.write(fd, str.encode("%d\n" % (output + 1))) + os.close(fd) + finally: + try: + os.close(fd) + except: + pass return output if '%(day_counter)' in template: diff --git a/tests/integration_tests/test_flux.py b/tests/integration_tests/test_flux.py index 25513865a..9beeee4e5 100755 --- a/tests/integration_tests/test_flux.py +++ b/tests/integration_tests/test_flux.py @@ -51,13 +51,12 @@ def test_flux_startup(): njobs = 10 events = dict() - def cb1(job_id, state, ts, context): + def cb1(job_id, state): - # print([job_id, state, ts, context]) if job_id not in events: - events[job_id] = [ts, state] + events[job_id] = [state] else: - events[job_id].append([ts, state]) + events[job_id].append(state) fh = ru.FluxHelper() @@ -99,13 +98,12 @@ def test_flux_pickup(): for k,v in outer_fh.env.items(): os.environ[k] = v - def cb1(job_id, state, ts, context): + def cb1(job_id, state): - # print([job_id, state, ts, context]) if job_id not in events: - events[job_id] = [ts, state] + events[job_id] = [state] else: - events[job_id].append([ts, state]) + events[job_id].append(state) fh = ru.FluxHelper() fh.start_flux() From 22285ca0535ff287c15e8fbf2991082de367c7b2 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Fri, 28 Mar 2025 08:35:47 +0100 Subject: [PATCH 13/44] snap --- src/radical/utils/flux.py | 112 ++++++++++++++++++-------------------- 1 file changed, 54 insertions(+), 58 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 837ef02bd..523fc398b 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -11,7 +11,7 @@ from .misc import as_list from .which import which -from .ids import generate_id +from .ids import generate_id, ID_SIMPLE from .logger import Logger from .modules import import_module @@ -81,6 +81,7 @@ def _proc_line_cb(self, prefix: str, if line.startswith('FLUX_URI:'): self._uri = line.strip().split(':', 1)[1] + # -------------------------------------------------------------------------- # def _proc_state_cb(self, proc: Process, state: str) -> None: @@ -134,11 +135,9 @@ def __init__(self, uri : str = None, self._jthread = None self._started = False - self._tasks = dict() # task ID -> task - self._task_ids = dict() # flux ID -> task ID - self._flux_ids = dict() # task ID -> flux ID - - self._elock = mt.Lock() # lock event dict + self._elock = mt.RLock() # lock event dict + self._task_ids = dict() # flux ID -> task ID + self._flux_ids = dict() # task ID -> flux ID self._events = defaultdict(list) # flux ID -> event list self._cbacks = list() # list of callbacks @@ -168,6 +167,20 @@ def start(self) -> None: self._started = True + # -------------------------------------------------------------------------- + # + def stop(self): + + if not self._started: + self._jterm.set() + self._jthread.join() + + # FIXME + self._handle = None + self._flux_service = None + self._uri = None + self._started = False + # -------------------------------------------------------------------------- # @@ -216,36 +229,37 @@ def unregister_cb(self, cb: callable) -> None: # -------------------------------------------------------------------------- # - def _handle_events(self, flux_id: 'flux.job.JobID', - event : 'flux.job.journal.JournalEvent' = None + def _handle_events(self, fid : 'flux.job.JobID', + event: 'flux.job.journal.JournalEvent' = None ) -> None: with self._elock: # if triggered by submit, check if we have anything to do if not event: - if flux_id not in self._events: + if fid not in self._events: return # check if we can handle the event - otherwise store it if not self._cbacks: - self._events[flux_id].append(event) + self._events[fid].append(event) return - # check if we already know the task - otherwise store the event - if flux_id not in self._task_ids: - self._events[flux_id].append(event) + # check if application knows the task - otherwise store the event + if fid not in self._task_ids: + self._events[fid].append(event) return - # task is known, process stored events - for ev in self._events[flux_id]: + # task is known, flush stored events + for ev in self._events[fid]: for cb in self._cbacks: - cb(self._task_ids[flux_id], ev) + cb(self._task_ids[fid], ev) + self._events[fid] = [] # process the current event if event: for cb in self._cbacks: - cb(self._task_ids[flux_id], event) + cb(self._task_ids[fid], event) # -------------------------------------------------------------------------- @@ -301,52 +315,32 @@ def spec_from_dict(self, td: dict) -> 'flux.job.JobspecV1': attributes=attributes, tasks=tasks, version=version) - return spec # -------------------------------------------------------------------------- # - def submit(self, descriptions: List[Dict[str, Any]]) -> List[str]: - - jobs = list() - - def _submit_cb(tid: str, f: _flux.future.Future) -> None: - # care must be taken on the order of some of the operations: the - # `journal_cb` will check if a flux_id is known, and if so will - # assume that the task id is known also and callbacks can be issued. - # So *first* register the task ID, *then* the flux ID, to avoid the - # need for additional locking. - - flux_id = _flux_job.submit_get_id(f) - self._flux_ids[tid] = flux_id - self._task_ids[flux_id] = tid - jobs.append(flux_id) - - # if we already got events we'll invoke the callbacks now - self._handle_events(flux_id) + def submit(self, specs: List[Dict[str, Any]]) -> List[str]: try: - # asynchronously submit jobspec files from a directory - for i, descr in enumerate(descriptions): - cb = partial(_submit_cb, 'task.%04d' % i) - spec = descr + futs = list() + tids = list() + for spec in specs: + tid = spec.attributes['system'].get('job',{}).get('name') \ + or generate_id('ru.flux.task', ID_SIMPLE) fut = _flux_job.submit_async(self._handle, spec, waitable=True) - fut.then(cb) - - # make sure we get jobid events - if self._handle.reactor_run() < 0: - self._log.error("reactor run failed") - self._handle.fatal_error("reactor start failed") + futs.append(fut) + tids.append(tid) - # wait for all jobs to get IDs - # FIXME: Do we need a timeout? Also, the jobid cb can signal on - # completion, so we could wait for that instead of polling. - while len(jobs) < len(descriptions): - time.sleep(0.001) + for fut, tid in zip(futs, tids): + fid = fut.get_id() + self._task_ids[fid] = tid + self._flux_ids[tid] = fid - return jobs + # check if we meanwhile got events to handle + self._handle_events(fid) + return tids except Exception: self._log.exception("exception") @@ -355,18 +349,20 @@ def _submit_cb(tid: str, f: _flux.future.Future) -> None: # -------------------------------------------------------------------------- # - def cancel(self, flux_ids: [str|List[str]]) -> None: + def cancel(self, tids: [str|List[str]]) -> None: - for flux_id in as_list(flux_ids): - _flux_job.cancel_async(self._handle, flux_id, reason='user cancel') + for tid in as_list(tids): + fid = self._flux_ids[tid] + _flux_job.cancel_async(self._handle, fid, reason='user cancel') # -------------------------------------------------------------------------- # - def wait(self, flux_ids: [str|List[str]]) -> None: + def wait(self, tids: [str|List[str]]) -> None: - for flux_id in as_list(flux_ids): - _flux_job.wait(self._handle, flux_id) + for tid in as_list(tids): + fid = self._flux_ids[tid] + _flux_job.wait(self._handle, fid) # ------------------------------------------------------------------------------ From de0a839d03ca526b97d635fdbf586ab4b6a1eaab Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Fri, 28 Mar 2025 10:31:56 +0100 Subject: [PATCH 14/44] snap --- src/radical/utils/flux.py | 161 +++++++++++++++++++++++++++----------- 1 file changed, 117 insertions(+), 44 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 523fc398b..c512f2f27 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -1,6 +1,7 @@ import time import shlex +import queue import threading as mt @@ -130,12 +131,22 @@ def __init__(self, uri : str = None, self._log = log or Logger('radical.utils.flux') self._uid = generate_id('ru.flux') self._handle = None - self._journal = None self._service = None + + # journal watcher self._jthread = None - self._started = False - self._elock = mt.RLock() # lock event dict + # event handle thread + self._ethread = None + self._equeue = queue.Queue() + + # submit thread + self._sthread = None + self._squeue = queue.Queue() + self._sevent = mt.Event() + + self._idlock = mt.Lock() # lock ID dicts + self._elock = mt.Lock() # lock event dict self._task_ids = dict() # flux ID -> task ID self._flux_ids = dict() # task ID -> flux ID self._events = defaultdict(list) # flux ID -> event list @@ -152,7 +163,7 @@ def __init__(self, uri : str = None, # def start(self) -> None: - if self._started: + if self._handle is not None: return if not self._uri: @@ -161,25 +172,31 @@ def start(self) -> None: self._handle = _flux.Flux(self._uri) - self._jthread = mt.Thread(target=self._watcher) + self._jthread = mt.Thread(target=self._jwatcher) self._jthread.daemon = True self._jthread.start() - self._started = True + self._ethread = mt.Thread(target=self._ewatcher) + self._ethread.daemon = True + self._ethread.start() + + self._sthread = mt.Thread(target=self._swatcher) + self._sthread.daemon = True + self._sthread.start() + # -------------------------------------------------------------------------- # def stop(self): - if not self._started: + if self._handle is None: self._jterm.set() self._jthread.join() - # FIXME - self._handle = None + # FIXME: shutdown flux instance self._flux_service = None self._uri = None - self._started = False + self._handle = None # -------------------------------------------------------------------------- @@ -191,26 +208,87 @@ def uid(self) -> str: # -------------------------------------------------------------------------- # - def _watcher(self): + def _jwatcher(self): # NOTE: *never* used self._handle in this thread, as it is not thread # safe. Instead, use the private handle created here - handle = _flux.Flux(self._uri) + fh = _flux.Flux(self._uri) # start watching the event journal - self._journal = _flux_job.JournalConsumer(handle) - self._journal.start() + journal = _flux_job.JournalConsumer(fh) + journal.start() while True: try: - event = self._journal.poll(timeout=1.0) - self._handle_events(event.jobid, event) + event = journal.poll(timeout=1.0) + self._handle_events(fh, event.jobid, event) except TimeoutError: pass + # -------------------------------------------------------------------------- + # + def _swatcher(self): + + # if we get new specs, submit them, return IDs to iqueue, and also + # forward ID to ewatcher + + fh = _flux.Flux(self._uri) + while True: + + try: + specs = self._squeue.get(timeout=1.0) + + except queue.Empty: + continue + + try: + with self._idlock: + + futs = list() + tids = list() + + for spec in specs: + tid = spec.attributes['system']['job']['name'] + fut = _flux_job.submit_async(fh, spec, waitable=True) + + futs.append([fut, tid]) + + for fut, tid in futs: + fid = fut.get_id() + self._task_ids[fid] = tid + self._flux_ids[tid] = fid + + # trigger an event check + self._equeue.put(fid) + + # trigger submit completion + self._sevent.set() + + except Exception: + self._log.exception("exception") + raise + + + # -------------------------------------------------------------------------- + # + def _ewatcher(self): + + # if we get a new job ID, check if we have events for it + + fh = _flux.Flux(self._uri) + while True: + + try: + fid = self._equeue.get(timeout=1.0) + self._handle_events(fh, fid) + + except queue.Empty: + continue + + # -------------------------------------------------------------------------- # def register_cb(self, cb: callable) -> None: @@ -229,7 +307,8 @@ def unregister_cb(self, cb: callable) -> None: # -------------------------------------------------------------------------- # - def _handle_events(self, fid : 'flux.job.JobID', + def _handle_events(self, fh : 'flux.Flux', + fid : 'flux.job.JobID', event: 'flux.job.journal.JournalEvent' = None ) -> None: @@ -322,47 +401,41 @@ def spec_from_dict(self, td: dict) -> 'flux.job.JobspecV1': # def submit(self, specs: List[Dict[str, Any]]) -> List[str]: - try: - futs = list() - tids = list() - for spec in specs: - tid = spec.attributes['system'].get('job',{}).get('name') \ - or generate_id('ru.flux.task', ID_SIMPLE) - fut = _flux_job.submit_async(self._handle, spec, waitable=True) - futs.append(fut) - tids.append(tid) - - for fut, tid in zip(futs, tids): - fid = fut.get_id() - self._task_ids[fid] = tid - self._flux_ids[tid] = fid - - # check if we meanwhile got events to handle - self._handle_events(fid) + # ensure we have a job name which we use as task ID + for spec in specs: + tid = spec.attributes['system'].get('job', {}).get('name') + if not tid: + tid = generate_id(ID_SIMPLE) + if 'job' not in spec.attributes['system']: + spec.attributes['system']['job'] = dict() + spec.attributes['system']['job']['name'] = tid - return tids + self._sevent.clear() + self._squeue.put(specs) + self._sevent.wait() # FIXME: timeout? - except Exception: - self._log.exception("exception") - raise + tids = [spec.attributes['system']['job']['name'] for spec in specs] + return tids # -------------------------------------------------------------------------- # def cancel(self, tids: [str|List[str]]) -> None: - for tid in as_list(tids): - fid = self._flux_ids[tid] - _flux_job.cancel_async(self._handle, fid, reason='user cancel') + with self._idlock: + for tid in as_list(tids): + fid = self._flux_ids[tid] + _flux_job.cancel_async(self._handle, fid, reason='user cancel') # -------------------------------------------------------------------------- # def wait(self, tids: [str|List[str]]) -> None: - for tid in as_list(tids): - fid = self._flux_ids[tid] - _flux_job.wait(self._handle, fid) + with self._idlock: + for tid in as_list(tids): + fid = self._flux_ids[tid] + _flux_job.wait(self._handle, fid) # ------------------------------------------------------------------------------ From 4101b05248006fde90131704ed2bb3117e1149c9 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sat, 29 Mar 2025 13:18:01 +0100 Subject: [PATCH 15/44] snap --- src/radical/utils/__init__.py | 2 +- src/radical/utils/flux.py | 277 ++++++++++++++++++++-------------- 2 files changed, 162 insertions(+), 117 deletions(-) diff --git a/src/radical/utils/__init__.py b/src/radical/utils/__init__.py index 1f8c861fc..8c5f0216b 100644 --- a/src/radical/utils/__init__.py +++ b/src/radical/utils/__init__.py @@ -66,7 +66,7 @@ from .zmq import PubSub, Publisher, Subscriber from .zmq import Server, Client -from .flux import FluxHelper +from .flux import FluxService, FluxHelper from .logger import DEBUG, INFO, WARNING, WARN, ERROR, CRITICAL, OFF from .logger import Logger diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index c512f2f27..3a737819c 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -1,4 +1,5 @@ +import math import time import shlex import queue @@ -28,21 +29,76 @@ _flux_exc = e +# -------------------------------------------------------------------------- +# +def spec_from_command(cmd: str) -> 'flux.job.JobspecV1': + + spec = _flux_job.JobspecV1.from_command(shlex.split(cmd)) + spec.attributes['user']['uid'] = generate_id(ID_SIMPLE) + + return spec + + +# -------------------------------------------------------------------------- +# +def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': + + version = 1 + user = {'uid' : td.get('uid', generate_id(ID_SIMPLE))} + system = {'duration': td.get('duration', 0.0)} + tasks = [{'command': [td['executable']] + td.get('arguments', []), + 'slot' : 'task', + 'count' : {'per_slot': 1}}] + + if 'environment' in td: system['environment'] = td['environment'] + if 'sandbox' in td: system['cwd'] = td['sandbox'] + if 'shell' in td: system['shell'] = td['shell'] + if 'stdin' in td: system['stdin'] = td['stdin'] + if 'stdout' in td: system['stdout'] = td['stdout'] + if 'stderr' in td: system['stderr'] = td['stderr'] + + attributes = {'system' : system, + 'user' : user} + resources = [{'count': td.get('ranks', 1), + 'type' : 'slot', + 'label': 'task', + 'with' : [{ + 'count': int(td.get('cores_per_rank', 1)), + 'type' : 'core'}]}] + # 'count': int(td.get('gpus_per_rank', 0)) or None, + # 'type' : 'gpu' + + gpr = td.get('gpus_per_rank', 0) + if gpr: + resources[0]['with'].append({'count': math.ceil(gpr), # flux needs int + 'type' : 'gpu'}) + + spec = _flux_job.JobspecV1(resources=resources, + attributes=attributes, + tasks=tasks, + version=version) + return spec + + # ------------------------------------------------------------------------------ # -class _FluxService(object): +class FluxService(object): # -------------------------------------------------------------------------- # - def __init__(self, uid : str, - log : Logger) -> None: + def __init__(self, uid : str = None, + log : Logger = None, + launcher: str = None + ) -> None: - self._uid = uid - self._log = log + self._uid = uid or generate_id('ru.flux') + self._log = log or Logger('radical.utils.flux') + self._launcher = launcher or '' - self._fexe = which('flux') - self._uri = None - self._tout = 60 + self._fexe = which('flux') + self._uri = None + self._proc = None + self._ready = mt.Event() if not _flux: raise RuntimeError('flux module not found') from self._exception @@ -53,22 +109,16 @@ def __init__(self, uid : str, if not self._fexe: raise RuntimeError('flux executable not found') - self._start() - # -------------------------------------------------------------------------- # @property - def uri(self) -> str: - return self._uri + def uid(self) -> str: + return self._uid @property - def timeout(self) -> int: - return self._tout - - @timeout.setter - def timeout(self, tout) -> None: - self._tout = tout + def uri(self) -> str: + return self._uri # -------------------------------------------------------------------------- @@ -81,21 +131,31 @@ def _proc_line_cb(self, prefix: str, for line in lines: if line.startswith('FLUX_URI:'): self._uri = line.strip().split(':', 1)[1] + self._log.info('%s: found flux uri: %s', self._uid, self.uri) + self._ready.set() # -------------------------------------------------------------------------- # def _proc_state_cb(self, proc: Process, state: str) -> None: + self._log.info('flux instance state update: %s', state) + if state in Process.FINAL: + + self._log.info('flux instance stopped: %s', state) + self.stop() # -------------------------------------------------------------------------- # - def _start(self) -> None: + def start(self, timeout: float = None) -> None: fcmd = 'echo FLUX_URI:\\$FLUX_URI && sleep inf' cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) + if self._launcher: + cmd = '%s %s' % (self._launcher, cmd) + self._log.info('%s: start flux instance: %s', self._uid, cmd) p = Process(cmd) @@ -105,15 +165,35 @@ def _start(self) -> None: p.polldelay = 0.1 p.start() - start = time.time() - while time.time() - start < self._tout: - time.sleep(0.1) - if self._uri is not None: - break + self._proc = p + self._ptime = time.time() + + return self.ready(timeout=timeout) + + + # -------------------------------------------------------------------------- + # + def ready(self, timeout: float = None) -> None: + + if timeout is not None: + if timeout < 0: self._ready.wait() + else : self._ready.wait(timeout) + + return self._ready.is_set() + + + # -------------------------------------------------------------------------- + # + def stop(self) -> None: + + if not self._proc: + return + + self._proc.cancel() + self._proc.wait() - if self.uri is None: - self._log.error('%s: flux instance did not start', self._uid) - raise RuntimeError('%s: flux instance did not start', self._uid) + self.uri = None + self._proc = None self._log.info('%s: found flux uri: %s', self._uid, self.uri) @@ -124,14 +204,15 @@ class FluxHelper(object): # -------------------------------------------------------------------------- # - def __init__(self, uri : str = None, - log : Logger = None) -> None: + def __init__(self, uri : str, + log : Logger = None) -> None: + + self._t0 = time.time() self._uri = uri self._log = log or Logger('radical.utils.flux') self._uid = generate_id('ru.flux') - self._handle = None - self._service = None + self._handle = _flux.Flux(self._uri) # journal watcher self._jthread = None @@ -161,17 +242,11 @@ def __init__(self, uri : str = None, # -------------------------------------------------------------------------- # - def start(self) -> None: + def start(self, launcher: str = None) -> None: - if self._handle is not None: + if self._jthread is not None: return - if not self._uri: - self._flux_service = _FluxService(uid=self._uid, log=self._log) - self._uri = self._flux_service.uri - - self._handle = _flux.Flux(self._uri) - self._jthread = mt.Thread(target=self._jwatcher) self._jthread.daemon = True self._jthread.start() @@ -205,6 +280,10 @@ def stop(self): def uid(self) -> str: return self._uid + @property + def uri(self) -> str: + return self._uri + # -------------------------------------------------------------------------- # @@ -234,24 +313,25 @@ def _swatcher(self): # if we get new specs, submit them, return IDs to iqueue, and also # forward ID to ewatcher - fh = _flux.Flux(self._uri) while True: try: - specs = self._squeue.get(timeout=1.0) + specs = self._squeue.get(block=True, timeout=1.0) except queue.Empty: continue - try: - with self._idlock: + except Exception: + self._log.exception("exception") + raise - futs = list() - tids = list() + with self._idlock: + try: + futs = list() for spec in specs: - tid = spec.attributes['system']['job']['name'] + tid = spec.attributes['user']['uid'] fut = _flux_job.submit_async(fh, spec, waitable=True) futs.append([fut, tid]) @@ -264,12 +344,14 @@ def _swatcher(self): # trigger an event check self._equeue.put(fid) - # trigger submit completion - self._sevent.set() - except Exception: - self._log.exception("exception") - raise + except Exception: + self._log.exception("exception") + raise + + finally: + # trigger submit completion + self._sevent.set() # -------------------------------------------------------------------------- @@ -294,6 +376,7 @@ def _ewatcher(self): def register_cb(self, cb: callable) -> None: with self._elock: + self._log.debug('==== register cb %s', cb) self._cbacks.append(cb) @@ -314,6 +397,8 @@ def _handle_events(self, fh : 'flux.Flux', with self._elock: + self._log.debug_9('==== event %s: %s', fid, event) + # if triggered by submit, check if we have anything to do if not event: if fid not in self._events: @@ -321,100 +406,52 @@ def _handle_events(self, fh : 'flux.Flux', # check if we can handle the event - otherwise store it if not self._cbacks: + self._log.debug('==== no cb %s: %s', fid, event) self._events[fid].append(event) return # check if application knows the task - otherwise store the event if fid not in self._task_ids: + self._log.debug('==== no id %s: %s', fid, event) self._events[fid].append(event) return # task is known, flush stored events for ev in self._events[fid]: + self._log.debug('==== play %s: %s - %s', fid, event, self._cbacks) for cb in self._cbacks: cb(self._task_ids[fid], ev) self._events[fid] = [] # process the current event if event: + self._log.debug('==== relay %s: %s - %s', fid, event, self._cbacks) for cb in self._cbacks: cb(self._task_ids[fid], event) # -------------------------------------------------------------------------- # - def spec_from_command(self, cmd: str) -> 'flux.job.JobspecV1': - - return _flux_job.JobspecV1.from_command(shlex.split(cmd)) - - - # -------------------------------------------------------------------------- - # - def spec_from_dict(self, td: dict) -> 'flux.job.JobspecV1': - - version = 1 - tasks = [{'command': [td['executable']] + td.get('arguments', []), - 'slot' : 'task', - 'count' : {'per_slot': 1}}] - - system = {'duration': td.get('duration', 0.0)} - - if 'environment' in td: system['environment'] = td['environment'] - if 'sandbox' in td: system['cwd'] = td['sandbox'] - if 'shell' in td: system['shell'] = td['shell'] - if 'stdin' in td: system['stdin'] = td['stdin'] - if 'stdout' in td: system['stdout'] = td['stdout'] - if 'stderr' in td: system['stderr'] = td['stderr'] - if 'uid' in td: system['job'] = {'name': td['uid']} - - attributes = {'system' : system} - resources = [{'count': td.get('ranks', 1), - 'type' : 'slot', - 'label': 'task', - 'with' : [{ - 'count': int(td.get('cores_per_rank', 1)), - 'type' : 'core'}]}] - # 'count': int(td.get('gpus_per_rank', 0)) or None, - # 'type' : 'gpu' - - if 'gpus_per_rank' in td: - resources[0]['with'].append({ - # flux likes integer GPU counts - 'count': math.ceil(td['gpus_per_rank']), - 'type' : 'gpu'}) - - # import json - # return json.dumps({ - # 'version' : version, - # 'resources' : resources, - # 'attributes': attributes, - # 'tasks' : tasks}) - - spec = _flux_job.JobspecV1(resources=resources, - attributes=attributes, - tasks=tasks, - version=version) - return spec + def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: + if not self._handle: + raise RuntimeError('flux instance not started') - # -------------------------------------------------------------------------- - # - def submit(self, specs: List[Dict[str, Any]]) -> List[str]: - - # ensure we have a job name which we use as task ID + self._log.debug('submit %d specs', len(specs)) + tids = list() for spec in specs: - tid = spec.attributes['system'].get('job', {}).get('name') + tid = spec.attributes['user'].get('uid') if not tid: tid = generate_id(ID_SIMPLE) - if 'job' not in spec.attributes['system']: - spec.attributes['system']['job'] = dict() - spec.attributes['system']['job']['name'] = tid + if 'user' not in spec.attributes: + spec.attributes['user'] = dict() + spec.attributes['user']['uid'] = tid + tids.append(tid) self._sevent.clear() self._squeue.put(specs) self._sevent.wait() # FIXME: timeout? - tids = [spec.attributes['system']['job']['name'] for spec in specs] return tids @@ -422,6 +459,9 @@ def submit(self, specs: List[Dict[str, Any]]) -> List[str]: # def cancel(self, tids: [str|List[str]]) -> None: + if not self._handle: + raise RuntimeError('flux instance not started') + with self._idlock: for tid in as_list(tids): fid = self._flux_ids[tid] @@ -432,10 +472,15 @@ def cancel(self, tids: [str|List[str]]) -> None: # def wait(self, tids: [str|List[str]]) -> None: + if not self._handle: + raise RuntimeError('flux instance not started') + + tids = as_list(tids) with self._idlock: - for tid in as_list(tids): - fid = self._flux_ids[tid] - _flux_job.wait(self._handle, fid) + fids = [self._flux_ids[tid] for tid in tids] + + for fid in fids: + _flux_job.wait(self._handle, fid) # ------------------------------------------------------------------------------ From c6b8ae7a946b49e516ce22c495c1d23ff308bedf Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sat, 29 Mar 2025 13:51:18 +0100 Subject: [PATCH 16/44] snap --- bin/radical-stack | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bin/radical-stack b/bin/radical-stack index a440bd634..a1c7bf17a 100755 --- a/bin/radical-stack +++ b/bin/radical-stack @@ -6,7 +6,7 @@ import radical.utils as ru namespaces = sys.argv[1:] if not namespaces: - namespaces = ['radical'] + namespaces = ['radical', 'rc'] stack = ru.stack(namespaces) From ee2b2e062fbeb68df17693b879b3d62a1f1e0b8a Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sat, 29 Mar 2025 22:07:43 +0100 Subject: [PATCH 17/44] snap --- src/radical/utils/env.py | 40 +++++++++++++++++++--------------------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/src/radical/utils/env.py b/src/radical/utils/env.py index 796f1756d..0cfdcfc5a 100644 --- a/src/radical/utils/env.py +++ b/src/radical/utils/env.py @@ -544,7 +544,12 @@ def __exit__(self, exc_type: Optional[Exception], exc_tb : Optional[Any] ) -> None: - if exc_type and self._child: + if self._parent: + while self._data is None: + try : self._data = self._q.get(timeout=1) + except queue.Empty: pass + + elif exc_type and self._child: stacktrace = ' '.join(traceback.format_exception( exc_type, exc_val, exc_tb)) self._q.put([None, exc_type, exc_val, stacktrace]) @@ -553,41 +558,34 @@ def __exit__(self, exc_type: Optional[Exception], os._exit(0) - if self._parent: - - while True: - try: - self._data = self._q.get(timeout=1) - break - except queue.Empty: - self._data = None - pass - - # -------------------------------------------------------------------------- # def put(self, data: str) -> None: - if self._child: - self._q.put([data, None, None, None]) - self._q.close() - self._q.join_thread() - os._exit(0) + assert self._child + + self._q.put([data, None, None, None]) + self._q.close() + self._q.join_thread() + os._exit(0) # -------------------------------------------------------------------------- # def get(self) -> Any: + assert self._parent + if self._data is None: return - data, exc_type, exc_val, stacktrace = self._data if exc_type: - sys.stdout.write('%s [%s]\n' % (exc_type, exc_val)) - sys.stdout.write('%s\n\n' % stacktrace) - raise exc_type # pylint: disable=raising-bad-type + sys.stderr.write('envp excepted %s(%s)\n' % (exc_type, exc_val)) + sys.stderr.write(stacktrace) + sys.stderr.flush() + raise RuntimeError('envp failed: %s(%s) - check stderr' + % (exc_type, exc_val)) return data From 925513983f75363ba8ea7ede12cbd5ab4ce0990f Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sun, 30 Mar 2025 09:48:14 +0200 Subject: [PATCH 18/44] snap --- src/radical/utils/flux.py | 227 +++++++++++++++++++++++++++----------- 1 file changed, 163 insertions(+), 64 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 3a737819c..d65b1ab00 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -44,7 +44,7 @@ def spec_from_command(cmd: str) -> 'flux.job.JobspecV1': def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': version = 1 - user = {'uid' : td.get('uid', generate_id(ID_SIMPLE))} + user = {'uid' : td.get('uid', generate_id('ru_flux', ID_SIMPLE))} system = {'duration': td.get('duration', 0.0)} tasks = [{'command': [td['executable']] + td.get('arguments', []), 'slot' : 'task', @@ -101,10 +101,10 @@ def __init__(self, uid : str = None, self._ready = mt.Event() if not _flux: - raise RuntimeError('flux module not found') from self._exception + raise RuntimeError('flux module not found') from _flux_exc if not _flux_job: - raise RuntimeError('flux.job module not found') from self._exception + raise RuntimeError('flux.job module not found') from _flux_exc if not self._fexe: raise RuntimeError('flux executable not found') @@ -210,9 +210,17 @@ def __init__(self, uri : str, self._t0 = time.time() self._uri = uri - self._log = log or Logger('radical.utils.flux') + self._log = log or Logger('radical.utils.flux') self._uid = generate_id('ru.flux') self._handle = _flux.Flux(self._uri) + self._api_lock = mt.Lock() + + if 'JournalConsumer' in dir(_flux_job): + self._version = 1 + else: + self._version = 0 + + self._version = 0 # FIXME # journal watcher self._jthread = None @@ -244,34 +252,38 @@ def __init__(self, uri : str, # def start(self, launcher: str = None) -> None: - if self._jthread is not None: - return + with self._api_lock: - self._jthread = mt.Thread(target=self._jwatcher) - self._jthread.daemon = True - self._jthread.start() + if self._jthread is not None: + return + + self._jthread = mt.Thread(target=self._jwatcher) + self._jthread.daemon = True + self._jthread.start() - self._ethread = mt.Thread(target=self._ewatcher) - self._ethread.daemon = True - self._ethread.start() + self._ethread = mt.Thread(target=self._ewatcher) + self._ethread.daemon = True + self._ethread.start() - self._sthread = mt.Thread(target=self._swatcher) - self._sthread.daemon = True - self._sthread.start() + self._sthread = mt.Thread(target=self._swatcher) + self._sthread.daemon = True + self._sthread.start() # -------------------------------------------------------------------------- # def stop(self): - if self._handle is None: - self._jterm.set() - self._jthread.join() + with self._api_lock: - # FIXME: shutdown flux instance - self._flux_service = None - self._uri = None - self._handle = None + if self._handle is None: + self._jterm.set() + self._jthread.join() + + # FIXME: shutdown flux instance + self._flux_service = None + self._uri = None + self._handle = None # -------------------------------------------------------------------------- @@ -289,6 +301,24 @@ def uri(self) -> str: # def _jwatcher(self): + if self._version == 0: self._jwatcher_v0() + elif self._version == 1: self._jwatcher_v1() + + + # -------------------------------------------------------------------------- + # + def _jwatcher_v0(self): + + # all event handling is done in the submission thread which attaches an + # event cb to each future + while True: + time.sleep(1.0) + + + # -------------------------------------------------------------------------- + # + def _jwatcher_v1(self): + # NOTE: *never* used self._handle in this thread, as it is not thread # safe. Instead, use the private handle created here fh = _flux.Flux(self._uri) @@ -311,6 +341,76 @@ def _jwatcher(self): # def _swatcher(self): + if self._version == 0: self._swatcher_v0() + elif self._version == 1: self._swatcher_v1() + + + # -------------------------------------------------------------------------- + # + def _swatcher_v0(self): + ''' + if we get new specs, submit them, return IDs to iqueue, and also + forward ID to ewatcher + ''' + + events = ['submit', 'depend', 'alloc', 'start', # 'cleanup', + 'finish', 'release', 'free', 'clean', 'priority', 'exception'] + + exe = _flux_job.executor.FluxExecutor(handle_kwargs={'url': self._uri}) + fh = _flux.Flux(self._uri) + fids = list() + + def event_cb(fid, fut, event): + self._handle_events(fh, fid, event) + + def jobid_cb(tid, fut): + fid = fut.jobid() + self._log.debug('jobid %s -> %s', tid, fid) + fids.append([fid, tid]) + self._flux_ids[tid] = fid + self._task_ids[fid] = tid + self._equeue.put(fid) + + for event in events: + fut.add_event_callback(event, partial(event_cb, fid)) + + while True: + + try: + specs = self._squeue.get(block=True, timeout=1.0) + + except queue.Empty: + continue + + except: + self._log.exception("exception") + raise + + with self._idlock: + + try: + + for spec in specs: + tid = spec.attributes['user']['uid'] + fut = exe.submit(spec, waitable=True) + fut.add_jobid_callback(partial(jobid_cb, tid)) + + while len(fids) < len(specs): + time.sleep(0.1) + + except Exception: + self._log.exception("exception") + raise + + finally: + # trigger submit completion + self._sevent.set() + + + # -------------------------------------------------------------------------- + # + def _swatcher_v1(self): + # if we get new specs, submit them, return IDs to iqueue, and also # forward ID to ewatcher fh = _flux.Flux(self._uri) @@ -322,11 +422,12 @@ def _swatcher(self): except queue.Empty: continue - except Exception: + except: self._log.exception("exception") raise with self._idlock: + try: futs = list() @@ -344,7 +445,6 @@ def _swatcher(self): # trigger an event check self._equeue.put(fid) - except Exception: self._log.exception("exception") raise @@ -375,16 +475,16 @@ def _ewatcher(self): # def register_cb(self, cb: callable) -> None: - with self._elock: - self._log.debug('==== register cb %s', cb) - self._cbacks.append(cb) + with self._api_lock, self._elock: + self._log.debug('register cb %s', cb) + self._cbacks.append(cb) # -------------------------------------------------------------------------- # def unregister_cb(self, cb: callable) -> None: - with self._elock: + with self._api_lock, self._elock: self._cbacks.remove(cb) @@ -397,7 +497,7 @@ def _handle_events(self, fh : 'flux.Flux', with self._elock: - self._log.debug_9('==== event %s: %s', fid, event) + # self._log.debug_9('event %s: %s', fid, event) # if triggered by submit, check if we have anything to do if not event: @@ -406,81 +506,80 @@ def _handle_events(self, fh : 'flux.Flux', # check if we can handle the event - otherwise store it if not self._cbacks: - self._log.debug('==== no cb %s: %s', fid, event) self._events[fid].append(event) return # check if application knows the task - otherwise store the event if fid not in self._task_ids: - self._log.debug('==== no id %s: %s', fid, event) self._events[fid].append(event) return + tid = self._task_ids[fid] + # task is known, flush stored events for ev in self._events[fid]: - self._log.debug('==== play %s: %s - %s', fid, event, self._cbacks) for cb in self._cbacks: - cb(self._task_ids[fid], ev) + try : cb(tid, ev) + except: self._log.exception('cb failed: %s') self._events[fid] = [] # process the current event if event: - self._log.debug('==== relay %s: %s - %s', fid, event, self._cbacks) for cb in self._cbacks: - cb(self._task_ids[fid], event) + try : cb(tid, event) + except: self._log.exception('cb failed: %s') # -------------------------------------------------------------------------- # def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: - if not self._handle: - raise RuntimeError('flux instance not started') + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') - self._log.debug('submit %d specs', len(specs)) - tids = list() - for spec in specs: - tid = spec.attributes['user'].get('uid') - if not tid: - tid = generate_id(ID_SIMPLE) - if 'user' not in spec.attributes: - spec.attributes['user'] = dict() - spec.attributes['user']['uid'] = tid - tids.append(tid) + self._log.debug('== submit %d specs', len(specs)) + tids = [spec.attributes['user']['uid'] for spec in specs] - self._sevent.clear() - self._squeue.put(specs) - self._sevent.wait() # FIXME: timeout? + self._sevent.clear() + self._squeue.put(specs) + self._sevent.wait() # FIXME: timeout? - return tids + return tids # -------------------------------------------------------------------------- # def cancel(self, tids: [str|List[str]]) -> None: - if not self._handle: - raise RuntimeError('flux instance not started') + with self._api_lock: - with self._idlock: - for tid in as_list(tids): - fid = self._flux_ids[tid] - _flux_job.cancel_async(self._handle, fid, reason='user cancel') + if not self._handle: + raise RuntimeError('flux instance not started') + + with self._idlock: + for tid in as_list(tids): + fid = self._flux_ids[tid] + _flux_job.cancel_async(self._handle, fid, reason='user cancel') # -------------------------------------------------------------------------- # def wait(self, tids: [str|List[str]]) -> None: - if not self._handle: - raise RuntimeError('flux instance not started') + with self._api_lock: - tids = as_list(tids) - with self._idlock: - fids = [self._flux_ids[tid] for tid in tids] + if not self._handle: + raise RuntimeError('flux instance not started') + + tids = as_list(tids) + with self._idlock: + fids = [self._flux_ids[tid] for tid in tids] - for fid in fids: - _flux_job.wait(self._handle, fid) + for fid in fids: + print('wait for %s [%s]' % (fid, tids)) + _flux_job.wait(self._handle, fid) # ------------------------------------------------------------------------------ From 268360fecf5bda10cbaaf5950c17f9105e6f0b15 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sun, 30 Mar 2025 21:11:40 +0200 Subject: [PATCH 19/44] snap --- src/radical/utils/flux.py | 318 ++++++++++++++++++++++++++++++-------- 1 file changed, 253 insertions(+), 65 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index d65b1ab00..72d5f30e1 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -200,7 +200,7 @@ def stop(self) -> None: # ------------------------------------------------------------------------------ # -class FluxHelper(object): +class FluxHelperV0(object): # -------------------------------------------------------------------------- # @@ -215,16 +215,6 @@ def __init__(self, uri : str, self._handle = _flux.Flux(self._uri) self._api_lock = mt.Lock() - if 'JournalConsumer' in dir(_flux_job): - self._version = 1 - else: - self._version = 0 - - self._version = 0 # FIXME - - # journal watcher - self._jthread = None - # event handle thread self._ethread = None self._equeue = queue.Queue() @@ -257,10 +247,6 @@ def start(self, launcher: str = None) -> None: if self._jthread is not None: return - self._jthread = mt.Thread(target=self._jwatcher) - self._jthread.daemon = True - self._jthread.start() - self._ethread = mt.Thread(target=self._ewatcher) self._ethread.daemon = True self._ethread.start() @@ -297,57 +283,9 @@ def uri(self) -> str: return self._uri - # -------------------------------------------------------------------------- - # - def _jwatcher(self): - - if self._version == 0: self._jwatcher_v0() - elif self._version == 1: self._jwatcher_v1() - - - # -------------------------------------------------------------------------- - # - def _jwatcher_v0(self): - - # all event handling is done in the submission thread which attaches an - # event cb to each future - while True: - time.sleep(1.0) - - - # -------------------------------------------------------------------------- - # - def _jwatcher_v1(self): - - # NOTE: *never* used self._handle in this thread, as it is not thread - # safe. Instead, use the private handle created here - fh = _flux.Flux(self._uri) - - # start watching the event journal - journal = _flux_job.JournalConsumer(fh) - journal.start() - - while True: - - try: - event = journal.poll(timeout=1.0) - self._handle_events(fh, event.jobid, event) - - except TimeoutError: - pass - - # -------------------------------------------------------------------------- # def _swatcher(self): - - if self._version == 0: self._swatcher_v0() - elif self._version == 1: self._swatcher_v1() - - - # -------------------------------------------------------------------------- - # - def _swatcher_v0(self): ''' if we get new specs, submit them, return IDs to iqueue, and also forward ID to ewatcher @@ -409,7 +347,249 @@ def jobid_cb(tid, fut): # -------------------------------------------------------------------------- # - def _swatcher_v1(self): + def _ewatcher(self): + + # if we get a new job ID, check if we have events for it + + fh = _flux.Flux(self._uri) + while True: + + try: + fid = self._equeue.get(timeout=1.0) + self._handle_events(fh, fid) + + except queue.Empty: + continue + + + # -------------------------------------------------------------------------- + # + def register_cb(self, cb: callable) -> None: + + with self._api_lock, self._elock: + self._log.debug('register cb %s', cb) + self._cbacks.append(cb) + + + # -------------------------------------------------------------------------- + # + def unregister_cb(self, cb: callable) -> None: + + with self._api_lock, self._elock: + self._cbacks.remove(cb) + + + # -------------------------------------------------------------------------- + # + def _handle_events(self, fh : 'flux.Flux', + fid : 'flux.job.JobID', + event: 'flux.job.journal.JournalEvent' = None + ) -> None: + + with self._elock: + + # self._log.debug_9('event %s: %s', fid, event) + + # if triggered by submit, check if we have anything to do + if not event: + if fid not in self._events: + return + + # check if we can handle the event - otherwise store it + if not self._cbacks: + self._events[fid].append(event) + return + + # check if application knows the task - otherwise store the event + if fid not in self._task_ids: + self._events[fid].append(event) + return + + tid = self._task_ids[fid] + + # task is known, flush stored events + for ev in self._events[fid]: + for cb in self._cbacks: + try : cb(tid, ev) + except: self._log.exception('cb failed: %s') + self._events[fid] = [] + + # process the current event + if event: + for cb in self._cbacks: + try : cb(tid, event) + except: self._log.exception('cb failed: %s') + + + # -------------------------------------------------------------------------- + # + def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + self._log.debug('== submit %d specs', len(specs)) + tids = [spec.attributes['user']['uid'] for spec in specs] + + self._sevent.clear() + self._squeue.put(specs) + self._sevent.wait() # FIXME: timeout? + + return tids + + + # -------------------------------------------------------------------------- + # + def cancel(self, tids: [str|List[str]]) -> None: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + with self._idlock: + for tid in as_list(tids): + fid = self._flux_ids[tid] + _flux_job.cancel_async(self._handle, fid, reason='user cancel') + + + # -------------------------------------------------------------------------- + # + def wait(self, tids: [str|List[str]]) -> None: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + tids = as_list(tids) + with self._idlock: + fids = [self._flux_ids[tid] for tid in tids] + + for fid in fids: + print('wait for %s [%s]' % (fid, tids)) + _flux_job.wait(self._handle, fid) + + +# ------------------------------------------------------------------------------ +# +class FluxHelperV1(object): + + # -------------------------------------------------------------------------- + # + def __init__(self, uri : str, + log : Logger = None) -> None: + + self._t0 = time.time() + + self._uri = uri + self._log = log or Logger('radical.utils.flux') + self._uid = generate_id('ru.flux') + self._handle = _flux.Flux(self._uri) + self._api_lock = mt.Lock() + + # journal watcher + self._jthread = None + + # event handle thread + self._ethread = None + self._equeue = queue.Queue() + + # submit thread + self._sthread = None + self._squeue = queue.Queue() + self._sevent = mt.Event() + + self._idlock = mt.Lock() # lock ID dicts + self._elock = mt.Lock() # lock event dict + self._task_ids = dict() # flux ID -> task ID + self._flux_ids = dict() # task ID -> flux ID + self._events = defaultdict(list) # flux ID -> event list + self._cbacks = list() # list of callbacks + + if not _flux: + raise RuntimeError('flux module not found') from _flux_exc + + if not _flux_job: + raise RuntimeError('flux.job module not found') from _flux_exc + + + # -------------------------------------------------------------------------- + # + def start(self, launcher: str = None) -> None: + + with self._api_lock: + + if self._jthread is not None: + return + + self._jthread = mt.Thread(target=self._jwatcher) + self._jthread.daemon = True + self._jthread.start() + + self._ethread = mt.Thread(target=self._ewatcher) + self._ethread.daemon = True + self._ethread.start() + + self._sthread = mt.Thread(target=self._swatcher) + self._sthread.daemon = True + self._sthread.start() + + + # -------------------------------------------------------------------------- + # + def stop(self): + + with self._api_lock: + + if self._handle is None: + self._jterm.set() + self._jthread.join() + + # FIXME: shutdown flux instance + self._flux_service = None + self._uri = None + self._handle = None + + + # -------------------------------------------------------------------------- + # + @property + def uid(self) -> str: + return self._uid + + @property + def uri(self) -> str: + return self._uri + + + # -------------------------------------------------------------------------- + # + def _jwatcher(self): + + # NOTE: *never* used self._handle in this thread, as it is not thread + # safe. Instead, use the private handle created here + fh = _flux.Flux(self._uri) + + # start watching the event journal + journal = _flux_job.JournalConsumer(fh) + journal.start() + + while True: + + try: + event = journal.poll(timeout=1.0) + self._handle_events(fh, event.jobid, event) + + except TimeoutError: + pass + + + # -------------------------------------------------------------------------- + # + def _swatcher(self): # if we get new specs, submit them, return IDs to iqueue, and also # forward ID to ewatcher @@ -578,9 +758,17 @@ def wait(self, tids: [str|List[str]]) -> None: fids = [self._flux_ids[tid] for tid in tids] for fid in fids: - print('wait for %s [%s]' % (fid, tids)) _flux_job.wait(self._handle, fid) # ------------------------------------------------------------------------------ +# +if 'JournalConsumer' in dir(_flux_job): + FluxHelper = FluxHelperV1 +else: + FluxHelper = FluxHelperV0 + +FluxHelper = FluxHelperV1 + +# ------------------------------------------------------------------------------ From 9ba3d4b025482d1c0e14a1eed9ca39ffb14ad7e5 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sun, 30 Mar 2025 21:55:45 +0200 Subject: [PATCH 20/44] snap --- src/radical/utils/flux.py | 26 ++++++++++---------------- 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 72d5f30e1..9b52b29e3 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -244,7 +244,7 @@ def start(self, launcher: str = None) -> None: with self._api_lock: - if self._jthread is not None: + if self._ethread is not None: return self._ethread = mt.Thread(target=self._ewatcher) @@ -301,17 +301,6 @@ def _swatcher(self): def event_cb(fid, fut, event): self._handle_events(fh, fid, event) - def jobid_cb(tid, fut): - fid = fut.jobid() - self._log.debug('jobid %s -> %s', tid, fid) - fids.append([fid, tid]) - self._flux_ids[tid] = fid - self._task_ids[fid] = tid - self._equeue.put(fid) - - for event in events: - fut.add_event_callback(event, partial(event_cb, fid)) - while True: try: @@ -328,13 +317,18 @@ def jobid_cb(tid, fut): try: + fids = list() for spec in specs: tid = spec.attributes['user']['uid'] fut = exe.submit(spec, waitable=True) - fut.add_jobid_callback(partial(jobid_cb, tid)) + fid = fut.jobid() + self._flux_ids[tid] = fid + self._task_ids[fid] = tid + fids.append(fid) + self._equeue.put(fid) - while len(fids) < len(specs): - time.sleep(0.1) + for event in events: + fut.add_event_callback(event, partial(event_cb, fid)) except Exception: self._log.exception("exception") @@ -768,7 +762,7 @@ def wait(self, tids: [str|List[str]]) -> None: else: FluxHelper = FluxHelperV0 -FluxHelper = FluxHelperV1 +FluxHelper = FluxHelperV0 # ------------------------------------------------------------------------------ From b21fe4015ff4e3e650d43ee99be03d6afe28d4d6 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sun, 30 Mar 2025 23:16:28 +0200 Subject: [PATCH 21/44] snap --- src/radical/utils/flux.py | 227 +++++++++++++------------------------- 1 file changed, 79 insertions(+), 148 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 9b52b29e3..833bd0a17 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -214,15 +214,8 @@ def __init__(self, uri : str, self._uid = generate_id('ru.flux') self._handle = _flux.Flux(self._uri) self._api_lock = mt.Lock() - - # event handle thread - self._ethread = None - self._equeue = queue.Queue() - - # submit thread - self._sthread = None - self._squeue = queue.Queue() - self._sevent = mt.Event() + self._exe = _flux_job.executor.FluxExecutor( + handle_kwargs={'url': self._uri}) self._idlock = mt.Lock() # lock ID dicts self._elock = mt.Lock() # lock event dict @@ -242,18 +235,7 @@ def __init__(self, uri : str, # def start(self, launcher: str = None) -> None: - with self._api_lock: - - if self._ethread is not None: - return - - self._ethread = mt.Thread(target=self._ewatcher) - self._ethread.daemon = True - self._ethread.start() - - self._sthread = mt.Thread(target=self._swatcher) - self._sthread.daemon = True - self._sthread.start() + pass # -------------------------------------------------------------------------- @@ -262,14 +244,10 @@ def stop(self): with self._api_lock: - if self._handle is None: - self._jterm.set() - self._jthread.join() - - # FIXME: shutdown flux instance - self._flux_service = None - self._uri = None - self._handle = None + # FIXME: shutdown flux instance + self._flux_service = None + self._uri = None + self._handle = None # -------------------------------------------------------------------------- @@ -283,86 +261,13 @@ def uri(self) -> str: return self._uri - # -------------------------------------------------------------------------- - # - def _swatcher(self): - ''' - if we get new specs, submit them, return IDs to iqueue, and also - forward ID to ewatcher - ''' - - events = ['submit', 'depend', 'alloc', 'start', # 'cleanup', - 'finish', 'release', 'free', 'clean', 'priority', 'exception'] - - exe = _flux_job.executor.FluxExecutor(handle_kwargs={'url': self._uri}) - fh = _flux.Flux(self._uri) - fids = list() - - def event_cb(fid, fut, event): - self._handle_events(fh, fid, event) - - while True: - - try: - specs = self._squeue.get(block=True, timeout=1.0) - - except queue.Empty: - continue - - except: - self._log.exception("exception") - raise - - with self._idlock: - - try: - - fids = list() - for spec in specs: - tid = spec.attributes['user']['uid'] - fut = exe.submit(spec, waitable=True) - fid = fut.jobid() - self._flux_ids[tid] = fid - self._task_ids[fid] = tid - fids.append(fid) - self._equeue.put(fid) - - for event in events: - fut.add_event_callback(event, partial(event_cb, fid)) - - except Exception: - self._log.exception("exception") - raise - - finally: - # trigger submit completion - self._sevent.set() - - - # -------------------------------------------------------------------------- - # - def _ewatcher(self): - - # if we get a new job ID, check if we have events for it - - fh = _flux.Flux(self._uri) - while True: - - try: - fid = self._equeue.get(timeout=1.0) - self._handle_events(fh, fid) - - except queue.Empty: - continue - - # -------------------------------------------------------------------------- # def register_cb(self, cb: callable) -> None: with self._api_lock, self._elock: - self._log.debug('register cb %s', cb) - self._cbacks.append(cb) + self._log.debug('register cb %s', cb) + self._cbacks.append(cb) # -------------------------------------------------------------------------- @@ -380,39 +285,35 @@ def _handle_events(self, fh : 'flux.Flux', event: 'flux.job.journal.JournalEvent' = None ) -> None: - with self._elock: - - # self._log.debug_9('event %s: %s', fid, event) - - # if triggered by submit, check if we have anything to do - if not event: - if fid not in self._events: - return + self._log.debug('event for %s: %s', fid, event) + # print('event %s: %s' % (fid, event)) - # check if we can handle the event - otherwise store it - if not self._cbacks: - self._events[fid].append(event) + # if triggered by submit, check if we have anything to do + if not event: + if fid not in self._events: + # print('no event') return - # check if application knows the task - otherwise store the event - if fid not in self._task_ids: - self._events[fid].append(event) - return - - tid = self._task_ids[fid] + # check if we can handle the event - otherwise store it + if not self._cbacks: + # print('no cbacks') + self._events[fid].append(event) + return - # task is known, flush stored events - for ev in self._events[fid]: - for cb in self._cbacks: - try : cb(tid, ev) - except: self._log.exception('cb failed: %s') - self._events[fid] = [] + # task is known, flush stored events + for ev in self._events[fid]: + # print('flush stored events') + for cb in self._cbacks: + try : cb(fid, ev) + except: self._log.exception('cb failed: %s') + self._events[fid] = [] - # process the current event - if event: - for cb in self._cbacks: - try : cb(tid, event) - except: self._log.exception('cb failed: %s') + # process the current event + if event: + # print('process current event') + for cb in self._cbacks: + try : cb(fid, event) + except: self._log.exception('cb failed: %s') # -------------------------------------------------------------------------- @@ -425,18 +326,53 @@ def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: raise RuntimeError('flux instance not started') self._log.debug('== submit %d specs', len(specs)) - tids = [spec.attributes['user']['uid'] for spec in specs] - self._sevent.clear() - self._squeue.put(specs) - self._sevent.wait() # FIXME: timeout? + events = ['submit', 'depend', 'alloc', 'start', # 'cleanup', + 'finish', 'release', 'free', 'clean', 'priority', 'exception'] - return tids + def event_cb(fid, fut, event): + self._handle_events(self._handle, fid, event) + + futures = list() + def id_cb(fut): + flux_id = fut.jobid() + idx = fut.ru_idx + for ev in events: + tmp_cb = partial(event_cb, flux_id) + fut.add_event_callback(ev, tmp_cb) + futures.append([flux_id, idx, fut]) + self._log.debug('got flux id: %s: %s', idx, flux_id) + + for idx, spec in enumerate(specs): + fut = self._exe.submit(spec, waitable=True) + fut.ru_idx = idx + self._log.debug('%s: submitted : %s', self._uid, idx) + fut.add_jobid_callback(id_cb) + + # wait until we saw all jobid callbacks (assume 10 tasks/sec) + timeout = len(specs) + timeout = max(100, timeout) + start = time.time() + self._log.debug('%s: wait %.2fsec for %d flux IDs', + self._uid, timeout, len(specs)) + while len(futures) < len(specs): + time.sleep(0.1) + self._log.debug('%s: wait %s / %s', self._uid, + len(futures), len(specs)) + if time.time() - start > timeout: + raise RuntimeError('%s: timeout on submission', self._uid) + self._log.info('got %d flux IDs', len(futures)) + + # get flux_ids sorted by submission order (idx) + flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] + + self._log.debug('%s: submitted: %s', self._uid, flux_ids) + return flux_ids # -------------------------------------------------------------------------- # - def cancel(self, tids: [str|List[str]]) -> None: + def cancel(self, fids: [str|List[str]]) -> None: with self._api_lock: @@ -444,26 +380,21 @@ def cancel(self, tids: [str|List[str]]) -> None: raise RuntimeError('flux instance not started') with self._idlock: - for tid in as_list(tids): - fid = self._flux_ids[tid] + for fid in as_list(fids): _flux_job.cancel_async(self._handle, fid, reason='user cancel') # -------------------------------------------------------------------------- # - def wait(self, tids: [str|List[str]]) -> None: + def wait(self, fids: [str|List[str]]) -> None: with self._api_lock: if not self._handle: raise RuntimeError('flux instance not started') - tids = as_list(tids) - with self._idlock: - fids = [self._flux_ids[tid] for tid in tids] - for fid in fids: - print('wait for %s [%s]' % (fid, tids)) + self._log.debug('wait for %s', fid) _flux_job.wait(self._handle, fid) @@ -650,8 +581,8 @@ def _ewatcher(self): def register_cb(self, cb: callable) -> None: with self._api_lock, self._elock: - self._log.debug('register cb %s', cb) - self._cbacks.append(cb) + self._log.debug('register cb %s', cb) + self._cbacks.append(cb) # -------------------------------------------------------------------------- From 11adc53e8e028490dc74817a4a63771dc92fb780 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Sun, 30 Mar 2025 23:30:16 +0200 Subject: [PATCH 22/44] snap --- src/radical/utils/flux.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 833bd0a17..075e6a0e2 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -693,7 +693,6 @@ def wait(self, tids: [str|List[str]]) -> None: else: FluxHelper = FluxHelperV0 -FluxHelper = FluxHelperV0 # ------------------------------------------------------------------------------ From d09c285cbf329c631b79d12cac2c94b4e66eb6d8 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 31 Mar 2025 23:59:07 +0200 Subject: [PATCH 23/44] fix uri detection; --- src/radical/utils/flux.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 075e6a0e2..8578b7e0b 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -130,7 +130,17 @@ def _proc_line_cb(self, prefix: str, for line in lines: if line.startswith('FLUX_URI:'): - self._uri = line.strip().split(':', 1)[1] + parts = line.strip().split(':', 1) + self._log.info('%s: found flux info: %s', self._uid, parts) + + uri = parts[1].split('=', 1)[1] + host = parts[2].split('=', 1)[1] + + url = Url(uri) + url.host = host + url.schema = 'ssh' + self._uri = str(flux_url) + self._log.info('%s: found flux uri: %s', self._uid, self.uri) self._ready.set() @@ -150,7 +160,7 @@ def _proc_state_cb(self, proc: Process, state: str) -> None: # def start(self, timeout: float = None) -> None: - fcmd = 'echo FLUX_URI:\\$FLUX_URI && sleep inf' + fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=$(hostname) && sleep inf' cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) if self._launcher: From 05663bb0fe2f9fc8f362d70e856673d1cea97652 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Tue, 1 Apr 2025 15:11:29 +0200 Subject: [PATCH 24/44] fix uri detection --- src/radical/utils/flux.py | 47 ++++++++++++++++++++++++++------------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 8578b7e0b..47b21bed1 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -11,6 +11,7 @@ from collections import defaultdict from typing import List, Dict, Any +from .url import Url from .misc import as_list from .which import which from .ids import generate_id, ID_SIMPLE @@ -22,6 +23,10 @@ _flux = import_module('flux') _flux_job = import_module('flux.job') _flux_exc = None + if 'JournalConsumer' in dir(_flux_job): + _flux_v = 1 + else: + _flux_v = 0 except Exception as e: _flux = None @@ -97,6 +102,8 @@ def __init__(self, uid : str = None, self._fexe = which('flux') self._uri = None + self._r_uri = None + self._host = None self._proc = None self._ready = mt.Event() @@ -121,6 +128,11 @@ def uri(self) -> str: return self._uri + @property + def r_uri(self) -> str: + return self._r_uri + + # -------------------------------------------------------------------------- # def _proc_line_cb(self, prefix: str, @@ -128,21 +140,26 @@ def _proc_line_cb(self, prefix: str, lines : List[str] ) -> None: - for line in lines: - if line.startswith('FLUX_URI:'): - parts = line.strip().split(':', 1) - self._log.info('%s: found flux info: %s', self._uid, parts) + try: + for line in lines: + self._log.info('=== line: %s', line) + if line.startswith('FLUX_URI='): + parts = line.strip().split(' ', 1) + self._log.info('%s: found flux info: %s', self._uid, parts) - uri = parts[1].split('=', 1)[1] - host = parts[2].split('=', 1)[1] + self._uri = parts[0].split('=', 1)[1] + self._host = parts[1].split('=', 1)[1] - url = Url(uri) - url.host = host - url.schema = 'ssh' - self._uri = str(flux_url) + url = Url(self._uri) + url.host = self._host + url.schema = 'ssh' + self._r_uri = str(url) - self._log.info('%s: found flux uri: %s', self._uid, self.uri) - self._ready.set() + self._log.info('%s: flux uri: %s', self._uid, self._uri) + self._log.info('%s: r uri: %s', self._uid, self._r_uri) + self._ready.set() + except: + self._log.exception('line processing failed') # -------------------------------------------------------------------------- @@ -698,10 +715,8 @@ def wait(self, tids: [str|List[str]]) -> None: # ------------------------------------------------------------------------------ # -if 'JournalConsumer' in dir(_flux_job): - FluxHelper = FluxHelperV1 -else: - FluxHelper = FluxHelperV0 +if _flux_v == 1: FluxHelper = FluxHelperV1 +else : FluxHelper = FluxHelperV0 # ------------------------------------------------------------------------------ From 9994d68cdee2d9bd33d2062a57c9412ac4d02abb Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 2 Apr 2025 08:23:22 +0200 Subject: [PATCH 25/44] add missing `_flux_v` --- src/radical/utils/flux.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 47b21bed1..1503b912a 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -32,6 +32,7 @@ _flux = None _flux_job = None _flux_exc = e + _flux_v = -1 # -------------------------------------------------------------------------- From c363b96e346fffddf03c31ca477d05b1554793a0 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 2 Apr 2025 04:29:01 -0400 Subject: [PATCH 26/44] snap --- src/radical/utils/flux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 8578b7e0b..982d78918 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -160,7 +160,7 @@ def _proc_state_cb(self, proc: Process, state: str) -> None: # def start(self, timeout: float = None) -> None: - fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=$(hostname) && sleep inf' + fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=\\$(hostname) && sleep inf' cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) if self._launcher: From 2062bd180975684a6427cf6cc6ad7e7d1221ee56 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Thu, 3 Apr 2025 10:53:21 +0200 Subject: [PATCH 27/44] recover module path hack for flux --- src/radical/utils/flux.py | 77 ++++++++++++++++++++++++++++++++------- 1 file changed, 63 insertions(+), 14 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 7705f8546..052f9d745 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -1,4 +1,6 @@ +import os +import sys import math import time import shlex @@ -17,22 +19,69 @@ from .ids import generate_id, ID_SIMPLE from .logger import Logger from .modules import import_module +from .shell import sh_callout -try: - _flux = import_module('flux') - _flux_job = import_module('flux.job') - _flux_exc = None - if 'JournalConsumer' in dir(_flux_job): - _flux_v = 1 - else: - _flux_v = 0 - -except Exception as e: - _flux = None - _flux_job = None - _flux_exc = e - _flux_v = -1 +# ------------------------------------------------------------------------------ +# +def import_flux(): + ''' + import the flux module, if available + + returns: flux : loaded python module (`None` if not available) + flux.job : loaded python module (`None` if not available) + exception: exception raised during import (`None` if no error) + version : what ru.FluxHelper version to use (`0` or `1`) + ''' + + flux = None + flux_job = None + flux_exc = None + flux_v = None + + try: + flux = import_module('flux') + flux_job = import_module('flux.job') + if 'JournalConsumer' in dir(flux_job): + flux_v = 1 + else: + flux_v = 0 + + except Exception as e: + flux_exc = e + + + # on failure, try to derive module path from flux executable + if flux is None or flux_job is None: + + to_pop = None + try: + cmd = 'flux python -c "import flux; print(flux.__file__)"' + out, err, ret = sh_callout(cmd) + + if not ret: + flux_path = os.path.dirname(out.strip()) + mod_path = os.path.dirname(flux_path) + sys.path.append(mod_path) + to_pop = mod_path + + flux = import_module('flux') + flux_job = import_module('flux.job') + if 'JournalConsumer' in dir(flux_job): + flux_v = 1 + else: + flux_v = 0 + + except Exception as e: + flux_exc = e + + if to_pop: + sys.path.remove(to_pop) + + return flux, flux_job, flux_exc, flux_v + + +_flux, _flux_job, _flux_exc, _flux_v = import_flux() # -------------------------------------------------------------------------- From efa166b42af1703aff627041df7a63c73c35607a Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Thu, 3 Apr 2025 05:00:23 -0400 Subject: [PATCH 28/44] fix an exception --- src/radical/utils/flux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 7705f8546..63da42ab1 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -388,7 +388,7 @@ def id_cb(fut): self._log.debug('%s: wait %s / %s', self._uid, len(futures), len(specs)) if time.time() - start > timeout: - raise RuntimeError('%s: timeout on submission', self._uid) + raise RuntimeError('%s: timeout on submission' % self._uid) self._log.info('got %d flux IDs', len(futures)) # get flux_ids sorted by submission order (idx) From 12ec9cd96cefb909bf2717fc43b279ce71a3029d Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Fri, 4 Apr 2025 10:31:36 +0200 Subject: [PATCH 29/44] shield from non-events --- src/radical/utils/flux.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 7361d6ab8..f086dd978 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -269,7 +269,7 @@ def stop(self) -> None: self._proc.cancel() self._proc.wait() - self.uri = None + self._uri = None self._proc = None self._log.info('%s: found flux uri: %s', self._uid, self.uri) @@ -583,7 +583,10 @@ def _jwatcher(self): try: event = journal.poll(timeout=1.0) - self._handle_events(fh, event.jobid, event) + if event: + # FIXME: How can that ever *not* be a journal event? + # But it has happened... + self._handle_events(fh, event.jobid, event) except TimeoutError: pass From 6f4494fe0f12a9ba7a7fbf461546dfc2a8b9676e Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 9 Apr 2025 11:36:27 +0200 Subject: [PATCH 30/44] better reporting --- src/radical/utils/flux.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index f086dd978..910d7c56d 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -227,8 +227,10 @@ def _proc_state_cb(self, proc: Process, state: str) -> None: # def start(self, timeout: float = None) -> None: - fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=\\$(hostname) && sleep inf' - cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) + fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=\\$(hostname) ' + fcmd += ' && flux resources list ' + fmcd += ' && sleep inf ' + cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) if self._launcher: cmd = '%s %s' % (self._launcher, cmd) @@ -382,7 +384,7 @@ def _handle_events(self, fh : 'flux.Flux', # print('flush stored events') for cb in self._cbacks: try : cb(fid, ev) - except: self._log.exception('cb failed: %s') + except: self._log.exception('cb failed') self._events[fid] = [] # process the current event @@ -390,7 +392,7 @@ def _handle_events(self, fh : 'flux.Flux', # print('process current event') for cb in self._cbacks: try : cb(fid, event) - except: self._log.exception('cb failed: %s') + except: self._log.exception('cb failed') # -------------------------------------------------------------------------- @@ -705,14 +707,14 @@ def _handle_events(self, fh : 'flux.Flux', for ev in self._events[fid]: for cb in self._cbacks: try : cb(tid, ev) - except: self._log.exception('cb failed: %s') + except: self._log.exception('cb failed') self._events[fid] = [] # process the current event if event: for cb in self._cbacks: try : cb(tid, event) - except: self._log.exception('cb failed: %s') + except: self._log.exception('cb failed') # -------------------------------------------------------------------------- From 1e0f4a5bebc3f74c47d1ab2c32d051f610f68ef1 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 9 Apr 2025 11:58:51 +0200 Subject: [PATCH 31/44] better reporting --- src/radical/utils/flux.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 910d7c56d..06a05d865 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -598,6 +598,8 @@ def _jwatcher(self): # def _swatcher(self): + self._log.debug('=== swatcher started') + # if we get new specs, submit them, return IDs to iqueue, and also # forward ID to ewatcher fh = _flux.Flux(self._uri) @@ -605,6 +607,7 @@ def _swatcher(self): try: specs = self._squeue.get(block=True, timeout=1.0) + self._log.debug('=== got %d specs', len(specs)) except queue.Empty: continue @@ -621,7 +624,6 @@ def _swatcher(self): for spec in specs: tid = spec.attributes['user']['uid'] fut = _flux_job.submit_async(fh, spec, waitable=True) - futs.append([fut, tid]) for fut, tid in futs: @@ -638,6 +640,7 @@ def _swatcher(self): finally: # trigger submit completion + self._log.debug('=== submit done') self._sevent.set() @@ -726,12 +729,13 @@ def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: if not self._handle: raise RuntimeError('flux instance not started') - self._log.debug('== submit %d specs', len(specs)) + self._log.debug('== submit %d specs start', len(specs)) tids = [spec.attributes['user']['uid'] for spec in specs] self._sevent.clear() self._squeue.put(specs) self._sevent.wait() # FIXME: timeout? + self._log.debug('== submit %d specs done', len(specs)) return tids From 1c27f2cb2889f4155d9512422ae8e6c3ee7db3e8 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 9 Apr 2025 12:01:36 +0200 Subject: [PATCH 32/44] typo --- src/radical/utils/flux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py index 06a05d865..2391ecd1c 100644 --- a/src/radical/utils/flux.py +++ b/src/radical/utils/flux.py @@ -229,7 +229,7 @@ def start(self, timeout: float = None) -> None: fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=\\$(hostname) ' fcmd += ' && flux resources list ' - fmcd += ' && sleep inf ' + fcmd += ' && sleep inf ' cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) if self._launcher: From c8343f265878670ec5f02799b7a071534655143e Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Thu, 10 Apr 2025 14:03:31 +0200 Subject: [PATCH 33/44] restructure flux code --- src/radical/utils/flux.py | 782 ----------------------- src/radical/utils/flux/__init__.py | 10 + src/radical/utils/flux/flux_helper_v0.py | 217 +++++++ src/radical/utils/flux/flux_helper_v1.py | 313 +++++++++ src/radical/utils/flux/flux_module.py | 183 ++++++ src/radical/utils/flux/flux_service.py | 153 +++++ 6 files changed, 876 insertions(+), 782 deletions(-) delete mode 100644 src/radical/utils/flux.py create mode 100644 src/radical/utils/flux/__init__.py create mode 100644 src/radical/utils/flux/flux_helper_v0.py create mode 100644 src/radical/utils/flux/flux_helper_v1.py create mode 100644 src/radical/utils/flux/flux_module.py create mode 100644 src/radical/utils/flux/flux_service.py diff --git a/src/radical/utils/flux.py b/src/radical/utils/flux.py deleted file mode 100644 index 2391ecd1c..000000000 --- a/src/radical/utils/flux.py +++ /dev/null @@ -1,782 +0,0 @@ - -import os -import sys -import math -import time -import shlex -import queue - -import threading as mt - -from rc.process import Process -from functools import partial -from collections import defaultdict -from typing import List, Dict, Any - -from .url import Url -from .misc import as_list -from .which import which -from .ids import generate_id, ID_SIMPLE -from .logger import Logger -from .modules import import_module -from .shell import sh_callout - - -# ------------------------------------------------------------------------------ -# -def import_flux(): - ''' - import the flux module, if available - - returns: flux : loaded python module (`None` if not available) - flux.job : loaded python module (`None` if not available) - exception: exception raised during import (`None` if no error) - version : what ru.FluxHelper version to use (`0` or `1`) - ''' - - flux = None - flux_job = None - flux_exc = None - flux_v = None - - try: - flux = import_module('flux') - flux_job = import_module('flux.job') - if 'JournalConsumer' in dir(flux_job): - flux_v = 1 - else: - flux_v = 0 - - except Exception as e: - flux_exc = e - - - # on failure, try to derive module path from flux executable - if flux is None or flux_job is None: - - to_pop = None - try: - cmd = 'flux python -c "import flux; print(flux.__file__)"' - out, err, ret = sh_callout(cmd) - - if not ret: - flux_path = os.path.dirname(out.strip()) - mod_path = os.path.dirname(flux_path) - sys.path.append(mod_path) - to_pop = mod_path - - flux = import_module('flux') - flux_job = import_module('flux.job') - if 'JournalConsumer' in dir(flux_job): - flux_v = 1 - else: - flux_v = 0 - - except Exception as e: - flux_exc = e - - if to_pop: - sys.path.remove(to_pop) - - return flux, flux_job, flux_exc, flux_v - - -_flux, _flux_job, _flux_exc, _flux_v = import_flux() - - -# -------------------------------------------------------------------------- -# -def spec_from_command(cmd: str) -> 'flux.job.JobspecV1': - - spec = _flux_job.JobspecV1.from_command(shlex.split(cmd)) - spec.attributes['user']['uid'] = generate_id(ID_SIMPLE) - - return spec - - -# -------------------------------------------------------------------------- -# -def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': - - version = 1 - user = {'uid' : td.get('uid', generate_id('ru_flux', ID_SIMPLE))} - system = {'duration': td.get('duration', 0.0)} - tasks = [{'command': [td['executable']] + td.get('arguments', []), - 'slot' : 'task', - 'count' : {'per_slot': 1}}] - - if 'environment' in td: system['environment'] = td['environment'] - if 'sandbox' in td: system['cwd'] = td['sandbox'] - if 'shell' in td: system['shell'] = td['shell'] - if 'stdin' in td: system['stdin'] = td['stdin'] - if 'stdout' in td: system['stdout'] = td['stdout'] - if 'stderr' in td: system['stderr'] = td['stderr'] - - attributes = {'system' : system, - 'user' : user} - resources = [{'count': td.get('ranks', 1), - 'type' : 'slot', - 'label': 'task', - 'with' : [{ - 'count': int(td.get('cores_per_rank', 1)), - 'type' : 'core'}]}] - # 'count': int(td.get('gpus_per_rank', 0)) or None, - # 'type' : 'gpu' - - gpr = td.get('gpus_per_rank', 0) - if gpr: - resources[0]['with'].append({'count': math.ceil(gpr), # flux needs int - 'type' : 'gpu'}) - - spec = _flux_job.JobspecV1(resources=resources, - attributes=attributes, - tasks=tasks, - version=version) - return spec - - -# ------------------------------------------------------------------------------ -# -class FluxService(object): - - # -------------------------------------------------------------------------- - # - def __init__(self, uid : str = None, - log : Logger = None, - launcher: str = None - ) -> None: - - self._uid = uid or generate_id('ru.flux') - self._log = log or Logger('radical.utils.flux') - self._launcher = launcher or '' - - self._fexe = which('flux') - self._uri = None - self._r_uri = None - self._host = None - self._proc = None - self._ready = mt.Event() - - if not _flux: - raise RuntimeError('flux module not found') from _flux_exc - - if not _flux_job: - raise RuntimeError('flux.job module not found') from _flux_exc - - if not self._fexe: - raise RuntimeError('flux executable not found') - - - # -------------------------------------------------------------------------- - # - @property - def uid(self) -> str: - return self._uid - - @property - def uri(self) -> str: - return self._uri - - - @property - def r_uri(self) -> str: - return self._r_uri - - - # -------------------------------------------------------------------------- - # - def _proc_line_cb(self, prefix: str, - proc : Process, - lines : List[str] - ) -> None: - - try: - for line in lines: - self._log.info('=== line: %s', line) - if line.startswith('FLUX_URI='): - parts = line.strip().split(' ', 1) - self._log.info('%s: found flux info: %s', self._uid, parts) - - self._uri = parts[0].split('=', 1)[1] - self._host = parts[1].split('=', 1)[1] - - url = Url(self._uri) - url.host = self._host - url.schema = 'ssh' - self._r_uri = str(url) - - self._log.info('%s: flux uri: %s', self._uid, self._uri) - self._log.info('%s: r uri: %s', self._uid, self._r_uri) - self._ready.set() - except: - self._log.exception('line processing failed') - - - # -------------------------------------------------------------------------- - # - def _proc_state_cb(self, proc: Process, state: str) -> None: - - self._log.info('flux instance state update: %s', state) - if state in Process.FINAL: - - self._log.info('flux instance stopped: %s', state) - self.stop() - - - # -------------------------------------------------------------------------- - # - def start(self, timeout: float = None) -> None: - - fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=\\$(hostname) ' - fcmd += ' && flux resources list ' - fcmd += ' && sleep inf ' - cmd = '%s start bash -c "%s"' % (self._fexe, fcmd) - - if self._launcher: - cmd = '%s %s' % (self._launcher, cmd) - - self._log.info('%s: start flux instance: %s', self._uid, cmd) - - p = Process(cmd) - p.register_cb(p.CB_OUT_LINE, partial(self._proc_line_cb, 'out')) - p.register_cb(p.CB_ERR_LINE, partial(self._proc_line_cb, 'err')) - p.register_cb(p.CB_STATE, self._proc_state_cb) - p.polldelay = 0.1 - p.start() - - self._proc = p - self._ptime = time.time() - - return self.ready(timeout=timeout) - - - # -------------------------------------------------------------------------- - # - def ready(self, timeout: float = None) -> None: - - if timeout is not None: - if timeout < 0: self._ready.wait() - else : self._ready.wait(timeout) - - return self._ready.is_set() - - - # -------------------------------------------------------------------------- - # - def stop(self) -> None: - - if not self._proc: - return - - self._proc.cancel() - self._proc.wait() - - self._uri = None - self._proc = None - - self._log.info('%s: found flux uri: %s', self._uid, self.uri) - - -# ------------------------------------------------------------------------------ -# -class FluxHelperV0(object): - - # -------------------------------------------------------------------------- - # - def __init__(self, uri : str, - log : Logger = None) -> None: - - self._t0 = time.time() - - self._uri = uri - self._log = log or Logger('radical.utils.flux') - self._uid = generate_id('ru.flux') - self._handle = _flux.Flux(self._uri) - self._api_lock = mt.Lock() - self._exe = _flux_job.executor.FluxExecutor( - handle_kwargs={'url': self._uri}) - - self._idlock = mt.Lock() # lock ID dicts - self._elock = mt.Lock() # lock event dict - self._task_ids = dict() # flux ID -> task ID - self._flux_ids = dict() # task ID -> flux ID - self._events = defaultdict(list) # flux ID -> event list - self._cbacks = list() # list of callbacks - - if not _flux: - raise RuntimeError('flux module not found') from _flux_exc - - if not _flux_job: - raise RuntimeError('flux.job module not found') from _flux_exc - - - # -------------------------------------------------------------------------- - # - def start(self, launcher: str = None) -> None: - - pass - - - # -------------------------------------------------------------------------- - # - def stop(self): - - with self._api_lock: - - # FIXME: shutdown flux instance - self._flux_service = None - self._uri = None - self._handle = None - - - # -------------------------------------------------------------------------- - # - @property - def uid(self) -> str: - return self._uid - - @property - def uri(self) -> str: - return self._uri - - - # -------------------------------------------------------------------------- - # - def register_cb(self, cb: callable) -> None: - - with self._api_lock, self._elock: - self._log.debug('register cb %s', cb) - self._cbacks.append(cb) - - - # -------------------------------------------------------------------------- - # - def unregister_cb(self, cb: callable) -> None: - - with self._api_lock, self._elock: - self._cbacks.remove(cb) - - - # -------------------------------------------------------------------------- - # - def _handle_events(self, fh : 'flux.Flux', - fid : 'flux.job.JobID', - event: 'flux.job.journal.JournalEvent' = None - ) -> None: - - self._log.debug('event for %s: %s', fid, event) - # print('event %s: %s' % (fid, event)) - - # if triggered by submit, check if we have anything to do - if not event: - if fid not in self._events: - # print('no event') - return - - # check if we can handle the event - otherwise store it - if not self._cbacks: - # print('no cbacks') - self._events[fid].append(event) - return - - # task is known, flush stored events - for ev in self._events[fid]: - # print('flush stored events') - for cb in self._cbacks: - try : cb(fid, ev) - except: self._log.exception('cb failed') - self._events[fid] = [] - - # process the current event - if event: - # print('process current event') - for cb in self._cbacks: - try : cb(fid, event) - except: self._log.exception('cb failed') - - - # -------------------------------------------------------------------------- - # - def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: - - with self._api_lock: - - if not self._handle: - raise RuntimeError('flux instance not started') - - self._log.debug('== submit %d specs', len(specs)) - - events = ['submit', 'depend', 'alloc', 'start', # 'cleanup', - 'finish', 'release', 'free', 'clean', 'priority', 'exception'] - - def event_cb(fid, fut, event): - self._handle_events(self._handle, fid, event) - - futures = list() - def id_cb(fut): - flux_id = fut.jobid() - idx = fut.ru_idx - for ev in events: - tmp_cb = partial(event_cb, flux_id) - fut.add_event_callback(ev, tmp_cb) - futures.append([flux_id, idx, fut]) - self._log.debug('got flux id: %s: %s', idx, flux_id) - - for idx, spec in enumerate(specs): - fut = self._exe.submit(spec, waitable=True) - fut.ru_idx = idx - self._log.debug('%s: submitted : %s', self._uid, idx) - fut.add_jobid_callback(id_cb) - - # wait until we saw all jobid callbacks (assume 10 tasks/sec) - timeout = len(specs) - timeout = max(100, timeout) - start = time.time() - self._log.debug('%s: wait %.2fsec for %d flux IDs', - self._uid, timeout, len(specs)) - while len(futures) < len(specs): - time.sleep(0.1) - self._log.debug('%s: wait %s / %s', self._uid, - len(futures), len(specs)) - if time.time() - start > timeout: - raise RuntimeError('%s: timeout on submission' % self._uid) - self._log.info('got %d flux IDs', len(futures)) - - # get flux_ids sorted by submission order (idx) - flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] - - self._log.debug('%s: submitted: %s', self._uid, flux_ids) - return flux_ids - - - # -------------------------------------------------------------------------- - # - def cancel(self, fids: [str|List[str]]) -> None: - - with self._api_lock: - - if not self._handle: - raise RuntimeError('flux instance not started') - - with self._idlock: - for fid in as_list(fids): - _flux_job.cancel_async(self._handle, fid, reason='user cancel') - - - # -------------------------------------------------------------------------- - # - def wait(self, fids: [str|List[str]]) -> None: - - with self._api_lock: - - if not self._handle: - raise RuntimeError('flux instance not started') - - for fid in fids: - self._log.debug('wait for %s', fid) - _flux_job.wait(self._handle, fid) - - -# ------------------------------------------------------------------------------ -# -class FluxHelperV1(object): - - # -------------------------------------------------------------------------- - # - def __init__(self, uri : str, - log : Logger = None) -> None: - - self._t0 = time.time() - - self._uri = uri - self._log = log or Logger('radical.utils.flux') - self._uid = generate_id('ru.flux') - self._handle = _flux.Flux(self._uri) - self._api_lock = mt.Lock() - - # journal watcher - self._jthread = None - - # event handle thread - self._ethread = None - self._equeue = queue.Queue() - - # submit thread - self._sthread = None - self._squeue = queue.Queue() - self._sevent = mt.Event() - - self._idlock = mt.Lock() # lock ID dicts - self._elock = mt.Lock() # lock event dict - self._task_ids = dict() # flux ID -> task ID - self._flux_ids = dict() # task ID -> flux ID - self._events = defaultdict(list) # flux ID -> event list - self._cbacks = list() # list of callbacks - - if not _flux: - raise RuntimeError('flux module not found') from _flux_exc - - if not _flux_job: - raise RuntimeError('flux.job module not found') from _flux_exc - - - # -------------------------------------------------------------------------- - # - def start(self, launcher: str = None) -> None: - - with self._api_lock: - - if self._jthread is not None: - return - - self._jthread = mt.Thread(target=self._jwatcher) - self._jthread.daemon = True - self._jthread.start() - - self._ethread = mt.Thread(target=self._ewatcher) - self._ethread.daemon = True - self._ethread.start() - - self._sthread = mt.Thread(target=self._swatcher) - self._sthread.daemon = True - self._sthread.start() - - - # -------------------------------------------------------------------------- - # - def stop(self): - - with self._api_lock: - - if self._handle is None: - self._jterm.set() - self._jthread.join() - - # FIXME: shutdown flux instance - self._flux_service = None - self._uri = None - self._handle = None - - - # -------------------------------------------------------------------------- - # - @property - def uid(self) -> str: - return self._uid - - @property - def uri(self) -> str: - return self._uri - - - # -------------------------------------------------------------------------- - # - def _jwatcher(self): - - # NOTE: *never* used self._handle in this thread, as it is not thread - # safe. Instead, use the private handle created here - fh = _flux.Flux(self._uri) - - # start watching the event journal - journal = _flux_job.JournalConsumer(fh) - journal.start() - - while True: - - try: - event = journal.poll(timeout=1.0) - if event: - # FIXME: How can that ever *not* be a journal event? - # But it has happened... - self._handle_events(fh, event.jobid, event) - - except TimeoutError: - pass - - - # -------------------------------------------------------------------------- - # - def _swatcher(self): - - self._log.debug('=== swatcher started') - - # if we get new specs, submit them, return IDs to iqueue, and also - # forward ID to ewatcher - fh = _flux.Flux(self._uri) - while True: - - try: - specs = self._squeue.get(block=True, timeout=1.0) - self._log.debug('=== got %d specs', len(specs)) - - except queue.Empty: - continue - - except: - self._log.exception("exception") - raise - - with self._idlock: - - try: - - futs = list() - for spec in specs: - tid = spec.attributes['user']['uid'] - fut = _flux_job.submit_async(fh, spec, waitable=True) - futs.append([fut, tid]) - - for fut, tid in futs: - fid = fut.get_id() - self._task_ids[fid] = tid - self._flux_ids[tid] = fid - - # trigger an event check - self._equeue.put(fid) - - except Exception: - self._log.exception("exception") - raise - - finally: - # trigger submit completion - self._log.debug('=== submit done') - self._sevent.set() - - - # -------------------------------------------------------------------------- - # - def _ewatcher(self): - - # if we get a new job ID, check if we have events for it - - fh = _flux.Flux(self._uri) - while True: - - try: - fid = self._equeue.get(timeout=1.0) - self._handle_events(fh, fid) - - except queue.Empty: - continue - - - # -------------------------------------------------------------------------- - # - def register_cb(self, cb: callable) -> None: - - with self._api_lock, self._elock: - self._log.debug('register cb %s', cb) - self._cbacks.append(cb) - - - # -------------------------------------------------------------------------- - # - def unregister_cb(self, cb: callable) -> None: - - with self._api_lock, self._elock: - self._cbacks.remove(cb) - - - # -------------------------------------------------------------------------- - # - def _handle_events(self, fh : 'flux.Flux', - fid : 'flux.job.JobID', - event: 'flux.job.journal.JournalEvent' = None - ) -> None: - - with self._elock: - - # self._log.debug_9('event %s: %s', fid, event) - - # if triggered by submit, check if we have anything to do - if not event: - if fid not in self._events: - return - - # check if we can handle the event - otherwise store it - if not self._cbacks: - self._events[fid].append(event) - return - - # check if application knows the task - otherwise store the event - if fid not in self._task_ids: - self._events[fid].append(event) - return - - tid = self._task_ids[fid] - - # task is known, flush stored events - for ev in self._events[fid]: - for cb in self._cbacks: - try : cb(tid, ev) - except: self._log.exception('cb failed') - self._events[fid] = [] - - # process the current event - if event: - for cb in self._cbacks: - try : cb(tid, event) - except: self._log.exception('cb failed') - - - # -------------------------------------------------------------------------- - # - def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: - - with self._api_lock: - - if not self._handle: - raise RuntimeError('flux instance not started') - - self._log.debug('== submit %d specs start', len(specs)) - tids = [spec.attributes['user']['uid'] for spec in specs] - - self._sevent.clear() - self._squeue.put(specs) - self._sevent.wait() # FIXME: timeout? - self._log.debug('== submit %d specs done', len(specs)) - - return tids - - - # -------------------------------------------------------------------------- - # - def cancel(self, tids: [str|List[str]]) -> None: - - with self._api_lock: - - if not self._handle: - raise RuntimeError('flux instance not started') - - with self._idlock: - for tid in as_list(tids): - fid = self._flux_ids[tid] - _flux_job.cancel_async(self._handle, fid, reason='user cancel') - - - # -------------------------------------------------------------------------- - # - def wait(self, tids: [str|List[str]]) -> None: - - with self._api_lock: - - if not self._handle: - raise RuntimeError('flux instance not started') - - tids = as_list(tids) - with self._idlock: - fids = [self._flux_ids[tid] for tid in tids] - - for fid in fids: - _flux_job.wait(self._handle, fid) - - -# ------------------------------------------------------------------------------ -# -if _flux_v == 1: FluxHelper = FluxHelperV1 -else : FluxHelper = FluxHelperV0 - - -# ------------------------------------------------------------------------------ - diff --git a/src/radical/utils/flux/__init__.py b/src/radical/utils/flux/__init__.py new file mode 100644 index 000000000..a7fbfbcce --- /dev/null +++ b/src/radical/utils/flux/__init__.py @@ -0,0 +1,10 @@ + +from .flux_service import FluxService +from .flux_helper_v0 import FluxHelperV0 as _FluxHelperV0 +from .flux_helper_v1 import FluxHelperV1 as _FluxHelperV1 +from .flux_module import FluxModule, spec_from_command, spec_from_dict + +_fm = FluxModule() +if _fm.version == 1: FluxHelper = _FluxHelperV1 +else : FluxHelper = _FluxHelperV0 + diff --git a/src/radical/utils/flux/flux_helper_v0.py b/src/radical/utils/flux/flux_helper_v0.py new file mode 100644 index 000000000..84b631814 --- /dev/null +++ b/src/radical/utils/flux/flux_helper_v0.py @@ -0,0 +1,217 @@ + +import time +import queue + +import threading as mt + +from functools import partial +from collections import defaultdict +from typing import List + +from ..misc import as_list +from ..ids import generate_id +from ..logger import Logger + +from .flux_module import FluxModule + + +# ------------------------------------------------------------------------------ +# +class FluxHelperV0(object): + + # -------------------------------------------------------------------------- + # + def __init__(self, uri : str, + log : Logger = None) -> None: + + print('=== v0 flux helper ===') + + self._uri = uri + self._log = log or Logger('radical.utils.flux') + self._uid = generate_id('ru.flux') + + self._fm = FluxModule() + self._handle = self._fm.core.Flux(self._uri) + self._api_lock = mt.Lock() + self._exe = self._fm.job.executor.FluxExecutor( + handle_kwargs={'url': self._uri}) + + self._idlock = mt.Lock() # lock ID dicts + self._elock = mt.Lock() # lock event dict + self._task_ids = dict() # flux ID -> task ID + self._flux_ids = dict() # task ID -> flux ID + self._events = defaultdict(list) # flux ID -> event list + self._cbacks = list() # list of callbacks + + self._fm.verify() + + + # -------------------------------------------------------------------------- + # + def start(self, launcher: str = None) -> None: + + pass + + + # -------------------------------------------------------------------------- + # + def stop(self): + + with self._api_lock: + + # FIXME: shutdown flux instance + self._flux_service = None + self._uri = None + self._handle = None + + + # -------------------------------------------------------------------------- + # + @property + def uid(self) -> str: + return self._uid + + @property + def uri(self) -> str: + return self._uri + + + # -------------------------------------------------------------------------- + # + def register_cb(self, cb: callable) -> None: + + with self._api_lock, self._elock: + self._log.debug('register cb %s', cb) + self._cbacks.append(cb) + + + # -------------------------------------------------------------------------- + # + def unregister_cb(self, cb: callable) -> None: + + with self._api_lock, self._elock: + self._cbacks.remove(cb) + + + # -------------------------------------------------------------------------- + # + def _handle_events(self, fh : 'flux.Flux', + fid : 'flux.job.JobID', + event: 'flux.job.journal.JournalEvent' = None + ) -> None: + + self._log.debug('event for %s: %s', fid, event) + # print('event %s: %s' % (fid, event)) + + # if triggered by submit, check if we have anything to do + if not event: + if fid not in self._events: + # print('no event') + return + + # check if we can handle the event - otherwise store it + if not self._cbacks: + # print('no cbacks') + self._events[fid].append(event) + return + + # task is known, flush stored events + for ev in self._events[fid]: + # print('flush stored events') + for cb in self._cbacks: + try : cb(fid, ev) + except: self._log.exception('cb failed') + self._events[fid] = [] + + # process the current event + if event: + # print('process current event') + for cb in self._cbacks: + try : cb(fid, event) + except: self._log.exception('cb failed') + + + # -------------------------------------------------------------------------- + # + def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + self._log.debug('== submit %d specs', len(specs)) + + events = ['submit', 'depend', 'alloc', 'start', # 'cleanup', + 'finish', 'release', 'free', 'clean', 'priority', 'exception'] + + def event_cb(fid, fut, event): + self._handle_events(self._handle, fid, event) + + futures = list() + def id_cb(fut): + flux_id = fut.jobid() + idx = fut.ru_idx + for ev in events: + tmp_cb = partial(event_cb, flux_id) + fut.add_event_callback(ev, tmp_cb) + futures.append([flux_id, idx, fut]) + self._log.debug('got flux id: %s: %s', idx, flux_id) + + for idx, spec in enumerate(specs): + fut = self._exe.submit(spec, waitable=True) + fut.ru_idx = idx + self._log.debug('%s: submitted : %s', self._uid, idx) + fut.add_jobid_callback(id_cb) + + # wait until we saw all jobid callbacks (assume 10 tasks/sec) + timeout = len(specs) + timeout = max(100, timeout) + start = time.time() + self._log.debug('%s: wait %.2fsec for %d flux IDs', + self._uid, timeout, len(specs)) + while len(futures) < len(specs): + time.sleep(0.1) + self._log.debug('%s: wait %s / %s', self._uid, + len(futures), len(specs)) + if time.time() - start > timeout: + raise RuntimeError('%s: timeout on submission' % self._uid) + self._log.info('got %d flux IDs', len(futures)) + + # get flux_ids sorted by submission order (idx) + flux_ids = [fut[0] for fut in sorted(futures, key=lambda x: x[1])] + + self._log.debug('%s: submitted: %s', self._uid, flux_ids) + return flux_ids + + + # -------------------------------------------------------------------------- + # + def cancel(self, fids: [str|List[str]]) -> None: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + with self._idlock: + for fid in as_list(fids): + self._fm.job.cancel_async(self._handle, fid, reason='user cancel') + + + # -------------------------------------------------------------------------- + # + def wait(self, fids: [str|List[str]]) -> None: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + for fid in fids: + self._log.debug('wait for %s', fid) + self._fm.job.wait(self._handle, fid) + + +# ------------------------------------------------------------------------------ + diff --git a/src/radical/utils/flux/flux_helper_v1.py b/src/radical/utils/flux/flux_helper_v1.py new file mode 100644 index 000000000..45357ee82 --- /dev/null +++ b/src/radical/utils/flux/flux_helper_v1.py @@ -0,0 +1,313 @@ + +import time +import queue + +import threading as mt + +from collections import defaultdict +from typing import List + +from ..misc import as_list +from ..ids import generate_id +from ..logger import Logger + +from .flux_module import FluxModule + + +# ------------------------------------------------------------------------------ +# +class FluxHelperV1(object): + + # -------------------------------------------------------------------------- + # + def __init__(self, uri : str, + log : Logger = None) -> None: + + print('=== v1 flux helper ===') + + self._uri = uri + self._log = log or Logger('radical.utils.flux') + self._uid = generate_id('ru.flux') + + self._fm = FluxModule() + self._handle = self._fm.core.Flux(self._uri) + self._api_lock = mt.Lock() + + # journal watcher + self._jthread = None + + # event handle thread + self._ethread = None + self._equeue = queue.Queue() + + # submit thread + self._sthread = None + self._squeue = queue.Queue() + self._sevent = mt.Event() + + self._idlock = mt.Lock() # lock ID dicts + self._elock = mt.Lock() # lock event dict + self._task_ids = dict() # flux ID -> task ID + self._flux_ids = dict() # task ID -> flux ID + self._events = defaultdict(list) # flux ID -> event list + self._cbacks = list() # list of callbacks + + self._fm.verify() + + + # -------------------------------------------------------------------------- + # + def start(self, launcher: str = None) -> None: + + with self._api_lock: + + if self._jthread is not None: + return + + self._jthread = mt.Thread(target=self._jwatcher) + self._jthread.daemon = True + self._jthread.start() + + self._ethread = mt.Thread(target=self._ewatcher) + self._ethread.daemon = True + self._ethread.start() + + self._sthread = mt.Thread(target=self._swatcher) + self._sthread.daemon = True + self._sthread.start() + + + # -------------------------------------------------------------------------- + # + def stop(self): + + with self._api_lock: + + if self._handle is None: + self._jterm.set() + self._jthread.join() + + # FIXME: shutdown flux instance + self._flux_service = None + self._uri = None + self._handle = None + + + # -------------------------------------------------------------------------- + # + @property + def uid(self) -> str: + return self._uid + + @property + def uri(self) -> str: + return self._uri + + + # -------------------------------------------------------------------------- + # + def _jwatcher(self): + + # NOTE: *never* used self._handle in this thread, as it is not thread + # safe. Instead, use the private handle created here + fh = self._fm.core.Flux(self._uri) + + # start watching the event journal + journal = self._fm.job.JournalConsumer(fh) + journal.start() + + while True: + + try: + event = journal.poll(timeout=1.0) + if event: + # FIXME: How can that ever *not* be a journal event? + # But it has happened... + self._handle_events(fh, event.jobid, event) + + except TimeoutError: + pass + + + # -------------------------------------------------------------------------- + # + def _swatcher(self): + + self._log.debug('=== swatcher started') + + # if we get new specs, submit them, return IDs to iqueue, and also + # forward ID to ewatcher + fh = self._fm.core.Flux(self._uri) + while True: + + try: + specs = self._squeue.get(block=True, timeout=1.0) + self._log.debug('=== got %d specs', len(specs)) + + except queue.Empty: + continue + + except: + self._log.exception("exception") + raise + + with self._idlock: + + try: + + futs = list() + for spec in specs: + tid = spec.attributes['user']['uid'] + fut = self._fm.job.submit_async(fh, spec, waitable=True) + futs.append([fut, tid]) + + for fut, tid in futs: + fid = fut.get_id() + self._task_ids[fid] = tid + self._flux_ids[tid] = fid + + # trigger an event check + self._equeue.put(fid) + + except Exception: + self._log.exception("exception") + raise + + finally: + # trigger submit completion + self._log.debug('=== submit done') + self._sevent.set() + + + # -------------------------------------------------------------------------- + # + def _ewatcher(self): + + # if we get a new job ID, check if we have events for it + + fh = self._fm.core.Flux(self._uri) + while True: + + try: + fid = self._equeue.get(timeout=1.0) + self._handle_events(fh, fid) + + except queue.Empty: + continue + + + # -------------------------------------------------------------------------- + # + def register_cb(self, cb: callable) -> None: + + with self._api_lock, self._elock: + self._log.debug('register cb %s', cb) + self._cbacks.append(cb) + + + # -------------------------------------------------------------------------- + # + def unregister_cb(self, cb: callable) -> None: + + with self._api_lock, self._elock: + self._cbacks.remove(cb) + + + # -------------------------------------------------------------------------- + # + def _handle_events(self, fh : 'flux.Flux', + fid : 'flux.job.JobID', + event: 'flux.job.journal.JournalEvent' = None + ) -> None: + + with self._elock: + + # self._log.debug_9('event %s: %s', fid, event) + + # if triggered by submit, check if we have anything to do + if not event: + if fid not in self._events: + return + + # check if we can handle the event - otherwise store it + if not self._cbacks: + self._events[fid].append(event) + return + + # check if application knows the task - otherwise store the event + if fid not in self._task_ids: + self._events[fid].append(event) + return + + tid = self._task_ids[fid] + + # task is known, flush stored events + for ev in self._events[fid]: + for cb in self._cbacks: + try : cb(tid, ev) + except: self._log.exception('cb failed') + self._events[fid] = [] + + # process the current event + if event: + for cb in self._cbacks: + try : cb(tid, event) + except: self._log.exception('cb failed') + + + # -------------------------------------------------------------------------- + # + def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + self._log.debug('== submit %d specs start', len(specs)) + tids = [spec.attributes['user']['uid'] for spec in specs] + + self._sevent.clear() + self._squeue.put(specs) + self._sevent.wait() # FIXME: timeout? + self._log.debug('== submit %d specs done', len(specs)) + + return tids + + + # -------------------------------------------------------------------------- + # + def cancel(self, tids: [str|List[str]]) -> None: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + with self._idlock: + for tid in as_list(tids): + fid = self._flux_ids[tid] + self._fm.job.cancel_async(self._handle, fid, reason='user cancel') + + + # -------------------------------------------------------------------------- + # + def wait(self, tids: [str|List[str]]) -> None: + + with self._api_lock: + + if not self._handle: + raise RuntimeError('flux instance not started') + + tids = as_list(tids) + with self._idlock: + fids = [self._flux_ids[tid] for tid in tids] + + for fid in fids: + self._fm.job.wait(self._handle, fid) + + # FIXME: remove tasks which have been waited for. + + +# ------------------------------------------------------------------------------ + diff --git a/src/radical/utils/flux/flux_module.py b/src/radical/utils/flux/flux_module.py new file mode 100644 index 000000000..8d86cd7af --- /dev/null +++ b/src/radical/utils/flux/flux_module.py @@ -0,0 +1,183 @@ + +import os +import sys +import math +import shlex + +from typing import Any + +from ..which import which +from ..ids import generate_id, ID_SIMPLE +from ..modules import import_module +from ..shell import sh_callout + + +# ------------------------------------------------------------------------------ +# +class FluxModule(object): + + _flux_core = None + _flux_job = None + _flux_exc = None + _flux_v = None + + + # -------------------------------------------------------------------------- + # + def __init__(self): + ''' + import the flux module, if available + ''' + + if self._flux_core or self._flux_job or self._flux_exc: + return + + flux = None + flux_job = None + flux_exc = None + flux_v = None + + try: + flux = import_module('flux') + flux_job = import_module('flux.job') + if 'JournalConsumer' in dir(flux_job): + flux_v = 1 + else: + flux_v = 0 + + except Exception as e: + flux_exc = e + + + # on failure, try to derive module path from flux executable + if flux is None or flux_job is None: + + to_pop = None + try: + cmd = 'flux python -c "import flux; print(flux.__file__)"' + out, err, ret = sh_callout(cmd) + + assert not ret, [cmd, err] + + flux_path = os.path.dirname(out.strip()) + mod_path = os.path.dirname(flux_path) + sys.path.append(mod_path) + to_pop = mod_path + + flux = import_module('flux') + flux_job = import_module('flux.job') + if 'JournalConsumer' in dir(flux_job): + flux_v = 1 + else: + flux_v = 0 + + except Exception as e: + flux_exc = e + + if to_pop: + sys.path.remove(to_pop) + + self._flux_core = flux + self._flux_job = flux_job + self._flux_exc = flux_exc + self._flux_v = flux_v + self._flux_exe = which('flux') + + + # -------------------------------------------------------------------------- + # + def verify(self) -> None: + ''' + verify that flux modules are available + ''' + + if self._flux_core is None: + raise RuntimeError('flux core module not found') from self._flux_exc + + if self._flux_job is None: + raise RuntimeError('flux.job module not found') from self._flux_exc + + if self._flux_exe is None: + raise RuntimeError('flux executable not found') from self._flux_exc + + + # -------------------------------------------------------------------------- + # + @property + def version(self) -> int: + return self._flux_v + + @property + def core(self) -> Any: + return self._flux_core + + @property + def job(self) -> Any: + return self._flux_job + + @property + def exc(self) -> Any: + return self._flux_exc + + @property + def exe(self) -> str: + return self._flux_exe + + +# ------------------------------------------------------------------------------ +# +def spec_from_command(cmd: str) -> 'flux.job.JobspecV1': + + fm = FluxModule() + + spec = fm.job.JobspecV1.from_command(shlex.split(cmd)) + spec.attributes['user']['uid'] = generate_id(ID_SIMPLE) + + return spec + + +# ------------------------------------------------------------------------------ +# +def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': + + fm = FluxModule() + + version = 1 + user = {'uid' : td.get('uid', generate_id('ru_flux', ID_SIMPLE))} + system = {'duration': td.get('duration', 0.0)} + tasks = [{'command': [td['executable']] + td.get('arguments', []), + 'slot' : 'task', + 'count' : {'per_slot': 1}}] + + if 'environment' in td: system['environment'] = td['environment'] + if 'sandbox' in td: system['cwd'] = td['sandbox'] + if 'shell' in td: system['shell'] = td['shell'] + if 'stdin' in td: system['stdin'] = td['stdin'] + if 'stdout' in td: system['stdout'] = td['stdout'] + if 'stderr' in td: system['stderr'] = td['stderr'] + + attributes = {'system' : system, + 'user' : user} + resources = [{'count': td.get('ranks', 1), + 'type' : 'slot', + 'label': 'task', + 'with' : [{ + 'count': int(td.get('cores_per_rank', 1)), + 'type' : 'core'}]}] + # 'count': int(td.get('gpus_per_rank', 0)) or None, + # 'type' : 'gpu' + + gpr = td.get('gpus_per_rank', 0) + if gpr: + resources[0]['with'].append({'count': math.ceil(gpr), # flux needs int + 'type' : 'gpu'}) + + spec = fm.job.JobspecV1(resources=resources, + attributes=attributes, + tasks=tasks, + version=version) + return spec + + +# ------------------------------------------------------------------------------ + diff --git a/src/radical/utils/flux/flux_service.py b/src/radical/utils/flux/flux_service.py new file mode 100644 index 000000000..0980d4833 --- /dev/null +++ b/src/radical/utils/flux/flux_service.py @@ -0,0 +1,153 @@ + +import time + +import threading as mt + +from rc.process import Process +from functools import partial +from typing import List + +from ..url import Url +from ..ids import generate_id +from ..logger import Logger + +from .flux_module import FluxModule + + +# ------------------------------------------------------------------------------ +# +class FluxService(object): + + # -------------------------------------------------------------------------- + # + def __init__(self, uid : str = None, + log : Logger = None, + launcher: str = None + ) -> None: + + self._uid = uid or generate_id('ru.flux') + self._log = log or Logger('radical.utils.flux') + self._launcher = launcher or '' + + self._fm = FluxModule() + self._uri = None + self._r_uri = None + self._host = None + self._proc = None + self._ready = mt.Event() + + self._fm.verify() + + + # -------------------------------------------------------------------------- + # + @property + def uid(self) -> str: + return self._uid + + @property + def uri(self) -> str: + return self._uri + + + @property + def r_uri(self) -> str: + return self._r_uri + + + # -------------------------------------------------------------------------- + # + def _proc_line_cb(self, prefix: str, + proc : Process, + lines : List[str] + ) -> None: + + try: + for line in lines: + self._log.info('=== line: %s', line) + if line.startswith('FLUX_URI='): + parts = line.strip().split(' ', 1) + self._log.info('%s: found flux info: %s', self._uid, parts) + + self._uri = parts[0].split('=', 1)[1] + self._host = parts[1].split('=', 1)[1] + + url = Url(self._uri) + url.host = self._host + url.schema = 'ssh' + self._r_uri = str(url) + + self._log.info('%s: flux uri: %s', self._uid, self._uri) + self._log.info('%s: r uri: %s', self._uid, self._r_uri) + self._ready.set() + except: + self._log.exception('line processing failed') + + + # -------------------------------------------------------------------------- + # + def _proc_state_cb(self, proc: Process, state: str) -> None: + + self._log.info('flux instance state update: %s', state) + if state in Process.FINAL: + + self._log.info('flux instance stopped: %s', state) + self.stop() + + + # -------------------------------------------------------------------------- + # + def start(self, timeout: float = None) -> None: + + fcmd = 'echo FLUX_URI=\\$FLUX_URI FLUX_HOST=\\$(hostname) ' + fcmd += ' && flux resource list ' + fcmd += ' && sleep inf ' + cmd = ' %s start bash -c "%s"' % (self._fm.exe, fcmd) + + if self._launcher: + cmd = '%s %s' % (self._launcher, cmd) + + self._log.info('%s: start flux instance: %s', self._uid, cmd) + + p = Process(cmd) + p.register_cb(p.CB_OUT_LINE, partial(self._proc_line_cb, 'out')) + p.register_cb(p.CB_ERR_LINE, partial(self._proc_line_cb, 'err')) + p.register_cb(p.CB_STATE, self._proc_state_cb) + p.polldelay = 0.1 + p.start() + + self._proc = p + self._ptime = time.time() + + return self.ready(timeout=timeout) + + + # -------------------------------------------------------------------------- + # + def ready(self, timeout: float = None) -> None: + + if timeout is not None: + if timeout < 0: self._ready.wait() + else : self._ready.wait(timeout) + + return self._ready.is_set() + + + # -------------------------------------------------------------------------- + # + def stop(self) -> None: + + if not self._proc: + return + + self._proc.cancel() + self._proc.wait() + + self._uri = None + self._proc = None + + self._log.info('%s: found flux uri: %s', self._uid, self.uri) + + +# ------------------------------------------------------------------------------ + From c3834a30a86dd7f799defca11e9994a23085b8d0 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 11 Jun 2025 09:27:46 +0200 Subject: [PATCH 34/44] type hint fixes --- src/radical/utils/flux/flux_helper_v0.py | 6 +++--- src/radical/utils/flux/flux_helper_v1.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/radical/utils/flux/flux_helper_v0.py b/src/radical/utils/flux/flux_helper_v0.py index 84b631814..9402b7287 100644 --- a/src/radical/utils/flux/flux_helper_v0.py +++ b/src/radical/utils/flux/flux_helper_v0.py @@ -6,7 +6,7 @@ from functools import partial from collections import defaultdict -from typing import List +from typing import List, Union from ..misc import as_list from ..ids import generate_id @@ -187,7 +187,7 @@ def id_cb(fut): # -------------------------------------------------------------------------- # - def cancel(self, fids: [str|List[str]]) -> None: + def cancel(self, fids: [Union[str, List[str]]]) -> None: with self._api_lock: @@ -201,7 +201,7 @@ def cancel(self, fids: [str|List[str]]) -> None: # -------------------------------------------------------------------------- # - def wait(self, fids: [str|List[str]]) -> None: + def wait(self, fids: [Union[str, List[str]]]) -> None: with self._api_lock: diff --git a/src/radical/utils/flux/flux_helper_v1.py b/src/radical/utils/flux/flux_helper_v1.py index 45357ee82..31b54efee 100644 --- a/src/radical/utils/flux/flux_helper_v1.py +++ b/src/radical/utils/flux/flux_helper_v1.py @@ -5,7 +5,7 @@ import threading as mt from collections import defaultdict -from typing import List +from typing import List, Union from ..misc import as_list from ..ids import generate_id @@ -277,7 +277,7 @@ def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: # -------------------------------------------------------------------------- # - def cancel(self, tids: [str|List[str]]) -> None: + def cancel(self, tids: [Union[str, List[str]]]) -> None: with self._api_lock: @@ -292,7 +292,7 @@ def cancel(self, tids: [str|List[str]]) -> None: # -------------------------------------------------------------------------- # - def wait(self, tids: [str|List[str]]) -> None: + def wait(self, tids: [Union[str, List[str]]]) -> None: with self._api_lock: From 1f7224441fb72f756d92f143b2ac7ec83f902622 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 7 Jul 2025 11:19:37 +0200 Subject: [PATCH 35/44] cleanup --- src/radical/utils/flux/flux_helper_v0.py | 2 +- src/radical/utils/flux/flux_helper_v1.py | 8 ++++---- src/radical/utils/flux/flux_service.py | 3 ++- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/radical/utils/flux/flux_helper_v0.py b/src/radical/utils/flux/flux_helper_v0.py index 9402b7287..80b43c2e6 100644 --- a/src/radical/utils/flux/flux_helper_v0.py +++ b/src/radical/utils/flux/flux_helper_v0.py @@ -24,7 +24,7 @@ class FluxHelperV0(object): def __init__(self, uri : str, log : Logger = None) -> None: - print('=== v0 flux helper ===') + # print('=== v0 flux helper ===') self._uri = uri self._log = log or Logger('radical.utils.flux') diff --git a/src/radical/utils/flux/flux_helper_v1.py b/src/radical/utils/flux/flux_helper_v1.py index 31b54efee..5e74f8fc5 100644 --- a/src/radical/utils/flux/flux_helper_v1.py +++ b/src/radical/utils/flux/flux_helper_v1.py @@ -23,7 +23,7 @@ class FluxHelperV1(object): def __init__(self, uri : str, log : Logger = None) -> None: - print('=== v1 flux helper ===') + # print('=== v1 flux helper ===') self._uri = uri self._log = log or Logger('radical.utils.flux') @@ -133,7 +133,7 @@ def _jwatcher(self): # def _swatcher(self): - self._log.debug('=== swatcher started') + self._log.debug('swatcher started') # if we get new specs, submit them, return IDs to iqueue, and also # forward ID to ewatcher @@ -142,7 +142,7 @@ def _swatcher(self): try: specs = self._squeue.get(block=True, timeout=1.0) - self._log.debug('=== got %d specs', len(specs)) + self._log.debug('got %d specs', len(specs)) except queue.Empty: continue @@ -175,7 +175,7 @@ def _swatcher(self): finally: # trigger submit completion - self._log.debug('=== submit done') + self._log.debug('submit done') self._sevent.set() diff --git a/src/radical/utils/flux/flux_service.py b/src/radical/utils/flux/flux_service.py index 0980d4833..1d288809e 100644 --- a/src/radical/utils/flux/flux_service.py +++ b/src/radical/utils/flux/flux_service.py @@ -64,7 +64,8 @@ def _proc_line_cb(self, prefix: str, try: for line in lines: - self._log.info('=== line: %s', line) + self._log.info('%s: flux io : %s', self._uid, line) + if line.startswith('FLUX_URI='): parts = line.strip().split(' ', 1) self._log.info('%s: found flux info: %s', self._uid, parts) From 9de2bb3d5af7d81d3b9cb1c4bb0a235188f88460 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 30 Jul 2025 15:06:16 +0200 Subject: [PATCH 36/44] fix timeout --- src/radical/utils/flux/flux_module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/radical/utils/flux/flux_module.py b/src/radical/utils/flux/flux_module.py index 8d86cd7af..7b29137fc 100644 --- a/src/radical/utils/flux/flux_module.py +++ b/src/radical/utils/flux/flux_module.py @@ -144,7 +144,7 @@ def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': version = 1 user = {'uid' : td.get('uid', generate_id('ru_flux', ID_SIMPLE))} - system = {'duration': td.get('duration', 0.0)} + system = {'duration': td.get('timeout', 0.0)} tasks = [{'command': [td['executable']] + td.get('arguments', []), 'slot' : 'task', 'count' : {'per_slot': 1}}] From bf6e71230f1aac8192884183f49d5108771613e8 Mon Sep 17 00:00:00 2001 From: Mikhail Titov Date: Wed, 13 Aug 2025 14:12:36 -0400 Subject: [PATCH 37/44] updated JobSpec creation --- src/radical/utils/flux/flux_module.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/radical/utils/flux/flux_module.py b/src/radical/utils/flux/flux_module.py index 7b29137fc..e2a7e40f8 100644 --- a/src/radical/utils/flux/flux_module.py +++ b/src/radical/utils/flux/flux_module.py @@ -152,9 +152,6 @@ def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': if 'environment' in td: system['environment'] = td['environment'] if 'sandbox' in td: system['cwd'] = td['sandbox'] if 'shell' in td: system['shell'] = td['shell'] - if 'stdin' in td: system['stdin'] = td['stdin'] - if 'stdout' in td: system['stdout'] = td['stdout'] - if 'stderr' in td: system['stderr'] = td['stderr'] attributes = {'system' : system, 'user' : user} @@ -176,6 +173,11 @@ def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': attributes=attributes, tasks=tasks, version=version) + + if td.get('stdin') : spec.stdin = td['stdin'] + if td.get('stdout'): spec.stdout = td['stdout'] + if td.get('stderr'): spec.stderr = td['stderr'] + return spec From 0367672fc56a504a0bf91496727aba3b6522f048 Mon Sep 17 00:00:00 2001 From: Mikhail Titov Date: Sun, 17 Aug 2025 00:47:14 -0400 Subject: [PATCH 38/44] allow env variable being referred as part of the value --- src/radical/utils/env.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/radical/utils/env.py b/src/radical/utils/env.py index 0cfdcfc5a..2b7063cd0 100644 --- a/src/radical/utils/env.py +++ b/src/radical/utils/env.py @@ -60,7 +60,8 @@ def env_read(fname: str) -> Dict[str, str]: # ------------------------------------------------------------------------------ # -def env_write(script_path, env, unset=None, blacklist=None, pre_exec=None): +def env_write(script_path, env, unset=None, blacklist=None, pre_exec=None, + extend=False): data = '\n' if unset: @@ -104,7 +105,7 @@ def env_write(script_path, env, unset=None, blacklist=None, pre_exec=None): continue if not re_snake_case.match(k): continue - data += "export %s=%s\n" % (k, _quote(env[k])) + data += "export %s=%s\n" % (k, _quote(env[k], extend=extend)) data += '\n' if funcs: @@ -194,14 +195,15 @@ def env_read_lines(lines: List[str]) -> Dict[str, str]: # ------------------------------------------------------------------------------ # -def _quote(data: str) -> str: +def _quote(data: str, extend: bool = False) -> str: if "'" in data or '$' in data or '`' in data: # cannot use single quote, so use double quote and escale all other # double quotes in the data # NOTE: we only support these three types of shell directives - data = data.replace('"', '\\"') \ - .replace('$', '\\$') + data = data.replace('"', '\\"') + if not extend: + data = data.replace('$', '\\$') data = '"' + data + '"' else: @@ -407,7 +409,7 @@ def env_prep(environment : Optional[Dict[str,str]] = None, _, tmp_name = tempfile.mkstemp(prefix=prefix, dir=tgt) env_write(tmp_name, env=environment, unset=unset, blacklist=blacklist, - pre_exec=pre_exec_cached) + pre_exec=pre_exec_cached, extend=False) cmd = '/bin/bash -c ". %s && /usr/bin/env"' % tmp_name out, err, ret = sh_callout(cmd) @@ -430,7 +432,7 @@ def env_prep(environment : Optional[Dict[str,str]] = None, # FIXME: files could also be cached and re-used (copied or linked) if script_path: env_write(script_path, env=env, unset=unset, blacklist=blacklist, - pre_exec=pre_exec) + pre_exec=pre_exec, extend=True) return env From c507ce29c1393b63d9613511ccd4bedf6505d71b Mon Sep 17 00:00:00 2001 From: Mikhail Titov Date: Tue, 9 Sep 2025 13:34:13 -0400 Subject: [PATCH 39/44] added control over exit procedures for Flux jobs/tasks (`exit-on-error`, `exit-timeout`) --- src/radical/utils/flux/flux_module.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/radical/utils/flux/flux_module.py b/src/radical/utils/flux/flux_module.py index e2a7e40f8..691e4ae63 100644 --- a/src/radical/utils/flux/flux_module.py +++ b/src/radical/utils/flux/flux_module.py @@ -155,7 +155,9 @@ def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': attributes = {'system' : system, 'user' : user} - resources = [{'count': td.get('ranks', 1), + + n_ranks = td.get('ranks', 1) + resources = [{'count': n_ranks, 'type' : 'slot', 'label': 'task', 'with' : [{ @@ -174,6 +176,13 @@ def spec_from_dict(td: dict) -> 'flux.job.JobspecV1': tasks=tasks, version=version) + if n_ranks > 1: + if td.get('use_mpi', True): + # ensure that all ranks exit if one rank fails + spec.setattr_shell_option('exit-on-error', 1) # defaults to 0 + else: + spec.setattr_shell_option('exit-timeout', 'none') # defaults to 30 + if td.get('stdin') : spec.stdin = td['stdin'] if td.get('stdout'): spec.stdout = td['stdout'] if td.get('stderr'): spec.stderr = td['stderr'] From 8f1b6ab6ed73379c2f18430df14fd586e8b7f3dc Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 24 Sep 2025 11:55:01 +0200 Subject: [PATCH 40/44] fix tests --- src/radical/utils/flux/flux_module.py | 2 + src/radical/utils/flux/flux_service.py | 3 +- tests/integration_tests/test_flux.py | 94 ++++---------------------- tests/unittests/test_env.py | 3 +- tests/unittests/test_typeddict.py | 55 ++++++++++----- 5 files changed, 57 insertions(+), 100 deletions(-) mode change 100644 => 100755 tests/unittests/test_typeddict.py diff --git a/src/radical/utils/flux/flux_module.py b/src/radical/utils/flux/flux_module.py index 691e4ae63..ca263e0e4 100644 --- a/src/radical/utils/flux/flux_module.py +++ b/src/radical/utils/flux/flux_module.py @@ -131,6 +131,8 @@ def spec_from_command(cmd: str) -> 'flux.job.JobspecV1': fm = FluxModule() spec = fm.job.JobspecV1.from_command(shlex.split(cmd)) + if not 'user' in spec.attributes: + spec.attributes['user'] = dict() spec.attributes['user']['uid'] = generate_id(ID_SIMPLE) return spec diff --git a/src/radical/utils/flux/flux_service.py b/src/radical/utils/flux/flux_service.py index 1d288809e..594388598 100644 --- a/src/radical/utils/flux/flux_service.py +++ b/src/radical/utils/flux/flux_service.py @@ -26,7 +26,7 @@ def __init__(self, uid : str = None, ) -> None: self._uid = uid or generate_id('ru.flux') - self._log = log or Logger('radical.utils.flux') + self._log = log or Logger('radical.utils.flux', level='DEBUG') self._launcher = launcher or '' self._fm = FluxModule() @@ -139,6 +139,7 @@ def ready(self, timeout: float = None) -> None: def stop(self) -> None: if not self._proc: + self._uri = None return self._proc.cancel() diff --git a/tests/integration_tests/test_flux.py b/tests/integration_tests/test_flux.py index 9beeee4e5..d997b0133 100755 --- a/tests/integration_tests/test_flux.py +++ b/tests/integration_tests/test_flux.py @@ -14,37 +14,12 @@ yaml = pytest.importorskip('yaml') flux = pytest.importorskip('flux') events = dict() -spec = { - "tasks": [{ - "slot": "task", - "count": { - "per_slot": 1 - }, - "command": [ - "/bin/date" - ] - }], - "attributes": { - "system": { - "duration": 10000 - } - }, - "version": 1, - "resources": [{ - "count": 1, - "type" : "slot", - "label": "task", - "with": [{ - "count": 1, - "type": "core" - }] - }] - } +spec = ru.flux.spec_from_command(cmd='/bin/date') # ------------------------------------------------------------------------------ # -def test_flux_startup(): +def test_flux(): global events @@ -59,60 +34,16 @@ def cb1(job_id, state): events[job_id].append(state) - fh = ru.FluxHelper() - fh.start_flux() + fs = ru.FluxService() + fs.start(timeout=10) - assert fh.uri - assert 'FLUX_URI' in fh.env + assert fs.uri - specs = [spec] * njobs - ids = fh.submit_jobs(specs, cb=cb1) - assert len(ids) == njobs, len(ids) - - time.sleep(5) - - assert len(events) == njobs, len(events) - for jid in events: - # we expect at least 4 events per job: - # 'submit', 'start', 'finish', 'clean', - assert len(events[jid]) >= 4, [jid, events[jid]] - - fh.reset() - assert fh.uri is None - - -# ------------------------------------------------------------------------------ -# -def test_flux_pickup(): - - global events - - njobs = 10 - events = dict() - outer_fh = None - - if 'FLUX_URI' not in os.environ: - outer_fh = ru.FluxHelper() - outer_fh.start_flux() - - for k,v in outer_fh.env.items(): - os.environ[k] = v - - def cb1(job_id, state): - - if job_id not in events: - events[job_id] = [state] - else: - events[job_id].append(state) - - fh = ru.FluxHelper() - fh.start_flux() - - assert fh.uri - assert 'FLUX_URI' in fh.env + fh = ru.FluxHelper(uri=fs.uri) + fh.register_cb(cb1) specs = [spec] * njobs - ids = fh.submit_jobs(specs, cb=cb1) + ids = fh.submit(specs) assert len(ids) == njobs, len(ids) time.sleep(5) @@ -123,19 +54,18 @@ def cb1(job_id, state): # 'submit', 'start', 'finish', 'clean', assert len(events[jid]) >= 4, [jid, events[jid]] - fh.reset() + fh.stop() assert fh.uri is None - if outer_fh: - outer_fh.reset() + fs.stop() + assert fs.uri is None # ------------------------------------------------------------------------------ # if __name__ == '__main__': - test_flux_startup() - test_flux_pickup() + test_flux() # ------------------------------------------------------------------------------ diff --git a/tests/unittests/test_env.py b/tests/unittests/test_env.py index e0fe4325c..88c135bf9 100755 --- a/tests/unittests/test_env.py +++ b/tests/unittests/test_env.py @@ -126,7 +126,8 @@ def test_env_proc(): env_proc = ru.EnvProcess(env=env) with env_proc: - env_proc.put(ru.sh_callout('echo -n $%s' % key, shell=True)) + if env_proc: + env_proc.put(ru.sh_callout('echo -n $%s' % key, shell=True)) out = str(env_proc.get()) assert isinstance(out, str) diff --git a/tests/unittests/test_typeddict.py b/tests/unittests/test_typeddict.py old mode 100644 new mode 100755 index d2c84df25..38d2e319f --- a/tests/unittests/test_typeddict.py +++ b/tests/unittests/test_typeddict.py @@ -1,3 +1,5 @@ +#!/usr/bin/env python3 + # pylint: disable=protected-access __author__ = 'RADICAL-Cybertools Team' @@ -654,22 +656,43 @@ class TDSchemaNotNone(TypedDict): self.assertIsInstance(td.any_data, TypedDict) self.assertIsNot(td.any_data, input_data['any_data']) - # -------------------------------------------------------------------------- - # - def test_pickle(self): - - import pickle + # # -------------------------------------------------------------------------- + # # + # def test_pickle(self): + # + # import pickle + # + # td = TDSimple({ + # 'attr_str' : 'foo', + # 'attr_dict': {'bar': 'buz'}}) + # + # ser = pickle.dumps(td) + # td2 = pickle.loads(ser) + # td2.verify() + # + # self.assertEqual(td2.attr_str, 'foo') + # self.assertEqual(td2.attr_int, 1) + # self.assertEqual(td2.attr_dict['bar'], 'buz') - td = TDSimple({ - 'attr_str' : 'foo', - 'attr_dict': {'bar': 'buz'}}) - - ser = pickle.dumps(td) - td2 = pickle.loads(ser) - td2.verify() +# ------------------------------------------------------------------------------ +# +if __name__ == '__main__': + + tc = TypedDictTestCase() + tc.setUpClass() + + tc.test_init() + tc.test_hash() + tc.test_self_default() + tc.test_verify() + tc.test_verify_setter() + tc.test_base_methods() + tc.test_pop() + tc.test_popitem() + tc.test_query() + tc.test_metaclass() + tc.test_tderrors() + tc.test_none() + # tc.test_pickle() - self.assertEqual(td2.attr_str, 'foo') - self.assertEqual(td2.attr_int, 1) - self.assertEqual(td2.attr_dict['bar'], 'buz') -# ------------------------------------------------------------------------------ From 8118c167d408ec1145b766b47e03f4a618fdb4cb Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Wed, 24 Sep 2025 12:23:16 +0200 Subject: [PATCH 41/44] fix tests --- tests/unittests/test_heartbeat.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/unittests/test_heartbeat.py b/tests/unittests/test_heartbeat.py index aa2f7af3f..9c11399de 100755 --- a/tests/unittests/test_heartbeat.py +++ b/tests/unittests/test_heartbeat.py @@ -153,6 +153,13 @@ def is_alive(pid): else: return True + def _join(proc, timeout=0.1): + proc.join(timeout=timeout) + try: + os.waitpid(test_proc.pid, os.WNOHANG) + except: + pass + # watcher process def _watcher(action): @@ -208,7 +215,7 @@ def _watcher(action): # after 1.2 seconds, the watcher should have exited time.sleep(1.2) - test_proc.join(timeout=0.0) + _join(test_proc, 0.0) assert not is_alive(pids[0]) assert not is_alive(pids[1]) assert not is_alive(pids[2]) @@ -222,7 +229,7 @@ def _watcher(action): # after 0.4 seconds, only second sleep should still be alive time.sleep(0.4) - test_proc.join(timeout=0.1) + _join(test_proc, 0.1) assert not is_alive(pids[0]) assert not is_alive(pids[1]) assert is_alive(pids[2]) @@ -243,14 +250,14 @@ def _watcher(action): # after 0.4 seconds, only second sleep should still be alive time.sleep(0.4) - test_proc.join(timeout=0.1) + _join(test_proc, 0.1) assert is_alive(pids[0]) assert not is_alive(pids[1]) assert not is_alive(pids[2]) # after 0.5 seconds, none of the processes should be alive time.sleep(0.5) - test_proc.join(timeout=0.1) + _join(test_proc, 0.1) assert not is_alive(pids[0]) assert not is_alive(pids[1]) assert not is_alive(pids[2]) @@ -265,14 +272,14 @@ def _watcher(action): # after 0.4 seconds, only second sleep should still be alive time.sleep(0.4) - test_proc.join(timeout=0.1) + _join(test_proc, 0.1) assert is_alive(pids[0]) assert not is_alive(pids[1]) assert not is_alive(pids[2]) # after 0.5 seconds, none of the processes should be alive time.sleep(0.5) - test_proc.join(timeout=0.1) + _join(test_proc, 0.1) assert not is_alive(pids[0]) assert not is_alive(pids[1]) assert not is_alive(pids[2]) From 72c36aaee7afb268c84815d0a88acfa6afc6d066 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 29 Sep 2025 11:53:59 +0200 Subject: [PATCH 42/44] Update src/radical/utils/flux/flux_helper_v0.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- src/radical/utils/flux/flux_helper_v0.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/radical/utils/flux/flux_helper_v0.py b/src/radical/utils/flux/flux_helper_v0.py index 80b43c2e6..5d9510bdb 100644 --- a/src/radical/utils/flux/flux_helper_v0.py +++ b/src/radical/utils/flux/flux_helper_v0.py @@ -60,7 +60,6 @@ def stop(self): with self._api_lock: # FIXME: shutdown flux instance - self._flux_service = None self._uri = None self._handle = None From cb30eae12b001b5b205e405b13fa8098cb70b57c Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 29 Sep 2025 11:56:00 +0200 Subject: [PATCH 43/44] respond to comments --- src/radical/utils/flux/flux_helper_v1.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/radical/utils/flux/flux_helper_v1.py b/src/radical/utils/flux/flux_helper_v1.py index 5e74f8fc5..674172ff8 100644 --- a/src/radical/utils/flux/flux_helper_v1.py +++ b/src/radical/utils/flux/flux_helper_v1.py @@ -35,6 +35,7 @@ def __init__(self, uri : str, # journal watcher self._jthread = None + self._jterm = mt.Event() # event handle thread self._ethread = None @@ -116,7 +117,7 @@ def _jwatcher(self): journal = self._fm.job.JournalConsumer(fh) journal.start() - while True: + while not self._jterm.is_set(): try: event = journal.poll(timeout=1.0) From 19d86b3dc48315d58ad571a7e71fcf09192d40b5 Mon Sep 17 00:00:00 2001 From: Andre Merzky Date: Mon, 29 Sep 2025 11:58:26 +0200 Subject: [PATCH 44/44] respond to comments --- src/radical/utils/flux/flux_helper_v1.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/radical/utils/flux/flux_helper_v1.py b/src/radical/utils/flux/flux_helper_v1.py index 674172ff8..fef643fe5 100644 --- a/src/radical/utils/flux/flux_helper_v1.py +++ b/src/radical/utils/flux/flux_helper_v1.py @@ -270,7 +270,11 @@ def submit(self, specs: List['flux.job.JobspecV1']) -> List[str]: self._sevent.clear() self._squeue.put(specs) - self._sevent.wait() # FIXME: timeout? + self._sevent.wait(timeout=60.0) + + if not self._sevent.is_set(): + raise RuntimeError('flux submit timeout') + self._log.debug('== submit %d specs done', len(specs)) return tids