Coverage for /home/agp/Documents/me/code/gutools/gutools/session.py : 0%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import asyncio
2import os
3import re
4import sys
5import signal
6import random
7import yaml
8import ujson as json
9import glob
10import time
11import shlex
12import subprocess
13import psutil
14# import warnings
15# warnings.simplefilter('always', DeprecationWarning)
17from datetime import datetime
18from logging import debug, info, warn, error
19from functools import partial
21# 3rd
22from daemon import daemon, pidfile
23from codenamize import codenamize
25# mine
26from aiopb.aiopb import Hub
27from gutools.tools import build_uri, parse_uri, snake_case, \
28 update_context, identity, yaml_encode, yaml_decode, wlikelyhood, \
29 fileiter, copyfile, dset, load_config, save_config, expandpath, \
30 _call, async_call, serializable_container, update_container, \
31 walk, rebuild
33from gutools.system import send_signals, SmartPIDLockFile, get_fingerprint, compare_fingerprint, find_process, test_process, test_alive
34from gutools.uobjects import UObject, kget, vget
35from gutools.persistence import DB, DBWorkspace, FSLayout
36from gutools.unet import TL
37from gutools.bash_completer import hasheable_cmdline
38from gutools.uapplication import UApplication # TODO: needed?
39from gutools.uagents import UAgent
40from gutools.colors import *
42async def dummy_main():
43 foo = 1
44 for i in range(20):
45 print(f"xxx {i}")
46 await asyncio.sleep(6)
47 foo = 1
50class UProcess(UObject):
51 __args__ = ('name', 'alias', 'uri', 'state',
52 'pid', 'restart', 'shutdown', 'date')
53 __kwargs__ = {}
54 __slots__ = __args__ + tuple(set(__kwargs__).difference(__args__))
56class SessionWorkspace(DBWorkspace):
57 scheme = """
58 CREATE TABLE IF NOT EXISTS {}
59 (
60 {},
61 date DATETIME DEFAULT (datetime(CURRENT_TIMESTAMP, 'localtime')),
62 PRIMARY KEY (name),
63 UNIQUE (alias)
64 );
65 """.format(
66 snake_case(UProcess.__name__),
67 ','.join(UProcess.__slots__[:-1]),
68 )
71ext_decoder = {
72 'yaml': yaml_decode,
73 'json': json.decode,
74 'pid': int,
75 'out': identity,
76 'err': identity,
77}
79ext_encoder = {
80 'yaml': yaml_encode,
81 'json': json.encode,
82 'pid': str,
83 'out': identity,
84 'err': identity,
85}
87ext_aliases = {
88 'fp': 'yaml',
89}
91for k, v in ext_aliases.items():
92 ext_decoder[k] = ext_decoder[v]
93 ext_encoder[k] = ext_encoder[v]
95def daemonize(f):
96 """Decorator for daemonized methods"""
97 return f
99ATTACHED_SCHEME = 'attached'
100class USession(object):
101 """Create a session workspace to store persistent information related to
102 a process group.
103 """
104 WORKSPACE = SessionWorkspace # Need to be overrided
106 # file config
107 default_config = {
108 'copy': {
109 'bin': 'bin/*.*',
110 'out': '',
111 'pid': '',
112 'fp': '',
113 'err': '',
114 },
115 'templates': {
116 'scripts': 'bin/*.sh',
117 },
118 'dependences': {},
119 }
120 """Default configuration for a USession. Can be overrided."""
122 def __init__(self, path):
123 self.path = expandpath(path)
124 self.layout = FSLayout(self.path)
126 self.db_file = self.layout.get_path('db', 'session', touch=True)
127 self.db = DB(self.db_file)
128 self.workspace = self.WORKSPACE(self.db)
129 self.config = self.layout.get_content('etc', 'session',
130 default=self.default_config)
132 def init(self):
133 """Initialize the file structure, adding missing files from
134 config['copy'] section.
135 """
136 info(f'Initialize structure in {self.path}')
138 top = os.path.dirname(os.path.abspath(__file__))
139 for name, pattern in self.config.get('copy', {}).items():
140 self.layout.get_path('<folder>', name, touch=True)
141 for path in fileiter(top, wildcard=pattern):
142 where = os.path.join(self.path, path.split(top)[-1][1:])
143 copyfile(path, where, override=False)
145 def activate(self):
146 """Activate Session Supervisor and optionally launch attached
147 processes based on configuration.
148 """
149 scheme = self.__class__.__name__.lstrip('U').lower()
150 uri = f"{scheme}://supervisor"
151 self.start_daemon(target=self._supervisor, uri=uri, alias='supervisor')
153 def deactivate(self):
154 """Deactivate Session Supervisor and optionally attached processes
155 based on configuration.
156 """
157 def get_candidates():
158 for d, filename, pid in self.layout.iter_file_content('pid'):
159 if test_alive(pid):
160 for process in self.workspace.find(UProcess, name=d['name'], shutdown=1):
161 # assert pid == process.pid
162 yield d['name'], pid
164 def done(name):
165 print(f' {RED}{name}\t\t[done]{RESET}')
167 # signals = [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM, signal.SIGKILL]
168 signals = [signal.SIGQUIT, signal.SIGTERM, signal.SIGKILL]
169 send_signals(get_candidates, signals, done_callback=done)
171 def start_daemon(self, target, **kw):
172 """Start a service (target) as a daemon process.
174 `target` could be a function or an UApplication instance.
175 """
176 context = self._get_daemon_context(**kw)
177 with _call(daemon.DaemonContext, **context) as bel:
178 try:
179 context = self._prepare_start(**context)
180 # run forever ...
181 async_call(target, **context)
182 except Exception as why:
183 print(f"{YELLOW}ERROR: {RED}{why}{RESET}")
185 def stop(self, uri=None, name=None, **kw):
186 """Stop a daemon process."""
187 # try to find the process in DB
188 assert False, "TODO: remove dependeces from DB"
189 for process in list(self.workspace.find(UProcess, uri=uri, name=name)):
190 if process.pid:
191 print(f"{BLUE}Killing '{process.name}' ", end='', flush=True)
192 for tries in range(100):
193 if test_process(process.pid, signal.SIGTERM):
194 time.sleep(0.2)
195 print(".", end='', flush=True)
196 else:
197 process.pid = None
198 process.state = 'dead'
199 # self.workspace.update(process)
200 print(f"{RED}[dead]", flush=True)
201 break
202 # give a chance to end and pid has been removed
203 else:
204 print(f"{YELLOW}Process '{process.name}' is already dead")
206 # try to find process in pid file
207 for filename in fileiter('pid', wildcard=f'{name}.pid'):
208 test_process(filename, signal.SIGTERM)
210 print(f"{RESET}")
212 # alias=None, label='running', shutdown=True):
213 def attach(self, **kw):
214 """Attach a process/service to the session.
216 - `proc`: psutil.Process instance.
217 - `alias`: human alias for the process.
219 As optional keywords:
221 - `restart`: whenever a process must restarted if dies.
222 - `shutdown`: whenever a process must be shutdown when session is \
223 deactivated.
224 """
225 proc = kw.get('proc')
226 uri = kw.get('uri')
227 name = kw.get('name')
228 state = kw.setdefault('state', 'unknown')
229 alias = kw.get('alias')
231 # test_alive(self.layout.get_content('etc'))
232 # if not proc:
233 # for process in self.workspace.find(**kw):
236 if proc:
237 update_context(kw, proc)
238 fp = get_fingerprint(proc)
239 kw.setdefault('cmdline', fp['cmdline'])
240 uri = build_uri(scheme=ATTACHED_SCHEME,
241 path=hasheable_cmdline(fp['cmdline']))
242 kw.setdefault('uri', uri)
244 if uri:
245 name = kw.setdefault('name', codenamize(uri))
247 if name:
248 # save process config file
249 wanted = ['alias', 'cmdline', 'name', ]
250 conf = dict([(k, v) for (k, v) in kw.items() if k in wanted])
251 self.layout.update_content('etc', conf, name)
252 if alias:
253 self.layout.set_alias('etc', alias, name)
255 if proc and name:
256 # create a fingerprint file for launching the file when needed
257 # and include a pidfile in pid structure
258 self.layout.set_content('pid', proc.pid, name)
259 filename = self.layout.set_content('fp', fp, name, state)
260 update_context(kw, filename=filename)
262 # insert process in DB
263 process = UProcess(**kw)
264 self.workspace.update(process)
265 update_context(kw, process=process)
266 return kw
268 def status(self):
269 """**TBD:** Return session processes and status.
270 """
271 raise NotImplemented("implementation pending")
273 def dettach_process(self, proc: psutil.Process = None, alias=None):
274 """**TBD:** Dettach a process from the session.
276 Any of the following arguments can be used:
278 - `proc`: psutil.Process instance.
279 - `alias`: human alias of the process.
280 """
281 raise NotImplemented("implementation pending")
283 def add_dependence(self, parent, child):
284 """Add a dependence between processes"""
285 modified = False
286 for parent in parent.split(','):
287 for child in child.split(','):
288 d = self.config.setdefault('dependences', dict())
289 # don't use set() to preserve serialization
290 s = d.setdefault(child, list())
291 if parent not in s:
292 s.append(parent)
293 modified = True
295 if modified:
296 self.layout.set_content('etc', self.config, 'session', )
298 def remove_dependence(self, parent, child):
299 """Remove a dependence between processes"""
300 modified = False
301 for parent in parent.split(','):
302 for child in child.split(','):
303 d = self.config.get('dependences', dict())
304 # don't use set() to preserve serialization
305 s = d.get(parent, list())
306 if child in s:
307 s.remove(child)
308 modified = True
310 if modified:
311 self.layout.set_content('etc', self.config, 'session', )
313 def launch_process(self, match_if_same_cmdline=True, **kw):
314 """Launch a previous captured process by name.
316 - `match_if_same_cmdline`: consider the process to be the same if \
317 cmdline is the same that stored one.
318 - `kw`: contains search criteria keywords:
319 Examples of search keywords are:
321 - name='adaptable-joint'
322 - state='operative'
323 - alias='master.tws'
324 - restart=1
325 - shutdown=0
326 - uri='session://supervisor'
328 #. try to find a process that match criteria.
329 #. if score is low but has the same command line (optional)
330 we consider that is the same process in another non-registered state
332 #. update DB when process is found
333 #. launch a new process using the command line from config.
334 """
335 for process, score, proc, state, fp in self.find_processes(**kw):
336 if proc:
337 if score >= 0.9 or \
338 match_if_same_cmdline and fp['cmdline'] == proc.cmdline():
339 # self.layout.set_content('fp', fp, process.name, 'running')
340 self.workspace.update(process, pid=process.pid, state=state)
341 continue
343 # launch a new process if all dependeces are matched
344 for dep in kget(self.config['dependences'], process.name, process.alias) or []:
345 parent, state = dep.split(':')
346 pname, parent = vget(processes, parent, 'name', 'alias')
347 if not (parent and parent.state == state):
348 break # any dependence is not matched
349 else:
350 # launch the process
351 conf = self.layout.get_content('etc', process.name)
352 conf = self._launch_process(**conf)
354 def _launch_process(self, **data):
355 if 'cmdline' not in data:
356 uri = parse_uri(data['uri'])
357 data['cmdline'] = cmdline = [
358 uri['path'],
359 ]
360 for k, v in uri['query_'].items():
361 cmdline.append(f'{k}={v}')
363 conf = self._prepare_start(**data)
364 name = conf['name']
366 stdout = open(self.layout.get_path('out', name), 'w')
367 stderr = open(self.layout.get_path('err', name), 'w')
368 # conf = self.layout.get_content('etc', name)
369 p = subprocess.Popen(conf['cmdline'], stdout=stdout, stderr=stderr)
370 self.layout.set_content('pid', p.pid, name)
372 process = UProcess(**conf)
373 self.workspace.update(process, pid=p.pid, state='starting')
375 # wait until process seems to be 'stable' or timeout
376 fp0, fp, state = None, None, conf['state']
377 for trie in range(15): # timeout=30 secs
378 proc = psutil.Process(p.pid)
379 fp0, fp = fp, get_fingerprint(proc)
380 self.layout.set_content('fp', fp, name, state)
381 if fp0 == fp:
382 break
383 print(f"[{trie}] waiting for process '{process.alias or name}' to be stable...({state})")
384 time.sleep(5)
385 state = 'running'
387 update_context(conf, proc=proc, fp=fp, process=process)
388 return conf
390 def _prepare_start(self, uri=None, name=None, alias=None, **conf):
391 name = name or codenamize(uri)
392 alias = alias or name
394 uri_ = parse_uri(uri)
395 update_context(conf, uri_, uri_['query_'])
396 realm = conf.setdefault('realm', uri_['path'])
398 conf.setdefault('cmdline', sys.argv)
399 update_context(conf, uri=uri, name=name, alias=alias)
400 self.layout.set_content('etc', conf, name)
401 self.layout.set_alias('etc', alias, name)
403 conf['pid'] = pid = os.getpid()
404 conf['state'] = 'starting'
406 # TODO: config, tl, hub, db, session, etc
407 # loop = conf['loop'] = asyncio.get_event_loop()
408 hub = conf['hub'] = Hub(realm=realm)
410 # TODO: add any maintenance/domestic procedures
411 procedures = [dummy_main, ]
412 app = _call(UApplication, procedures=procedures,
413 fixtures=conf, **conf)
414 app.tl = TL(app=app, uid=conf.get('uid', name))
416 proc = psutil.Process(pid)
417 fp = get_fingerprint(proc)
418 # self.layout.set_content('fp', fp, name, conf['state'])
419 process = UProcess(**conf)
420 self.workspace.update(process)
422 update_context(conf, fp=fp, proc=proc, process=process,
423 app=app, tl=app.tl, hub=app.hub)
426 # app.run()
427 # run(app.tl.start())
429 return conf
431 def launch_service(self, uri, alias=None):
432 """Launch a service (a method instance) parsing calling arguments\
433 from uri.
435 Example:
437 ``session.launch_service(
438 uri="ib://tws:9090/download-historical?uid=historical&account=0&tws=tws.demo",
439 alias="downloader")``
440 """
441 uri_ = parse_uri(uri)
442 func = self
443 for name in uri_['path'].split('/'):
444 name = name.replace('-', '_')
445 func = getattr(func, name, self)
447 alias = alias or name
448 if func:
449 self.start_daemon(target=func, uri=uri, alias=alias)
450 else:
451 raise RuntimeError(f"can't find a function linked with f{uri}")
453 def _supervisor(self):
454 "Supervisor daemon tasks"
456 def check_attached_processes():
457 self.launch_process(restart=1)
459 def hide_check_attached_processes():
460 processes = dict()
462 # update fps states
463 for process in list(self.workspace.find(UProcess, restart=1)):
464 fps = self._get_available_fps(name=process.name)
465 score, proc, label, fp = self._best_matching_fps(
466 fps, process.pid)
467 if score < 0.90:
468 process.state, process.pid = None, None
469 else:
470 process.state, process.pid = label, proc.pid
471 self.workspace.update(process)
472 processes[process.name] = process
474 for name, process in processes.items():
475 if not process.pid:
476 # check dependeces/pre-requsites
477 for dep in kget(self.config['dependences'], process.name, process.alias) or []:
478 parent, state = dep.split(':')
479 pname, parent = vget(processes,
480 parent, 'name', 'alias')
481 if not (parent and parent.state == state):
482 break
483 else:
484 # launch the process/service
485 self.launch_process(name=process.name)
486 # # TODO: process and service are the same?
487 # if process.uri.startswith(ATTACHED_SCHEME):
488 # self.launch_process(name=process.name)
489 # else:
490 # self.launch_service(process.uri)
491 foo = 1
493 def clean_dead_pid():
494 count = 0
495 for filename in fileiter('pid', wildcard='*.pid'):
496 print(f" {YELLOW}Testing {filename:40}", end='')
497 proc = test_alive(filename)
498 name = os.path.basename(filename)
499 name = os.path.splitext(name)[0]
500 path = self.layout.get_path('etc', name)
501 if os.path.exists(path) and proc:
502 count += 1
503 print(f"{GREEN}[ok] {proc.pid}")
504 else:
505 print(f"{RED}[dead]")
506 os.remove(filename)
507 print(f" {GREEN}[{count}]{YELLOW} running processes{RESET}")
509 FUNCTIONS = [clean_dead_pid, check_attached_processes]
510 functions = []
512 # for tries in range(n):
513 while True:
514 if not functions:
515 functions.extend(FUNCTIONS)
516 # random.shuffle(functions)
517 if functions:
518 func = functions.pop()
519 print(f"{BLUE}- Running: {func.__name__}{RESET}")
520 func()
521 time.sleep(5)
523 print("Bye!")
525 def _get_daemon_context(self, **kw):
526 gpath = self.layout.get_path
527 s = partial(dset, kw)
528 name = s('name', codenamize(kw['uri']))
530 s('pid', gpath('pid', name))
531 s('pidfile', SmartPIDLockFile(path=kw['pid']))
532 s('files_preserve', list(range(0, 255)))
533 s('stdin', None)
534 s('stdout', open, gpath('out', name), 'w')
535 s('stderr', open, gpath('err', name), 'w')
536 s('working_directory', self.path)
537 s('shutdown', 0)
538 s('restart', 0)
539 # s('signal_map', {})
540 return kw
542 def _demo_task(self, n):
543 for i in range(n):
544 print(f"doing something cool: {i}")
545 time.sleep(5)
546 print("-End demo task-")
548 def _get_available_fps(self, name=None):
549 fingerprints = dict()
550 if isinstance(name, str):
551 regexp = re.compile(name, re.DOTALL)
552 else:
553 regexp = None
554 for info, filename, fp0 in self.layout.iter_file_content('fp'):
555 name = info['name']
556 if regexp and not regexp.match(name):
557 continue
558 fps = fingerprints.setdefault(name, dict())
559 fps[info['label']] = fp0
561 return fingerprints
563 def _best_matching_fp(self, fp0, *pids):
564 pids = [p for p in pids if p] or psutil.pids()
566 best_score, best_proc, best_fp = 0, None, None
567 for pid_ in pids:
568 proc = psutil.Process(pid_)
569 fp = get_fingerprint(proc)
570 s = compare_fingerprint(fp0, fp)
571 if s > best_score:
572 best_score, best_proc, best_fp = s, proc, fp
574 return best_score, best_proc, best_fp
576 def _best_matching_fps(self, fps, pids=[], name=None, deep_serch=False):
577 candidates = set([])
578 if not isinstance(pids, list):
579 pids = [pids]
580 pids = [p for p in pids if p]
581 if not pids:
582 deep_serch = True
583 candidates.update(pids)
585 if deep_serch:
586 candidates.update(psutil.pids())
588 def do_search(pids):
589 best_score, best_proc, best_label, best_fp = 0, None, None, None
590 for pid_ in pids:
591 # if pid_ < 1024: # not in Linux
592 # continue # system proceces
593 try:
594 proc = psutil.Process(pid_)
595 fp = get_fingerprint(proc)
596 except Exception as why:
597 continue # may be a zomby/protected process
599 for pname, fps_ in fps.items():
600 if name not in (pname, None):
601 continue
603 for label, fp0 in fps_.items():
604 # if not label and exclude_no_label:
605 # continue
606 s = compare_fingerprint(fp0, fp)
607 if s > best_score:
608 best_score, best_proc, best_label, best_fp = s, proc, label, fp
609 if s == 1.0:
610 return best_score, best_proc, best_label, best_fp
612 return best_score, best_proc, best_label, best_fp
614 return do_search(candidates)
617 def locate_process(self, **data):
618 for process in self.workspace.find(UProcess, **data):
619 fps = self._get_available_fps(name=process.name)
620 score, proc, state, fp = self._best_matching_fps(fps, name=process.name)
621 if score > 0.9:
622 return score, proc, state, fp
623 return [None] * 4
625 def find_processes(self, join='AND', **query):
626 """Find running processes that match a query"""
627 for process in self.workspace.find(UProcess, join=join, **query):
628 fps = self._get_available_fps(name=process.name)
629 score, proc, state, fp = self._best_matching_fps(
630 fps, name=process.name)
631 yield process, score, proc, state, fp
633 def _get_preferred_fp(self, name):
634 desired_labels = ['operative', 'running', 'launch', ]
635 candidates = dict()
636 # for d, filename, fp0 in self.layout.iter_file_content('fp', name):
637 # candidates[d['label']] = d, filename, fp0
639 for d, filename, fp0 in self.layout.iter_file_content('fp', name):
640 candidates[d['label']] = d, filename, fp0
642 if not candidates:
643 return {}, '', {}
645 desired_labels.extend(candidates.keys())
646 for label in desired_labels:
647 if label not in candidates:
648 continue
649 d, filename, fp0 = candidates[label]
650 break
652 return d, filename, fp0
654 def check_structure(self):
655 """Check is a USession folder has been Initialized."""
656 debug(f'Checking structure in {self.path}')
657 return os.path.exists(self.db_file)
659 # def load(self):
660 # warnings.warn("Load is not longer necessary", DeprecationWarning, stacklevel=2)
661 # self.config = load_config(self.session_config_file)
662 # if self.config:
663 # self.db = DB(self.db_file)
664 # self.workspace = self.WORKSPACE(self.db)
665 # else:
666 # raise RuntimeError("Config not found: maybe need session init here ?")
668 def update_config(self, value, *keys):
669 """Helper to save the session configuration."""
670 container = self.config
671 if update_container(container, value, *keys):
672 self.layout.set_content('etc', container, 'session')