libpipe 0.5__tar.gz → 0.6__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,8 +1,9 @@
1
- Metadata-Version: 2.3
1
+ Metadata-Version: 2.4
2
2
  Name: libpipe
3
- Version: 0.5
3
+ Version: 0.6
4
4
  Summary: Pipeline orchestration and node management utilities
5
5
  License: GPL-3.0-or-later
6
+ License-File: LICENSE
6
7
  Author: Florent Mertens
7
8
  Author-email: florent.mertens@gmail.com
8
9
  Requires-Python: >=3.10
@@ -12,10 +13,12 @@ Classifier: Programming Language :: Python :: 3.10
12
13
  Classifier: Programming Language :: Python :: 3.11
13
14
  Classifier: Programming Language :: Python :: 3.12
14
15
  Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Programming Language :: Python :: 3.14
15
17
  Requires-Dist: asyncssh (>=2.0)
16
18
  Requires-Dist: click (>=8.0)
17
19
  Requires-Dist: progressbar2 (>=3.40)
18
20
  Requires-Dist: toml (>=0.10)
21
+ Project-URL: Repository, https://gitlab.com/flomertens/libpipe
19
22
  Description-Content-Type: text/markdown
20
23
 
21
24
  Libpipe
@@ -6,6 +6,7 @@ import os
6
6
  import re
7
7
  import sys
8
8
  import time
9
+ import shlex
9
10
  import socket
10
11
  import signal
11
12
  import getpass
@@ -70,8 +71,7 @@ class WorkerStatus(Enum):
70
71
  IDLE = 1
71
72
  COMMITED = 2
72
73
  RUNNING = 3
73
- ERROR = 4
74
- CLOSED = 5
74
+ CLOSED = 4
75
75
 
76
76
 
77
77
  class Process(object):
@@ -146,7 +146,7 @@ class Task(object):
146
146
 
147
147
  async def terminate(self):
148
148
  if self.process is not None and hasattr(self.process.process, 'pid'):
149
- if self.process.client.host is localhost_shortname:
149
+ if self.process.client.host == localhost_shortname:
150
150
  self.log(f'Killing local process PID {self.process.process.pid}\n', 'local', err=True)
151
151
  try:
152
152
  os.killpg(os.getpgid(self.process.process.pid), signal.SIGTERM)
@@ -169,12 +169,13 @@ class Task(object):
169
169
  async def log_stream(self, stream, host, err=False):
170
170
  while True:
171
171
  try:
172
- async for line in stream:
173
- self.log(line, host, err=err)
172
+ line = await stream.readline()
174
173
  except ValueError:
175
174
  self.log('(line too long)', host, err=True)
176
175
  continue
177
- break
176
+ if not line:
177
+ break
178
+ self.log(line, host, err=err)
178
179
 
179
180
  def close_log(self):
180
181
  if self.pid_filename and os.path.exists(self.pid_filename):
@@ -198,6 +199,11 @@ class Task(object):
198
199
 
199
200
  ret_status, err_msg = await worker.run(self)
200
201
 
202
+ if not ret_status:
203
+ # mark host temporarily bad so next retry likely uses a different host
204
+ host = worker.client.host
205
+ client_pool.bad_hosts_until[host] = time.time() + client_pool.host_cooldown_s
206
+
201
207
  if err_msg:
202
208
  self.log(err_msg, worker.client.host, err=True)
203
209
 
@@ -230,10 +236,11 @@ class SSHClient(asyncssh.SSHClient):
230
236
 
231
237
  class Client(object):
232
238
 
233
- def __init__(self, host, user=None, force_sync=True, password=None):
239
+ def __init__(self, host, user=None, force_sync=True, password=None, connect_timeout=15):
234
240
  self.host = host
235
241
  self.user = user
236
242
  self.password = password
243
+ self.connect_timeout = connect_timeout
237
244
  self.conn = None
238
245
  self.client = None
239
246
  self.starting = False
@@ -241,18 +248,45 @@ class Client(object):
241
248
  self.closing = False
242
249
  self.force_sync = force_sync
243
250
  self.creating_session = asyncio.Lock()
251
+ self.start_lock = asyncio.Lock()
244
252
 
245
253
  async def start(self):
246
- if not self.starting:
254
+ # Only one concurrent connection attempt per Client
255
+ async with self.start_lock:
256
+ # If already connected, nothing to do
257
+ if self.connected():
258
+ self.started.set()
259
+ return
260
+
261
+ # Allow callers to wait on "started" for this attempt
262
+ if self.started.is_set():
263
+ self.started.clear()
264
+
247
265
  self.starting = True
248
266
  print(f'Starting client {self.host} ...')
267
+
249
268
  try:
250
- self.conn, self.client = await asyncssh.create_connection(SSHClient, self.host, username=self.user,
251
- password=self.password, agent_forwarding=True)
269
+ self.conn, self.client = await asyncssh.create_connection(
270
+ SSHClient,
271
+ self.host,
272
+ username=self.user,
273
+ password=self.password,
274
+ agent_forwarding=True,
275
+ connect_timeout=self.connect_timeout
276
+ )
252
277
  self.conn.set_keepalive(120, 5)
278
+
253
279
  except Exception:
280
+ # IMPORTANT: reset state so a later retry actually retries
281
+ self.conn = None
282
+ self.client = None
283
+ self.starting = False
284
+ self.started.set()
254
285
  raise
255
- finally:
286
+
287
+ else:
288
+ # Connected successfully
289
+ self.starting = False
256
290
  self.started.set()
257
291
 
258
292
  async def execute(self, task):
@@ -269,8 +303,10 @@ class Client(object):
269
303
 
270
304
  task.set_process(Process(self, process))
271
305
 
272
- await task.log_stream(process.stdout, self.host)
273
- await task.log_stream(process.stderr, self.host, err=True)
306
+ await asyncio.gather(
307
+ task.log_stream(process.stdout, self.host),
308
+ task.log_stream(process.stderr, self.host, err=True),
309
+ )
274
310
 
275
311
  process.channel.close()
276
312
  task.returncode = process.returncode
@@ -304,9 +340,10 @@ class LocalClient(object):
304
340
  stderr=asyncio.subprocess.PIPE, preexec_fn=os.setpgrp)
305
341
  task.set_process(Process(self, process))
306
342
 
307
- await task.log_stream(process.stdout, self.host)
308
- await task.log_stream(process.stderr, self.host, err=True)
309
-
343
+ await asyncio.gather(
344
+ task.log_stream(process.stdout, self.host),
345
+ task.log_stream(process.stderr, self.host, err=True),
346
+ )
310
347
  task.returncode = await process.wait()
311
348
 
312
349
  if self.force_sync:
@@ -340,7 +377,12 @@ class Worker(object):
340
377
  await self.client.start()
341
378
  except Exception as exc:
342
379
  err_msg = f'Error starting client: {str(exc)}\n'
343
- self.status = WorkerStatus.ERROR
380
+ # transient failure: allow task retry to pick a worker again
381
+ self.status = WorkerStatus.IDLE
382
+ try:
383
+ self.client.close() # ensure clean slate
384
+ except Exception:
385
+ pass
344
386
  return False, err_msg
345
387
 
346
388
  await self.client.started.wait()
@@ -349,7 +391,11 @@ class Worker(object):
349
391
  await self.client.execute(task)
350
392
  except ConnectionError as exc:
351
393
  err_msg = f'Error connecting client: {str(exc)}\n'
352
- self.status = WorkerStatus.ERROR
394
+ self.status = WorkerStatus.IDLE
395
+ try:
396
+ self.client.close()
397
+ except Exception:
398
+ pass
353
399
  return False, err_msg
354
400
  except asyncio.CancelledError:
355
401
  err_msg = f'Task {self.name} has been canceled\n'
@@ -360,6 +406,12 @@ class Worker(object):
360
406
  return False, err_msg
361
407
 
362
408
  err_msg = task.done_callback()
409
+ # err_msg = ''
410
+ # try:
411
+ # loop = asyncio.get_running_loop()
412
+ # err_msg = await loop.run_in_executor(None, task.done_callback)
413
+ # except Exception as exc:
414
+ # err_msg = f'Error executing callback of task {self.name}: {exc}\n'
363
415
 
364
416
  self.status = WorkerStatus.IDLE
365
417
 
@@ -370,11 +422,42 @@ class Worker(object):
370
422
  self.client.close()
371
423
 
372
424
 
425
+ class WorkerPoolReport:
426
+
427
+ def __init__(self, all_tasks):
428
+ self.all_tasks = list(all_tasks)
429
+ self.n_total = len(self.all_tasks)
430
+
431
+ # Buckets by status
432
+ self.success = [t for t in self.all_tasks if t.status == TaskStatus.SUCCESS]
433
+ self.failed = [t for t in self.all_tasks if t.status != TaskStatus.SUCCESS]
434
+
435
+ # Within SUCCESS, split by return code
436
+ self.success_rc0 = [t for t in self.success if t.returncode == 0]
437
+ self.success_nonzero = [t for t in self.success if t.returncode not in (None, 0)]
438
+ self.success_unknown_rc = [t for t in self.success if t.returncode is None]
439
+
440
+ # Counts
441
+ self.n_success = len(self.success)
442
+ self.n_failed = len(self.failed)
443
+ self.n_success_rc0 = len(self.success_rc0)
444
+ self.n_success_nonzero = len(self.success_nonzero)
445
+ self.n_success_unknown_rc = len(self.success_unknown_rc)
446
+
447
+ self.ok = (self.n_failed == 0 and self.n_success_nonzero == 0)
448
+
449
+ def host_of(self, task):
450
+ if task.process and task.process.client:
451
+ return getattr(task.process.client, "host", "?")
452
+ return task.run_on_host or "?"
453
+
454
+
373
455
  class WorkerPool(object):
374
456
 
375
457
  def __init__(self, hosts, name='Worker', max_tasks_per_worker=4,
376
458
  env_source_file=None, user=None, max_time=None, force_sync=False,
377
- debug=False, dry_run=False, password=None):
459
+ debug=False, dry_run=False, password=None, connect_timeout=15,
460
+ host_cooldown_s=10):
378
461
  """Initiate a worker pool.
379
462
 
380
463
  Args:
@@ -383,6 +466,8 @@ class WorkerPool(object):
383
466
  max_tasks_per_worker (int, optional): Maximum number of tasks to execute concurrently on an host.
384
467
  env_source_file (str, optional): Name of a file to source before executing a task.
385
468
  user (str, optional): User to connect
469
+ connect_timeout (int, optional): Per-attempt SSH connection timeout, in seconds.
470
+ host_cooldown_s (int, optional): Time a just-failed host is avoided while a healthy one exists.
386
471
  """
387
472
  if asyncio.get_event_loop().is_closed():
388
473
  asyncio.set_event_loop(asyncio.new_event_loop())
@@ -399,13 +484,17 @@ class WorkerPool(object):
399
484
  self.dry_run = dry_run
400
485
  self.hosts = hosts
401
486
  self.futures = None
402
- self.active_hosts = hosts
487
+ self.last_report = None
488
+ # host -> epoch time until which a just-failed host is avoided (steering only)
489
+ self.host_cooldown_s = host_cooldown_s
490
+ self.bad_hosts_until = {}
403
491
 
404
492
  for host in self.hosts:
405
493
  if host == localhost_shortname or host == 'localhost':
406
494
  client = LocalClient(localhost_shortname, force_sync=force_sync)
407
495
  else:
408
- client = Client(host, user, force_sync=force_sync, password=password)
496
+ client = Client(host, user, force_sync=force_sync, password=password,
497
+ connect_timeout=connect_timeout)
409
498
 
410
499
  for i in range(max_tasks_per_worker):
411
500
  worker = Worker(client, f'{client.host}:{i}')
@@ -423,15 +512,26 @@ class WorkerPool(object):
423
512
 
424
513
  async def get_worker(self, host=None):
425
514
  while True:
515
+ now = time.time()
516
+
517
+ # hosts not currently in cooldown
518
+ good_hosts = {h for h in self.hosts if self.bad_hosts_until.get(h, 0) <= now}
519
+
426
520
  for worker in self.workers:
427
- if host is not None and worker.client.host != host:
428
- if host not in self.active_hosts:
429
- print(f'Requested run_on_host "{host}" not available anymore. Releasing constraint ...')
430
- host = None
521
+ w_host = worker.client.host
522
+
523
+ # respect run_on_host if requested
524
+ if host is not None and w_host != host:
525
+ continue
526
+
527
+ # if we have at least one good host, avoid bad ones
528
+ if host is None and good_hosts and w_host not in good_hosts:
431
529
  continue
530
+
432
531
  if worker.status is WorkerStatus.IDLE:
433
532
  worker.status = WorkerStatus.COMMITED
434
533
  return worker
534
+
435
535
  await asyncio.sleep(0.1)
436
536
 
437
537
  def add(self, command, name=None, output_file=None, task_done_callback=None, done_callback=None,
@@ -448,9 +548,11 @@ class WorkerPool(object):
448
548
  run_on_host = None
449
549
 
450
550
  if self.env_file:
451
- command = f"bash -c '. {self.env_file}; {command}'"
551
+ script = f". {shlex.quote(self.env_file)}; {command}"
452
552
  else:
453
- command = f"bash -c '{command}'"
553
+ script = command
554
+
555
+ command = f"bash -lc {shlex.quote(script)}"
454
556
 
455
557
  if name is None:
456
558
  name = f'T{len(self.tasks)}'
@@ -480,26 +582,16 @@ class WorkerPool(object):
480
582
  n_todo = len(tasks[TaskStatus.WAITING]) + len(tasks[TaskStatus.RUNNING])
481
583
  self.pbar.update(n_done)
482
584
 
483
- self.active_hosts = set(w.client.host for w in self.workers if w.status != WorkerStatus.ERROR)
484
- n_active_workers = 0
485
-
486
- for worker in self.workers:
487
- n_active_workers += worker.status != WorkerStatus.ERROR
488
-
489
- if worker.status == WorkerStatus.RUNNING and self.max_time is not None:
490
- running_time = time.time() - worker.execute_start_time
491
- if running_time > (self.max_time + update_interval):
492
- task = worker.running_task
493
- task.timed_out = True
494
- task.log(f'Task timeout: waited for {running_time:.2f} s', worker.client.host, err=True)
495
- task.log(f'Command was: {task.command}', worker.client.host, err=True)
496
- await task.terminate()
497
-
498
- if n_active_workers == 0:
499
- secho("No active workers, canceling all remaining tasks ...", fg='red')
500
- if self.futures is not None:
501
- for future in self.futures[::-1]:
502
- future.cancel()
585
+ if self.max_time is not None:
586
+ for worker in self.workers:
587
+ if worker.status == WorkerStatus.RUNNING:
588
+ running_time = time.time() - worker.execute_start_time
589
+ if running_time > (self.max_time + update_interval):
590
+ task = worker.running_task
591
+ task.timed_out = True
592
+ task.log(f'Task timeout: waited for {running_time:.2f} s', worker.client.host, err=True)
593
+ task.log(f'Command was: {task.command}', worker.client.host, err=True)
594
+ await task.terminate()
503
595
 
504
596
  except asyncio.CancelledError:
505
597
  pass
@@ -530,7 +622,10 @@ class WorkerPool(object):
530
622
 
531
623
  await asyncio.gather(*self.futures, return_exceptions=False)
532
624
 
533
- def execute(self):
625
+ def get_report(self):
626
+ return self.last_report
627
+
628
+ def execute(self, print_summary=False):
534
629
  if not self.tasks:
535
630
  self.pbar.finish()
536
631
  return [], []
@@ -556,5 +651,26 @@ class WorkerPool(object):
556
651
  self.pbar.update(n_success)
557
652
  self.pbar.finish(dirty=n_success != len(self.tasks))
558
653
 
654
+ self.last_report = WorkerPoolReport(self.tasks)
655
+ report = self.last_report
656
+
657
+ if print_summary:
658
+ print("\n=== WorkerPool summary ===")
659
+ print(f"Successful tasks (rc=0): {report.n_success_rc0} / {report.n_total}")
660
+ print(f"Completed with non-zero rc: {report.n_success_nonzero} / {report.n_total}")
661
+ print(f"Failed tasks (exception/kill): {report.n_failed} / {report.n_total}")
662
+
663
+ if report.success_nonzero:
664
+ print("\nTasks completed with non-zero return code:")
665
+ for t in report.success_nonzero:
666
+ print(f" - {t.name} | host={report.host_of(t)} | rc={t.returncode}")
667
+
668
+ if report.failed:
669
+ print("\nFailed tasks:")
670
+ for t in report.failed:
671
+ print(f" - {t.name} | status={t.status.name} | host={report.host_of(t)} | rc={t.returncode}")
672
+
673
+ print("==========================\n")
674
+
559
675
  return tasks[TaskStatus.SUCCESS], tasks[TaskStatus.ERROR] + tasks[TaskStatus.CANCELED] \
560
676
  + tasks[TaskStatus.WAITING] + tasks[TaskStatus.RUNNING]
@@ -4,14 +4,13 @@ build-backend = "poetry.core.masonry.api"
4
4
 
5
5
  [project]
6
6
  name = "libpipe"
7
- version = "0.5"
7
+ version = "0.6"
8
8
  description = "Pipeline orchestration and node management utilities"
9
9
  readme = "README.md"
10
10
  license = { text = "GPL-3.0-or-later" }
11
11
  authors = [
12
12
  { name = "Florent Mertens", email = "florent.mertens@gmail.com" }
13
13
  ]
14
- repository = "https://gitlab.com/flomertens/libpipe"
15
14
  requires-python = ">=3.10"
16
15
 
17
16
  dependencies = [
@@ -21,9 +20,15 @@ dependencies = [
21
20
  "click>=8.0",
22
21
  ]
23
22
 
23
+ [project.urls]
24
+ Repository = "https://gitlab.com/flomertens/libpipe"
25
+
24
26
  [project.scripts]
25
27
  nodetool = "libpipe.tools.nodetool:main"
26
28
 
27
29
  [tool.poetry.group.dev.dependencies]
28
30
  pytest = "^8.0"
29
31
  pytest-cov = "^5.0"
32
+
33
+ [tool.pytest.ini_options]
34
+ testpaths = ["tests"]
@@ -1,73 +0,0 @@
1
- data_dir = '.'
2
- default_settings = ''
3
- debug = false
4
- longitude = 6.57
5
-
6
- [worker]
7
- nodes = 'localhost'
8
- max_concurrent = 4
9
- env_file = ''
10
-
11
- [merge_ms]
12
- apply_aoflagger = true
13
- data_col = 'CORRECTED_DATA'
14
- obs_id_out_suffix = 'flagged'
15
- blmin = 0
16
- blmax = 1000
17
- flag_baselines = ''
18
- avg_timestep = 1
19
- time_start_index = 0
20
- time_end_index = 0
21
-
22
- [ssins]
23
- apply_ssins = true
24
- n_time_avg = [1, 4, 8]
25
- nsigmas.1 = [8, 5, 5]
26
- nsigmas.4 = [8, 5, 5]
27
- nsigmas.8 = [8, 5, 5]
28
- percentage_freq_full_flag = 0.4
29
- percentage_time_full_flag = 0.4
30
- time_freq_threshold = 0.6
31
- baseline_threshold = 0.75
32
- snapshot_threshold = 0.5
33
-
34
- [image]
35
- name = 'default'
36
- umin = 50
37
- umax = 250
38
- scale = '0.5amin'
39
- size = '512 512'
40
- data_col = 'CORRECTED_DATA'
41
- channels_out = 'all'
42
- time_start_index = 0
43
- time_end_index = 0
44
- lst_bins = []
45
- stokes = 'IV'
46
- split_even_odd = true
47
- weight = 'natural'
48
-
49
- [vis_cube]
50
- fov = 4
51
- umin = 50
52
- umax = 500
53
- win_fct = '(tukey,0.2)'
54
-
55
- [vis_to_sph]
56
- obs_id_out_suffix = 'sph'
57
- pre_flag = ''
58
-
59
- [combine]
60
- obs_id_out_prefix = '001'
61
- scale_with_noise = false
62
- weights_mode = 'full'
63
- inhomogeneous = false
64
- pre_flag = ''
65
-
66
- [power_spectra]
67
- eor_bin_list = ''
68
- ps_config = ''
69
- flagger = ''
70
-
71
- [gpr]
72
- config_i = ''
73
- config_v = ''
@@ -1,305 +0,0 @@
1
- import os
2
- import pathlib
3
- from unittest import mock
4
-
5
- from libpipe import settings
6
-
7
-
8
- class Settings(settings.BaseSettings):
9
-
10
- DEFAULT_SETTINGS = pathlib.Path(__file__).absolute().parent / 'settings_test.toml'
11
-
12
-
13
- def test_settings(tmp_path):
14
- c = Settings(None, {'vis': {'umin': 20, 'umax': 40, 'list': ['a', 'b']}, 'data_path': 'some_path'})
15
-
16
- assert c.vis.umin == 20
17
- assert c.vis.umax == 40
18
- assert c.data_path == 'some_path'
19
- assert c['vis'] == {'umin': 20, 'umax': 40, 'list': ['a', 'b']}
20
- assert c.vis.list == ['a', 'b']
21
- assert c.vis.list[1] == 'b'
22
-
23
- c.image = {'lmin': 100, 'data_col': 'DATA'}
24
- c.obs_ids = [1, 2, 3, 4]
25
-
26
- assert c.image.lmin == 100
27
- assert c.image.data_col == 'DATA'
28
- assert c.obs_ids == [1, 2, 3, 4]
29
- assert c.obs_ids[0] == 1
30
-
31
- s = '''
32
- db_path = 'some_other_path'
33
-
34
- [vis]
35
- list = ['c', 'd']
36
-
37
- [worker]
38
- nodes = 'node1[0-1][1-5]'
39
- concurrent = 2
40
- '''
41
-
42
- filename = tmp_path / 'test.toml'
43
- with open(str(filename), 'w') as f:
44
- f.write(s)
45
-
46
- c2 = Settings.load(str(filename))
47
-
48
- assert c2.vis.list == ['c', 'd']
49
- assert c2.worker.nodes == 'node1[0-1][1-5]'
50
-
51
- c3 = c2.copy()
52
- c3.worker.concurrent = 3
53
- c3.new = 'a'
54
-
55
- assert c2.worker.concurrent == 2
56
- assert c3.worker.concurrent == 3
57
- assert hasattr(c3, 'new')
58
- assert not hasattr(c2, 'new')
59
-
60
- c4 = c + c2
61
-
62
- assert c4.vis.umin == 20
63
- assert c4.vis.umax == 40
64
- assert c4.vis.list == ['c', 'd']
65
- assert c4.worker.nodes == 'node1[0-1][1-5]'
66
- assert c4.worker.concurrent == 2
67
-
68
- c5 = c + c2 + c3
69
-
70
- assert c5.worker.concurrent == 3
71
- assert c5.new == 'a'
72
-
73
- c = Settings.load_from_string(s)
74
-
75
- assert c.vis.list == ['c', 'd']
76
- assert c.worker.nodes == 'node1[0-1][1-5]'
77
-
78
- c = Settings.load_from_string(['db_path="a"', 'vis.c=2'])
79
-
80
- assert c.db_path == 'a'
81
- assert c.vis.c == 2
82
-
83
-
84
- def test_default_settings(tmp_path):
85
- s = '''
86
- db_path = 'some_path'
87
- default_settings = ''
88
-
89
- [vis_cube]
90
- umax = 200
91
- win_fct = 'hann'
92
-
93
- [worker]
94
- nodes = 'node1[0-1][1-5]'
95
- concurrent = 2
96
- '''
97
-
98
- filename = tmp_path / 'config.toml'
99
- with open(str(filename), 'w') as f:
100
- f.write(s)
101
-
102
- c = Settings.load_with_defaults(filename)
103
-
104
- assert c.db_path == 'some_path'
105
- assert c.default_settings == ''
106
- assert c.vis_cube.umax == 200
107
- assert c.vis_cube.win_fct == 'hann'
108
- assert c.worker.nodes == 'node1[0-1][1-5]'
109
- assert c.worker.concurrent == 2
110
- assert c.data_dir == '.'
111
- assert c.vis_cube.umin == 50
112
-
113
- s = f'''
114
- db_path = 'some_other_path'
115
- default_settings = '{str(filename)}'
116
-
117
- [vis_cube]
118
- umax = 400
119
-
120
- [image]
121
- umax = 5000
122
-
123
- [worker]
124
- concurrent = 2
125
- '''
126
-
127
- filename2 = tmp_path / 'config2.toml'
128
- with open(str(filename2), 'w') as f:
129
- f.write(s)
130
-
131
- with mock.patch.object(Settings, 'DEFAULT_SETTINGS', filename):
132
- c = Settings.load_with_defaults(filename2)
133
-
134
- assert c.db_path == 'some_other_path'
135
- assert os.path.abspath(c.default_settings) == os.path.abspath(filename)
136
- assert c.vis_cube.umax == 400
137
- assert c.vis_cube.win_fct == 'hann'
138
- assert c.worker.nodes == 'node1[0-1][1-5]'
139
- assert c.worker.concurrent == 2
140
- assert c.image.umax == 5000
141
-
142
- c.save(strip_parents_keys=True)
143
-
144
- c = Settings.load_with_defaults(filename2)
145
-
146
- assert c.db_path == 'some_other_path'
147
- assert os.path.abspath(c.default_settings) == os.path.abspath(filename)
148
- assert c.vis_cube.umax == 400
149
- assert c.vis_cube.win_fct == 'hann'
150
- assert c.worker.nodes == 'node1[0-1][1-5]'
151
- assert c.worker.concurrent == 2
152
- assert c.image.umax == 5000
153
-
154
- c = Settings.load(filename2)
155
-
156
- print(open(filename2).readlines())
157
- assert c.db_path == 'some_other_path'
158
- assert os.path.abspath(c.default_settings) == os.path.abspath(filename)
159
- assert c.image.umax == 5000
160
- assert c.vis_cube.umax == 400
161
- assert not hasattr(c, 'worker')
162
- assert not hasattr(c.image, 'win_fct')
163
-
164
-
165
- def test_duplicate(tmp_path):
166
- s = '''
167
- db_path = 'some_path'
168
- default_settings = ''
169
-
170
- [worker]
171
- nodes = 'node1[0-1][1-5]'
172
- concurrent = 2
173
- some_parset = 'config/some_file.parset'
174
- '''
175
- os.chdir(tmp_path)
176
-
177
- filename = tmp_path / 'config.toml'
178
- with open(str(filename), 'w') as f:
179
- f.write(s)
180
-
181
- os.makedirs(tmp_path / 'config')
182
-
183
- (tmp_path / 'config' / 'some_file.parset').touch()
184
-
185
- s = Settings.load_with_defaults(filename)
186
-
187
- new_file = tmp_path / 'new_config.toml'
188
-
189
- ns = s.duplicate(new_file)
190
-
191
- assert ns.db_path == 'some_path'
192
- assert ns.worker.some_parset == 'config/some_file.parset'
193
- assert not (tmp_path / 'new_dir' / 'config' / 'some_file.parset').is_file()
194
-
195
- new_file = tmp_path / 'new_dir' / 'new_config.toml'
196
-
197
- ns = s.duplicate(new_file, copy_parset_files=False)
198
-
199
- assert ns.db_path == 'some_path'
200
- assert ns.worker.some_parset == str(tmp_path / 'config' / 'some_file.parset')
201
-
202
- new_file = tmp_path / 'new_dir' / 'new_config.toml'
203
-
204
- ns = s.duplicate(new_file, copy_parset_files=True)
205
-
206
- assert ns.db_path == 'some_path'
207
- assert ns.worker.some_parset == 'config/some_file.parset'
208
- assert (tmp_path / 'new_dir' / 'config' / 'some_file.parset').is_file()
209
-
210
- ns = s.duplicate(new_file, update_relative_path=True, **{'db_path': 'added_path', 'worker': {'concurrent': 4}})
211
-
212
- assert ns.db_path == 'added_path'
213
- assert ns.worker.concurrent == 4
214
-
215
-
216
- def test_get_path(tmp_path):
217
- c = Settings(None, {'data_dir': 'a'})
218
- assert c.get_path('data_dir') == str(pathlib.Path.cwd() / 'a')
219
-
220
- c = Settings('/some/file.toml', {'data_dir': 'a'})
221
- assert c.get_path('data_dir') == '/some/a'
222
-
223
- c = Settings(None, {'data_dir': '/a'})
224
- assert c.get_path('data_dir') == '/a'
225
-
226
- c = Settings('/some/file.toml', {'data_dir': '/a'})
227
- assert c.get_path('data_dir') == '/a'
228
-
229
- c = Settings.load_from_string('a.b.data_dir = "/some_dir"')
230
- assert c.get_path('a.b.data_dir') == '/some_dir'
231
-
232
- c = Settings.load_from_string('a.b.data_dir = "some_dir"')
233
- assert c.get_path('a.b.data_dir') == str(pathlib.Path.cwd() / 'some_dir')
234
-
235
- c = Settings.load_from_string('a.b.c.data_dir = "/some_dir"')
236
- assert c.a.b.get_path('c.data_dir') == '/some_dir'
237
-
238
-
239
- def test_default_type(tmp_path):
240
- s = '''
241
- [type1]
242
- a = 200
243
- b = 'h'
244
-
245
- [type2]
246
- c = [1, 2, 3]
247
- d.k = 2
248
-
249
- [type3]
250
- pa.b = 3
251
- pa.c = '2'
252
- v = [1]
253
-
254
- [type4]
255
- a = 0
256
-
257
- [type5]
258
- u = 0
259
- '''
260
-
261
- filename = tmp_path / 'default.toml'
262
- with open(str(filename), 'w') as f:
263
- f.write(s)
264
-
265
- s = '''
266
- [action1]
267
- type = 'type1'
268
- a = 100
269
-
270
- [action2]
271
- type = 'type1'
272
- c = 200
273
-
274
- [action3]
275
- type = 'type2'
276
- c = [2, 3, 1, 2]
277
- d = 4
278
-
279
- [type3]
280
- pa.a = 1
281
- pa.c = '5'
282
-
283
- [type5]
284
-
285
- '''
286
-
287
- filename2 = tmp_path / 'config.toml'
288
- with open(str(filename2), 'w') as f:
289
- f.write(s)
290
-
291
- with mock.patch.object(Settings, 'DEFAULT_SETTINGS', filename):
292
- c = Settings.load_with_defaults(filename2, with_type_keyword=True)
293
-
294
- assert c.action1.a == 100
295
- assert c.action1.b == 'h'
296
- assert c.action2.a == 200
297
- assert c.action2.c == 200
298
- assert c.action3.d == 4
299
- assert c.action3.c == [2, 3, 1, 2]
300
- assert c.type3.pa.a == 1
301
- assert c.type3.pa.b == 3
302
- assert c.type3.pa.c == '5'
303
- assert c.type3.v == [1]
304
- assert 'type4' not in c
305
- assert 'type5' in c
@@ -1,274 +0,0 @@
1
- import os
2
- import string
3
- import random
4
- import functools
5
-
6
- from unittest import mock
7
-
8
- from libpipe import worker
9
-
10
-
11
- def test_worker(tmp_path):
12
- w = worker.WorkerPool([worker.localhost_shortname, 'test_node'], 'Test', max_tasks_per_worker=5)
13
-
14
- assert len(w.workers) == 10
15
- assert isinstance(w.workers[0].client, worker.LocalClient)
16
- assert isinstance(w.workers[6].client, worker.Client)
17
-
18
- w = worker.WorkerPool([worker.localhost_shortname], 'Test')
19
-
20
- for i in range(20):
21
- w.add('touch %s/t%d' % (str(tmp_path), i))
22
-
23
- success, errors = w.execute()
24
-
25
- assert len(success) == 20
26
- assert len(errors) == 0
27
-
28
- for i in range(20):
29
- assert os.path.isfile('%s/t%d' % (str(tmp_path), i))
30
-
31
-
32
- def test_worker_log(tmp_path):
33
- w = worker.WorkerPool(['wrong_server', worker.localhost_shortname], 'Test')
34
-
35
- for i in range(20):
36
- w.add(f'echo Task {i}', output_file=str(tmp_path / f't{i}.log'))
37
- print(str(tmp_path / f't{i}.log'))
38
-
39
- success, errors = w.execute()
40
-
41
- assert len(success) == 20
42
- assert len(errors) == 0
43
-
44
- for i, task in enumerate(success):
45
- assert os.path.isfile(tmp_path / f't{i}.log')
46
- with open(tmp_path / f't{i}.log') as f:
47
- s = f.readlines()
48
- assert len(s) == 10 or len(s) == 8
49
- assert s[0].startswith('# Logging starting at ')
50
- assert s[1] == f"# Input command: bash -c 'echo Task {i}'\n"
51
- if len(s) == 8:
52
- assert s[2] == f'Starting on host {task.process.client.host} with PID {task.process.process.pid}\n'
53
- assert s[4] == f'Task {i}\n'
54
- if len(s) == 10:
55
- assert s[2].startswith('Error')
56
- assert s[3] == f'Will retry 1 / 3 ...\n'
57
- assert s[4] == f'Starting on host {task.process.client.host} with PID {task.process.process.pid}\n'
58
- assert s[6] == f'Task {i}\n'
59
- assert s[-1].startswith('# Logging stopped at ')
60
-
61
-
62
- def test_pidfile(tmp_path):
63
- w = worker.WorkerPool([worker.localhost_shortname], 'Test')
64
-
65
- for i in range(20):
66
- w.add(f'echo Task {i}; exit {i}', pid_filename=f't{i}.pid', keep_pid_file_on_error=True)
67
-
68
- success, errors = w.execute()
69
-
70
- assert len(success) == 20
71
- assert len(errors) == 0
72
-
73
- for task in success:
74
- assert os.path.exists(task.pid_filename) != (task.returncode == 0)
75
- if task.returncode > 0:
76
- with open(task.pid_filename) as fd:
77
- s = fd.read()
78
- assert s == f'{task.process.client.host} {task.process.process.pid}'
79
-
80
-
81
- def test_worker_wrong_server(tmp_path):
82
- w = worker.WorkerPool(['wrong_server'] * 6, 'Test')
83
-
84
- for i in range(20):
85
- w.add(f'echo Task {i}')
86
-
87
- success, errors = w.execute()
88
-
89
- assert len(success) == 0
90
- assert len(errors) == 20
91
-
92
-
93
- def test_worker_env_file(tmp_path):
94
- with open(tmp_path / 'env', 'w') as f:
95
- f.write('export WORKER_TEST=test_env')
96
-
97
- w = worker.WorkerPool([worker.localhost_shortname], 'Test', env_source_file=str(tmp_path / 'env'))
98
-
99
- os.listdir(tmp_path)
100
-
101
- for i in range(20):
102
- w.add('echo env:${WORKER_TEST}; touch %s/t_${WORKER_TEST}_%d' % (str(tmp_path), i))
103
-
104
- success, errors = w.execute()
105
-
106
- assert len(success) == 20
107
- assert len(errors) == 0
108
-
109
- for i in range(20):
110
- assert os.path.isfile('%s/t_test_env_%d' % (str(tmp_path), i))
111
-
112
-
113
- def test_worker_retcode():
114
- w = worker.WorkerPool([worker.localhost_shortname], 'Test')
115
-
116
- w.add('set -e; cat non_existing_file', name=f't0')
117
-
118
- success, errors = w.execute()
119
-
120
- assert len(success) == 1
121
- assert len(errors) == 0
122
- assert success[0].returncode == 1
123
-
124
- w = worker.WorkerPool([worker.localhost_shortname], 'Test')
125
-
126
- for i in range(10):
127
- w.add(f'echo Task {i}; exit {i}', name=f't{i}')
128
-
129
- success, errors = w.execute()
130
-
131
- print([k.process.client for k in success])
132
-
133
- assert len(success) == 10
134
- assert len(errors) == 0
135
- assert [f't{k.returncode}' for k in success] == [k.name for k in success]
136
-
137
-
138
- def test_timeout():
139
- w = worker.WorkerPool([worker.localhost_shortname], 'Test', max_time=2)
140
-
141
- w.add('echo 0')
142
-
143
- success, errors = w.execute()
144
-
145
- assert len(success) == 1
146
- assert len(errors) == 0
147
-
148
- w = worker.WorkerPool([worker.localhost_shortname], 'Test', max_time=1)
149
-
150
- w.add('sleep 100')
151
-
152
- success, errors = w.execute()
153
-
154
- # Currently falling. Need to be fixed
155
- assert len(success) == 0
156
- assert len(errors) == 1
157
-
158
-
159
- def test_done_callback():
160
- w = worker.WorkerPool([worker.localhost_shortname], 'Test', max_time=2)
161
-
162
- done = False
163
-
164
- def callback():
165
- nonlocal done
166
- done = True
167
-
168
- w.add('echo 0', done_callback=callback)
169
-
170
- success, errors = w.execute()
171
-
172
- assert len(success) == 1
173
- assert len(errors) == 0
174
- assert done
175
-
176
-
177
- def test_task_done_callback():
178
- w = worker.WorkerPool([worker.localhost_shortname], 'Test', max_time=2)
179
-
180
- done = 0
181
-
182
- def callback(task):
183
- print(task)
184
- nonlocal done
185
- assert isinstance(task, worker.Task)
186
- if task.name == 'Error':
187
- assert task.returncode == 1
188
- else:
189
- assert task.returncode == 0
190
- done += 1
191
-
192
- def callback_with_extra(name, task):
193
- assert name == 'test'
194
- callback(task)
195
-
196
- w.add('echo 0', task_done_callback=callback, name='Success1')
197
- w.add('echo 0', task_done_callback=callback, name='Success2')
198
- w.add('echo 3; cat non_existing_file', task_done_callback=functools.partial(callback_with_extra, 'test'), name='Error')
199
- w.add('echo 4;', task_done_callback=functools.partial(callback_with_extra, 'test_err'), name='CBK Error')
200
-
201
- success, errors = w.execute()
202
-
203
- assert len(success) == 4
204
- assert len(errors) == 0
205
- assert done == 3
206
-
207
-
208
- class FakeSSHClient(worker.LocalClient):
209
-
210
- def __init__(self, host, user=None, force_sync=True, password=None):
211
- worker.LocalClient.__init__(self, force_sync=force_sync)
212
- self.host = host
213
-
214
- async def execute(self, task):
215
- await worker.LocalClient.execute(self, task)
216
- task.remote_host = self.host
217
-
218
-
219
- def test_run_on_host():
220
- with mock.patch.object(worker, 'Client', FakeSSHClient):
221
- n_hosts = 10
222
- hosts = [f'host{k}' for k in range(n_hosts)]
223
- w = worker.WorkerPool(hosts, 'Test', max_time=2)
224
-
225
- for i in range(20 * n_hosts):
226
- w.add('echo 0', run_on_host=random.choice(hosts))
227
- w.add('echo 0', run_on_host='wrong_host')
228
- task_last = w.tasks[-1]
229
-
230
- success, errors = w.execute()
231
-
232
- assert len(errors) == 0
233
-
234
- for task in success:
235
- if task == task_last:
236
- assert task.run_on_host == None
237
- else:
238
- assert task.remote_host == task.run_on_host
239
-
240
-
241
- def test_line_too_long(tmp_path):
242
- w = worker.WorkerPool([worker.localhost_shortname], 'Test')
243
-
244
- N = 100000
245
- s = ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(N))
246
-
247
- log_file = str(tmp_path / f'text.log')
248
-
249
- w.add(f'echo Start; echo Other task; echo Done')
250
- w.add(f'echo Start; echo {s}; echo Done', output_file=log_file)
251
- w.add(f'echo Start; echo Other task; echo Done')
252
-
253
- success, errors = w.execute()
254
-
255
- has_long_line = False
256
- has_done = False
257
-
258
- with open(log_file) as f:
259
- for line in f.readlines():
260
- if '(line too long)' in line:
261
- has_long_line = True
262
- if 'Done' in line:
263
- has_done = True
264
-
265
- assert has_long_line
266
- assert has_done
267
- assert len(errors) == 0
268
- assert len(success) == 3
269
-
270
-
271
- def test_client():
272
- client = worker.Client('localhost')
273
-
274
- task = worker.Task('Check size', 'du -hs *')
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes