tinytaskmanager 1.0.0__py3-none-any.whl

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.
tinytaskmanager.py ADDED
@@ -0,0 +1,1311 @@
1
+ #!/usr/bin/env python3
2
+
3
+ # tinytaskmanager
4
+ # Tiny task manager for Linux, MacOS and Unix-like systems.
5
+ # Written as a single Python script.
6
+ # MIT License
7
+ # Copyright (c) 2022 Yuri Escalianti <yuriescl@gmail.com>
8
+ # Homepage: https://github.com/yuriescl/tinytaskmanager
9
+
10
+ from datetime import datetime
11
+ from fcntl import LOCK_EX, LOCK_UN, lockf
12
+ import io
13
+ from io import SEEK_END, SEEK_SET
14
+ import json
15
+ from multiprocessing.dummy import Pool as ThreadPool
16
+ import os
17
+ from os.path import abspath, join
18
+ from pathlib import Path
19
+ import re
20
+ import shlex
21
+ from shutil import get_terminal_size, rmtree
22
+ from signal import SIGINT, SIGKILL, SIGTERM, Signals, signal
23
+ from subprocess import DEVNULL, Popen, check_output
24
+ from sys import argv, exit, stderr, stdout, version_info
25
+ import tempfile
26
+ import time
27
+ from time import sleep
28
+ from typing import Dict, List, Optional, Tuple, Union
29
+
30
+ if version_info[0] < 3 or version_info[1] < 8:
31
+ raise Exception("Python >=3.8 is required to run this program")
32
+
33
+
34
+ LOCK_FILE_NAME = "lock"
35
+ CACHE_DIR = Path.home() / ".tinytaskmanager"
36
+ LOCK_PATH = Path(CACHE_DIR / LOCK_FILE_NAME)
37
+
38
+ RESERVED_FILE_NAMES = [LOCK_FILE_NAME]
39
+
40
+ VERSION = "0.14.0"
41
+ BUSY_LOOP_INTERVAL = 0.1 # seconds
42
+ TIMESTAMP_FMT = "%Y%m%d%H%M%S"
43
+
44
+ TERMINATE = False
45
+
46
+ Task = dict
47
+
48
+
49
+ class TtmException(Exception):
50
+ pass
51
+
52
+
53
+ class Tailer(object):
54
+ """
55
+ Code obtained from https://github.com/GreatFruitOmsk/tailhead/blob/master/tailhead/__init__.py
56
+ Copyright (c) 2012 Mike Thornton
57
+ Implements tailing and heading functionality like GNU tail and head
58
+ commands.
59
+ """
60
+
61
+ LINE_TERMINATORS = (b"\r\n", b"\n", b"\r")
62
+
63
+ def __init__(self, file, read_size=1024, end=False):
64
+ if not isinstance(file, io.IOBase) or isinstance(file, io.TextIOBase):
65
+ raise ValueError("io object must be in the binary mode")
66
+
67
+ self.read_size = read_size
68
+ self.file = file
69
+
70
+ if end:
71
+ self.file.seek(0, SEEK_END)
72
+
73
+ def splitlines(self, data):
74
+ return re.split(b"|".join(self.LINE_TERMINATORS), data)
75
+
76
+ def read(self, read_size=-1):
77
+ read_str = self.file.read(read_size)
78
+ return len(read_str), read_str
79
+
80
+ def prefix_line_terminator(self, data):
81
+ for t in self.LINE_TERMINATORS:
82
+ if data.startswith(t):
83
+ return t
84
+
85
+ return None
86
+
87
+ def suffix_line_terminator(self, data):
88
+ for t in self.LINE_TERMINATORS:
89
+ if data.endswith(t):
90
+ return t
91
+
92
+ return None
93
+
94
+ def seek_next_line(self):
95
+ where = self.file.tell()
96
+ offset = 0
97
+
98
+ while True:
99
+ data_len, data = self.read(self.read_size)
100
+ data_where = 0
101
+
102
+ if not data_len:
103
+ break
104
+
105
+ # Consider the following example: Foo\r | \nBar where " | " denotes current position,
106
+ # 'Foo\r' is the read part and '\nBar' is the remaining part.
107
+ # We should completely consume terminator "\r\n" by reading one extra byte.
108
+ if b"\r\n" in self.LINE_TERMINATORS and data[-1] == b"\r"[0]:
109
+ terminator_where = self.file.tell()
110
+ terminator_len, terminator_data = self.read(1)
111
+
112
+ if terminator_len and terminator_data[0] == b"\n"[0]:
113
+ data_len += 1
114
+ data += b"\n"
115
+ else:
116
+ self.file.seek(terminator_where)
117
+
118
+ while data_where < data_len:
119
+ terminator = self.prefix_line_terminator(data[data_where:])
120
+ if terminator:
121
+ self.file.seek(where + offset + data_where + len(terminator))
122
+ return self.file.tell()
123
+ else:
124
+ data_where += 1
125
+
126
+ offset += data_len
127
+ self.file.seek(where + offset)
128
+
129
+ return -1
130
+
131
+ def seek_previous_line(self):
132
+ where = self.file.tell()
133
+ offset = 0
134
+
135
+ while True:
136
+ if offset == where:
137
+ break
138
+
139
+ read_size = self.read_size if self.read_size <= where else where
140
+ self.file.seek(where - offset - read_size, SEEK_SET)
141
+ data_len, data = self.read(read_size)
142
+
143
+ # Consider the following example: Foo\r | \nBar where " | " denotes current position,
144
+ # '\nBar' is the read part and 'Foo\r' is the remaining part.
145
+ # We should completely consume terminator "\r\n" by reading one extra byte.
146
+ if b"\r\n" in self.LINE_TERMINATORS and data[0] == b"\n"[0]:
147
+ terminator_where = self.file.tell()
148
+ if terminator_where > data_len + 1:
149
+ self.file.seek(where - offset - data_len - 1, SEEK_SET)
150
+ terminator_len, terminator_data = self.read(1)
151
+
152
+ if terminator_data[0] == b"\r"[0]:
153
+ data_len += 1
154
+ data = b"\r" + data
155
+
156
+ self.file.seek(terminator_where)
157
+
158
+ data_where = data_len
159
+
160
+ while data_where > 0:
161
+ terminator = self.suffix_line_terminator(data[:data_where])
162
+ if terminator and offset == 0 and data_where == data_len:
163
+ # The last character is a line terminator that finishes current line. Ignore it.
164
+ data_where -= len(terminator)
165
+ elif terminator:
166
+ self.file.seek(where - offset - (data_len - data_where))
167
+ return self.file.tell()
168
+ else:
169
+ data_where -= 1
170
+
171
+ offset += data_len
172
+
173
+ if where == 0:
174
+ # Nothing more to read.
175
+ return -1
176
+ else:
177
+ # Very first line.
178
+ self.file.seek(0)
179
+ return 0
180
+
181
+ def tail(self, lines=10):
182
+ self.file.seek(0, SEEK_END)
183
+
184
+ for i in range(lines):
185
+ if self.seek_previous_line() == -1:
186
+ break
187
+
188
+ data = self.file.read()
189
+
190
+ for t in self.LINE_TERMINATORS:
191
+ if data.endswith(t):
192
+ # Only terminators _between_ lines should be preserved.
193
+ # Otherwise terminator of the last line will be treated as separtaing line and empty line.
194
+ data = data[: -len(t)]
195
+ break
196
+
197
+ if data:
198
+ return self.splitlines(data)
199
+ else:
200
+ return []
201
+
202
+ def head(self, lines=10):
203
+ if lines < 0:
204
+ self.file.seek(0, SEEK_END)
205
+ for i in range(-lines):
206
+ if self.seek_previous_line() == -1:
207
+ break
208
+ else:
209
+ self.file.seek(0)
210
+ for i in range(lines):
211
+ if self.seek_next_line() == -1:
212
+ break
213
+
214
+ end_pos = self.file.tell()
215
+
216
+ self.file.seek(0)
217
+ data = self.file.read(end_pos)
218
+
219
+ for t in self.LINE_TERMINATORS:
220
+ if data.endswith(t):
221
+ # Only terminators _between_ lines should be preserved.
222
+ # Otherwise terminator of the last line will be treated as separtaing line and empty line.
223
+ data = data[: -len(t)]
224
+ break
225
+
226
+ if data:
227
+ return self.splitlines(data)
228
+ else:
229
+ return []
230
+
231
+ def follow(self):
232
+ trailing = True
233
+
234
+ while True:
235
+ where = self.file.tell()
236
+
237
+ if where > os.fstat(self.file.fileno()).st_size:
238
+ # File was truncated.
239
+ where = 0
240
+ self.file.seek(where)
241
+
242
+ line = self.file.readline()
243
+
244
+ if line:
245
+ if trailing and line in self.LINE_TERMINATORS:
246
+ # This is just the line terminator added to the end of the file
247
+ # before a new line, ignore.
248
+ trailing = False
249
+ continue
250
+
251
+ terminator = self.suffix_line_terminator(line)
252
+ if terminator:
253
+ line = line[: -len(terminator)]
254
+
255
+ trailing = False
256
+ yield line
257
+ else:
258
+ trailing = True
259
+ self.file.seek(where)
260
+ yield None
261
+
262
+
263
+ class AtomicOpen:
264
+ """https://stackoverflow.com/a/46407326/3705710"""
265
+
266
+ def __init__(self, path, *args, noop=False, **kwargs):
267
+ if noop is False:
268
+ self.file = open(path, *args, **kwargs)
269
+ self.lock_file(self.file)
270
+ self.noop = noop
271
+
272
+ @staticmethod
273
+ def lock_file(f):
274
+ if f.writable():
275
+ lockf(f, LOCK_EX)
276
+
277
+ @staticmethod
278
+ def unlock_file(f):
279
+ if f.writable():
280
+ lockf(f, LOCK_UN)
281
+
282
+ def __enter__(self, *args, **kwargs):
283
+ if self.noop is False:
284
+ return self.file
285
+
286
+ def __exit__(self, exc_type=None, exc_value=None, traceback=None):
287
+ if self.noop is False:
288
+ self.file.flush()
289
+ os.fsync(self.file.fileno())
290
+ self.unlock_file(self.file)
291
+ self.file.close()
292
+ if exc_type is not None:
293
+ return False
294
+ else:
295
+ return True
296
+
297
+
298
+ class bcolors:
299
+ """https://stackoverflow.com/a/287944/3705710"""
300
+
301
+ HEADER = "\033[95m"
302
+ OKBLUE = "\033[94m"
303
+ OKCYAN = "\033[96m"
304
+ OKGREEN = "\033[92m"
305
+ WARNING = "\033[93m"
306
+ LIGHTGREY = "\033[0;37m"
307
+ FAIL = "\033[91m"
308
+ ENDC = "\033[0m"
309
+ BOLD = "\033[1m"
310
+ UNDERLINE = "\033[4m"
311
+
312
+
313
+ ##############
314
+ # ARG PARSING
315
+
316
+
317
+ def arg_requires_value(arg: str, option: Optional[str] = None) -> bool:
318
+ def dashes(a: str):
319
+ return "-" if len(a) == 1 else "--"
320
+
321
+ if arg in ["cache-dir"]:
322
+ return True
323
+ if arg in ["h", "help"]:
324
+ return False
325
+ if option is None:
326
+ if arg in ["version"]:
327
+ return False
328
+ elif option == "run":
329
+ if arg in ["s", "shell", "split-output"]:
330
+ return False
331
+ if arg in ["n", "name"]:
332
+ return True
333
+ elif option == "start":
334
+ pass
335
+ elif option == "stop":
336
+ if arg in ["k", "kill"] + signals_list():
337
+ return False
338
+ elif option == "rm":
339
+ if arg in ["a", "all"]:
340
+ return False
341
+ elif option == "ls":
342
+ if arg in ["a", "all"]:
343
+ return False
344
+ elif option == "logs":
345
+ if arg in ["f", "follow", "head"]:
346
+ return False
347
+ raise TtmException(f"Unrecognized argument {dashes(arg)}{arg}")
348
+
349
+
350
+ def is_value_next(args: List[str], pos: int) -> bool:
351
+ return pos + 1 < len(args) and not args[pos + 1].startswith("-")
352
+
353
+
354
+ def parse_args(
355
+ args_to_parse: List[str],
356
+ ) -> Tuple[Dict, Optional[str], Dict, Optional[List[str]]]:
357
+ args = args_to_parse[1:]
358
+ global_args: Dict[str, Union[str, bool]] = {}
359
+ option = None
360
+ pos = 0
361
+ while True:
362
+ if pos >= len(args):
363
+ break
364
+ current_arg = args[pos]
365
+ if current_arg in ["run", "start", "stop", "rm", "ls", "logs"]:
366
+ option = current_arg
367
+ pos += 1
368
+ break
369
+ elif current_arg.startswith("--"):
370
+ current_arg = current_arg[2:]
371
+ if arg_requires_value(current_arg, option):
372
+ if not is_value_next(args, pos):
373
+ raise TtmException(f"Argument --{current_arg} requires a value")
374
+ global_args[current_arg] = args[pos + 1]
375
+ pos += 2
376
+ continue
377
+ else:
378
+ global_args[current_arg] = True
379
+ pos += 1
380
+ continue
381
+ elif current_arg.startswith("-"):
382
+ current_arg = current_arg[1:]
383
+ if len(current_arg) == 1:
384
+ if arg_requires_value(current_arg, option):
385
+ if not is_value_next(args, pos):
386
+ raise TtmException(f"Argument -{current_arg} requires a value")
387
+ global_args[current_arg] = args[pos + 1]
388
+ pos += 2
389
+ continue
390
+ else:
391
+ global_args[current_arg] = True
392
+ pos += 1
393
+ continue
394
+ else:
395
+ for letter in current_arg:
396
+ if arg_requires_value(letter, option):
397
+ raise TtmException(
398
+ f"Argument -{letter} cannot be grouped with other arguments"
399
+ )
400
+ global_args[letter] = True
401
+ pos += 1
402
+ continue
403
+ else:
404
+ raise TtmException(f"Unrecognized option {current_arg}")
405
+ pos += 1
406
+
407
+ option_args: Dict[str, Union[str, bool]] = {}
408
+ command = None
409
+
410
+ if option is not None:
411
+ if pos >= len(args) and option not in ["ls"]:
412
+ raise TtmException(f"Missing arguments for option '{option}'")
413
+ while True:
414
+ if pos >= len(args):
415
+ break
416
+ current_arg = args[pos]
417
+ if current_arg.startswith("--"):
418
+ current_arg = current_arg[2:]
419
+ if arg_requires_value(current_arg, option):
420
+ if not is_value_next(args, pos):
421
+ raise TtmException(f"Argument --{current_arg} requires a value")
422
+ option_args[current_arg] = args[pos + 1]
423
+ pos += 2
424
+ continue
425
+ else:
426
+ option_args[current_arg] = True
427
+ pos += 1
428
+ continue
429
+ elif current_arg.startswith("-"):
430
+ current_arg = current_arg[1:]
431
+ if len(current_arg) == 1:
432
+ if arg_requires_value(current_arg, option):
433
+ if not is_value_next(args, pos):
434
+ raise TtmException(
435
+ f"Argument -{current_arg} requires a value"
436
+ )
437
+ option_args[current_arg] = args[pos + 1]
438
+ pos += 2
439
+ continue
440
+ else:
441
+ option_args[current_arg] = True
442
+ pos += 1
443
+ continue
444
+ else:
445
+ for letter in current_arg:
446
+ if arg_requires_value(letter, option) and not is_value_next(
447
+ args, pos
448
+ ):
449
+ raise TtmException(
450
+ f"Argument -{letter} cannot be grouped with other arguments"
451
+ )
452
+ option_args[letter] = True
453
+ pos += 1
454
+ continue
455
+ else:
456
+ command = args[pos:]
457
+ break
458
+
459
+ return global_args, option, option_args, command
460
+
461
+
462
+ ##################
463
+ # FILE OPERATIONS
464
+
465
+
466
+ def init_cache_dir(cache_dir: Optional[Union[str, Path]]):
467
+ global CACHE_DIR
468
+ global LOCK_PATH
469
+ if cache_dir is not None:
470
+ CACHE_DIR = Path(cache_dir)
471
+ os.makedirs(CACHE_DIR, exist_ok=True)
472
+ LOCK_FILE_NAME = "lock"
473
+ LOCK_PATH = Path(CACHE_DIR / LOCK_FILE_NAME)
474
+ LOCK_PATH.touch(exist_ok=True)
475
+
476
+
477
+ def get_task_label(task: Task):
478
+ if task["name"] is not None:
479
+ return f"{task['name']}-{task['id']}"
480
+ else:
481
+ return task["id"]
482
+
483
+
484
+ def parse_task_id_or_name(task_name_or_id: str) -> Tuple[Optional[str], Optional[str]]:
485
+ try:
486
+ task_id = str(int(task_name_or_id))
487
+ name = None
488
+ except (ValueError, TypeError):
489
+ task_id = None
490
+ name = task_name_or_id
491
+ return task_id, name
492
+
493
+
494
+ def create_pidfile(task: Task):
495
+ if task.get("pid") is not None:
496
+ if task.get("pidfile"):
497
+ Path(task["pidfile"]).unlink(missing_ok=True)
498
+ fh, task["pidfile"] = tempfile.mkstemp()
499
+ with os.fdopen(fh, "w") as f:
500
+ f.write(task["pid"])
501
+
502
+
503
+ def create_task_cache(task: Task, split_output=False) -> Task:
504
+ dir_name = get_task_label(task)
505
+ dir_path = CACHE_DIR / dir_name
506
+ os.makedirs(dir_path, exist_ok=True)
507
+ filepath = dir_path / "task.json"
508
+ timestamp = datetime.now().strftime(TIMESTAMP_FMT)
509
+ if split_output:
510
+ stdout_path = dir_path / f"{dir_name}-{timestamp}.out"
511
+ stderr_path = dir_path / f"{dir_name}-{timestamp}.err"
512
+ task.update(
513
+ {
514
+ "stdout": str(stdout_path),
515
+ "stderr": str(stderr_path),
516
+ "started_at": timestamp,
517
+ }
518
+ )
519
+ else:
520
+ logs_path = dir_path / f"{dir_name}-{timestamp}.log"
521
+ task.update(
522
+ {
523
+ "logs": str(logs_path),
524
+ "started_at": timestamp,
525
+ }
526
+ )
527
+ create_pidfile(task)
528
+ with open(filepath, "w") as f:
529
+ task_to_dump = dict(task)
530
+ task_to_dump.pop("pid", None)
531
+ json.dump(task_to_dump, f)
532
+ return task
533
+
534
+
535
+ def update_task_cache(task: Task):
536
+ dir_name = get_task_label(task)
537
+ dir_path = CACHE_DIR / dir_name
538
+ filepath = dir_path / "task.json"
539
+ create_pidfile(task)
540
+ with open(filepath, "w") as f:
541
+ task_to_dump = dict(task)
542
+ task_to_dump.pop("pid", None)
543
+ json.dump(task_to_dump, f)
544
+
545
+
546
+ def is_task_running(task: Task) -> bool:
547
+ output = check_output(["ps", "-ax", "-o", "pid,args"], start_new_session=True)
548
+ for line in output.splitlines():
549
+ decoded = line.decode().strip()
550
+ ps_pid, cmdline = decoded.split(" ", 1)
551
+ if task.get("pid") is not None and ps_pid == task["pid"]:
552
+ return True
553
+ return False
554
+
555
+
556
+ def get_task_from_cache_file(cache_file_path: str):
557
+ with open(cache_file_path) as f:
558
+ task = json.load(f)
559
+ if task.get("pidfile") and Path(task["pidfile"]).exists():
560
+ with open(task["pidfile"], "r") as f:
561
+ task["pid"] = f.read()
562
+ return task
563
+
564
+
565
+ def find_task_by_name(name: str) -> Optional[Task]:
566
+ for filename in os.listdir(CACHE_DIR):
567
+ if filename not in RESERVED_FILE_NAMES and filename.split("-")[0] == name:
568
+ path = abspath(join(CACHE_DIR, filename, "task.json"))
569
+ return get_task_from_cache_file(path)
570
+ return None
571
+
572
+
573
+ def find_task_by_id(task_id: str) -> Optional[Dict]:
574
+ for filename in os.listdir(CACHE_DIR):
575
+ if filename in RESERVED_FILE_NAMES:
576
+ continue
577
+ try:
578
+ filename_split = filename.split("-")
579
+ if len(filename_split) == 1:
580
+ filename_task_id = filename_split[0]
581
+ else:
582
+ filename_task_id = filename_split[1]
583
+ if filename_task_id == task_id:
584
+ path = abspath(join(CACHE_DIR, filename, "task.json"))
585
+ return get_task_from_cache_file(path)
586
+ except IndexError:
587
+ pass
588
+ return None
589
+
590
+
591
+ def delete_pidfile(task: Task):
592
+ if task.get("pidfile"):
593
+ Path(task["pidfile"]).unlink(missing_ok=True)
594
+
595
+
596
+ def remove_task_by_name(name: str):
597
+ with AtomicOpen(LOCK_PATH):
598
+ for filename in os.listdir(CACHE_DIR):
599
+ filename_split = filename.split("-")
600
+ if len(filename_split) == 1:
601
+ continue
602
+ else:
603
+ filename_task_name = filename_split[0]
604
+
605
+ if filename_task_name == name:
606
+ task = find_task_by_name(name)
607
+ if task is None:
608
+ raise TtmException("Failed to find task by name")
609
+ if is_task_running(task):
610
+ raise TtmException(
611
+ "Cannot remove task while it's running.\n"
612
+ "To stop it, run:\n"
613
+ f"tinytaskmanager stop {name}"
614
+ )
615
+ dir_path = abspath(join(CACHE_DIR, filename))
616
+ rmtree(dir_path)
617
+ delete_pidfile(task)
618
+ return
619
+ raise TtmException(f"No task with name {name}")
620
+
621
+
622
+ def remove_task_by_id(task_id: str):
623
+ with AtomicOpen(LOCK_PATH):
624
+ for filename in os.listdir(CACHE_DIR):
625
+ try:
626
+ filename_split = filename.split("-")
627
+ if len(filename_split) == 1:
628
+ filename_task_id = filename_split[0]
629
+ else:
630
+ filename_task_id = filename_split[1]
631
+
632
+ if filename_task_id == task_id:
633
+ task = find_task_by_id(task_id)
634
+ if task is None:
635
+ raise TtmException("Failed to find task by id")
636
+ if is_task_running(task):
637
+ raise TtmException(
638
+ "Cannot remove task while it's running.\n"
639
+ "To stop it, run:\n"
640
+ f"tinytaskmanager stop {task_id}"
641
+ )
642
+ dir_path = abspath(join(CACHE_DIR, filename))
643
+ rmtree(dir_path)
644
+ delete_pidfile(task)
645
+ return
646
+ except IndexError:
647
+ pass
648
+ raise TtmException(f"No task with ID {task_id}")
649
+
650
+
651
+ def generate_id():
652
+ existing_ids = []
653
+ for filename in os.listdir(CACHE_DIR):
654
+ try:
655
+ existing_ids.append(str(int(filename)))
656
+ except ValueError:
657
+ try:
658
+ existing_ids.append(str(int(filename.split("-")[1])))
659
+ except (ValueError, IndexError):
660
+ pass
661
+ for i in range(1, 10000):
662
+ str_i = str(i)
663
+ if str_i not in existing_ids:
664
+ return str_i
665
+ raise TtmException("Failed to generated task ID")
666
+
667
+
668
+ def get_child_pids(parent_pid: int):
669
+ output = check_output(["ps", "-ax", "-o", "pid,ppid"], start_new_session=True)
670
+ ppid = str(parent_pid)
671
+ child_pids = []
672
+ for line in output.splitlines():
673
+ decoded = line.decode().strip()
674
+ ps_child_pid, ps_ppid = decoded.split(None, 1)
675
+ if ps_ppid == ppid:
676
+ child_pids.append(int(ps_child_pid))
677
+ return child_pids
678
+
679
+
680
+ def kill_recursively(pid: int, sig: int):
681
+ for child_pid in get_child_pids(pid):
682
+ kill_recursively(child_pid, sig)
683
+ os.kill(pid, sig)
684
+
685
+
686
+ #############
687
+ # OPERATIONS
688
+
689
+
690
+ def run(
691
+ command: List[str],
692
+ name: Optional[str] = None,
693
+ split_output=False,
694
+ shell=False,
695
+ ) -> Task:
696
+ with AtomicOpen(LOCK_PATH):
697
+ if name is not None:
698
+ task = find_task_by_name(name)
699
+ if task:
700
+ if is_task_running(task):
701
+ raise TtmException(
702
+ f"Task {name} is already running with PID {task['pid']}"
703
+ )
704
+ raise TtmException(
705
+ f"Task {name} already exists and it's not running.\n"
706
+ "To remove it, run:\n"
707
+ f"tinytaskmanager rm {name}"
708
+ )
709
+ task = {
710
+ "id": generate_id(),
711
+ "name": name,
712
+ "cwd": os.getcwd(),
713
+ "command": command,
714
+ "shell": shell,
715
+ }
716
+ if split_output:
717
+ task = create_task_cache(task, split_output=split_output)
718
+ stdout_path = task["stdout"]
719
+ stderr_path = task["stderr"]
720
+ logs_path = ""
721
+ else:
722
+ task = create_task_cache(task, split_output=split_output)
723
+ stdout_path = ""
724
+ stderr_path = ""
725
+ logs_path = task["logs"]
726
+ if split_output:
727
+ with open(stdout_path, "wb") as out:
728
+ with open(stderr_path, "wb") as err:
729
+ proc = Popen(
730
+ build_cmd(command, shell),
731
+ shell=shell,
732
+ cwd=task["cwd"],
733
+ stdin=DEVNULL,
734
+ stdout=out,
735
+ stderr=err,
736
+ start_new_session=True,
737
+ )
738
+ else:
739
+ with open(logs_path, "wb") as output:
740
+ proc = Popen(
741
+ build_cmd(command, shell),
742
+ shell=shell,
743
+ cwd=task["cwd"],
744
+ stdin=DEVNULL,
745
+ stdout=output,
746
+ stderr=output,
747
+ start_new_session=True,
748
+ )
749
+ task["pid"] = str(proc.pid)
750
+ update_task_cache(task)
751
+ return task
752
+
753
+
754
+ def start_task(task_id: Optional[str] = None, name: Optional[str] = None):
755
+ with AtomicOpen(LOCK_PATH):
756
+ if name is not None:
757
+ task = find_task_by_name(name)
758
+ if task is not None:
759
+ if is_task_running(task):
760
+ raise TtmException(
761
+ f"Task {name} is already running with PID {task['pid']}"
762
+ )
763
+ else:
764
+ raise TtmException(f"No task with name {name}")
765
+ elif task_id is not None:
766
+ task = find_task_by_id(task_id)
767
+ if task is not None:
768
+ if is_task_running(task):
769
+ raise TtmException(
770
+ f"Task with ID {task_id} is already running with PID {task['pid']}"
771
+ )
772
+ else:
773
+ raise TtmException(f"No task with ID {task_id}")
774
+ else:
775
+ raise ValueError("Either task_id or name must be set")
776
+
777
+ if task["name"] is not None:
778
+ dir_name = f"{task['name']}-{task['id']}"
779
+ else:
780
+ dir_name = task["id"]
781
+ dir_path = CACHE_DIR / dir_name
782
+ timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
783
+ if task.get("stdout") is not None:
784
+ task["stdout"] = str(dir_path / f"{dir_name}-{timestamp}.out")
785
+ task["stderr"] = str(dir_path / f"{dir_name}-{timestamp}.err")
786
+ with open(task["stdout"], "wb") as out:
787
+ with open(task["stderr"], "wb") as err:
788
+ proc = Popen(
789
+ build_cmd(task["command"], task["shell"]),
790
+ shell=task["shell"],
791
+ cwd=task["cwd"],
792
+ stdin=DEVNULL,
793
+ stdout=out,
794
+ stderr=err,
795
+ start_new_session=True,
796
+ )
797
+ else:
798
+ task["logs"] = str(dir_path / f"{dir_name}-{timestamp}.log")
799
+ with open(task["logs"], "wb") as output:
800
+ proc = Popen(
801
+ build_cmd(task["command"], task["shell"]),
802
+ shell=task["shell"],
803
+ cwd=task["cwd"],
804
+ stdin=DEVNULL,
805
+ stdout=output,
806
+ stderr=output,
807
+ start_new_session=True,
808
+ )
809
+ task["pid"] = str(proc.pid)
810
+ task["started_at"] = timestamp
811
+ update_task_cache(task)
812
+
813
+
814
+ def stop_task(
815
+ task_id: Optional[str] = None, name: Optional[str] = None, sig: int = SIGTERM
816
+ ):
817
+ with AtomicOpen(LOCK_PATH):
818
+ if name is not None:
819
+ task = find_task_by_name(name)
820
+ if task is not None:
821
+ if not is_task_running(task):
822
+ raise TtmException(f"Task {name} is not running")
823
+ else:
824
+ raise TtmException(f"No task with name {name}")
825
+ elif task_id is not None:
826
+ task = find_task_by_id(task_id)
827
+ if task is not None:
828
+ if not is_task_running(task):
829
+ raise TtmException(f"Task with ID {task_id} is not running")
830
+ else:
831
+ raise TtmException(f"No task with ID {task_id}")
832
+ else:
833
+ raise ValueError("Either task_id or name must be set")
834
+
835
+ # We kill and busy wait outside the above file lock for better parallel performance
836
+ kill_recursively(int(task["pid"]), sig)
837
+ while True:
838
+ # TODO add timeout
839
+ with AtomicOpen(LOCK_PATH):
840
+ if not is_task_running(task):
841
+ delete_pidfile(task)
842
+ break
843
+ sleep(BUSY_LOOP_INTERVAL)
844
+
845
+
846
+ def remove_all_tasks():
847
+ with AtomicOpen(LOCK_PATH):
848
+ for filename in os.listdir(CACHE_DIR):
849
+ if TERMINATE:
850
+ return
851
+ if filename in RESERVED_FILE_NAMES:
852
+ continue
853
+ path = abspath(join(CACHE_DIR, filename, "task.json"))
854
+ try:
855
+ task = get_task_from_cache_file(path)
856
+ if is_task_running(task):
857
+ print_error(f"Task {task['id']}: cannot remove while it's running")
858
+ else:
859
+ dir_path = abspath(join(CACHE_DIR, filename))
860
+ rmtree(dir_path)
861
+ except (NotADirectoryError, FileNotFoundError, ValueError):
862
+ pass
863
+
864
+
865
+ def rm(task_name_or_id: Optional[str], rm_all=False) -> bool:
866
+ try:
867
+ if rm_all:
868
+ remove_all_tasks()
869
+ else:
870
+ if task_name_or_id is None:
871
+ raise ValueError("task_name_or_id is None")
872
+ task_id, name = parse_task_id_or_name(task_name_or_id)
873
+ if task_id is not None:
874
+ remove_task_by_id(task_id)
875
+ elif name is not None:
876
+ remove_task_by_name(name)
877
+ except TtmException as e:
878
+ print_error(str(e))
879
+ return False
880
+ return True
881
+
882
+
883
+ def logs(task_name_or_id: str, follow=False, head=False):
884
+ def print_lines(lines):
885
+ if isinstance(lines, list):
886
+ for line in lines:
887
+ stdout.buffer.write(line)
888
+ stdout.buffer.write("\n".encode())
889
+ elif isinstance(lines, bytes):
890
+ stdout.buffer.write(lines)
891
+ stdout.buffer.write("\n".encode())
892
+ stdout.buffer.flush()
893
+
894
+ if follow and head:
895
+ raise TtmException("--follow and --head cannot be used together")
896
+
897
+ task_id, name = parse_task_id_or_name(task_name_or_id)
898
+
899
+ if task_id is not None:
900
+ task = find_task_by_id(task_id)
901
+ if task is None:
902
+ raise TtmException(f"No task with ID {task_id}")
903
+ elif name is not None:
904
+ task = find_task_by_name(name)
905
+ if task is None:
906
+ raise TtmException(f"No task with name {name}")
907
+ else:
908
+ raise ValueError("task_id and name are None")
909
+
910
+ logs_path = task.get("logs")
911
+ if logs_path is None:
912
+ raise TtmException(
913
+ "Task was created using --split-output, use 'stdout' or 'stderr' instead of 'logs'"
914
+ )
915
+
916
+ line_count = 15
917
+
918
+ if not head and follow:
919
+ with open(logs_path, "rb") as file:
920
+ print_grey(f"{logs_path} last {line_count} lines:")
921
+ print_lines(Tailer(file).tail(lines=line_count))
922
+
923
+ with open(logs_path, "rb") as file:
924
+ if head:
925
+ print_grey(f"{logs_path} first {line_count} lines:")
926
+ print_lines(Tailer(file).head(lines=line_count))
927
+ elif follow:
928
+ print_grey(f"{logs_path} followed tail:")
929
+ for line in Tailer(file, end=True).follow():
930
+ if line is None:
931
+ time.sleep(0.01)
932
+ continue
933
+ print_lines(line)
934
+ else:
935
+ print_grey(f"{logs_path} last {line_count} lines:")
936
+ print_lines(Tailer(file).tail(lines=line_count))
937
+
938
+ print_lines([])
939
+
940
+
941
+ def start(task_name_or_id: str) -> bool:
942
+ task_id, name = parse_task_id_or_name(task_name_or_id)
943
+
944
+ try:
945
+ start_task(task_id=task_id, name=name)
946
+ return True
947
+ except TtmException as e:
948
+ print_error(str(e))
949
+ return False
950
+
951
+
952
+ def stop(task_name_or_id: str, sig: int):
953
+ task_id, name = parse_task_id_or_name(task_name_or_id)
954
+
955
+ try:
956
+ stop_task(task_id=task_id, name=name, sig=sig)
957
+ return True
958
+ except TtmException as e:
959
+ print_error(str(e))
960
+ return False
961
+
962
+
963
+ def ls(ls_all=False, command: Optional[List[str]] = None):
964
+ tasks = []
965
+ with AtomicOpen(LOCK_PATH):
966
+ for filename in os.listdir(CACHE_DIR):
967
+ if filename in RESERVED_FILE_NAMES:
968
+ continue
969
+ path = abspath(join(CACHE_DIR, filename, "task.json"))
970
+ force_list = False
971
+ if command:
972
+ for task_name_or_id in command:
973
+ task_id, name = parse_task_id_or_name(task_name_or_id)
974
+ filename_split = filename.split("-")
975
+ if task_id in filename_split or name in filename_split:
976
+ force_list = True
977
+ if not force_list:
978
+ continue
979
+ try:
980
+ task = get_task_from_cache_file(path)
981
+ task["started_at"] = datetime.strptime(
982
+ task["started_at"], TIMESTAMP_FMT
983
+ )
984
+ if is_task_running(task):
985
+ diff = datetime.now() - task["started_at"]
986
+ task["uptime"] = format_seconds(int(diff.total_seconds()))
987
+ tasks.append(task)
988
+ elif ls_all or force_list:
989
+ task["pid"] = "-"
990
+ task["uptime"] = "-"
991
+ tasks.append(task)
992
+ except (NotADirectoryError, FileNotFoundError, ValueError):
993
+ pass
994
+
995
+ name_len_max = 4
996
+ for task in tasks:
997
+ if task["name"] is not None and len(task["name"]) > name_len_max:
998
+ name_len_max = len(task["name"])
999
+
1000
+ tasks = sorted(tasks, key=lambda d: d["started_at"])
1001
+
1002
+ columns = get_terminal_size((80, 20)).columns
1003
+ name_size = min(name_len_max, 16)
1004
+ command_size = columns - 21 - name_size
1005
+ template = r"{0:4} {1:NAME_SIZE} {2:6} {3:7} {4:COMMAND_SIZE}"
1006
+ template = template.replace("NAME_SIZE", str(name_size))
1007
+ template = template.replace("COMMAND_SIZE", str(command_size))
1008
+ print(template.format("ID", "NAME", "UPTIME", "PID", "COMMAND"))
1009
+ for task in tasks:
1010
+ print(
1011
+ template.format(
1012
+ task["id"],
1013
+ task["name"] or "-",
1014
+ task["uptime"],
1015
+ task["pid"],
1016
+ shlex.join(task["command"]),
1017
+ )
1018
+ )
1019
+
1020
+
1021
+ #######
1022
+ # MISC
1023
+
1024
+
1025
+ def build_cmd(command: List[str], shell: bool):
1026
+ if shell:
1027
+ if len(command) == 1:
1028
+ return command[0]
1029
+ else:
1030
+ return shlex.join(command)
1031
+ else:
1032
+ return command
1033
+
1034
+
1035
+ def format_seconds(seconds, long=False):
1036
+ if long:
1037
+ raise NotImplementedError()
1038
+ else:
1039
+ days, remainder = divmod(seconds, 86400)
1040
+ hours, remainder = divmod(remainder, 3600)
1041
+ minutes, seconds = divmod(remainder, 60)
1042
+ s = ""
1043
+ if days > 0:
1044
+ s += f"{days}d"
1045
+ elif hours > 0:
1046
+ s += f"{hours}h"
1047
+ elif minutes > 0:
1048
+ s += f"{minutes}m"
1049
+ else:
1050
+ s += f"{seconds}s"
1051
+ return s
1052
+
1053
+
1054
+ def signal_handler(signum, frame):
1055
+ global TERMINATE
1056
+ if signum in [SIGINT, SIGTERM]:
1057
+ TERMINATE = True
1058
+
1059
+
1060
+ def signals_list():
1061
+ return [str(sig.value) for sig in Signals]
1062
+
1063
+
1064
+ def print_error(msg: str, *args, **kwargs):
1065
+ if "file" not in kwargs:
1066
+ kwargs["file"] = stderr
1067
+ print(f"{bcolors.FAIL}{msg}{bcolors.ENDC}", *args, **kwargs)
1068
+
1069
+
1070
+ def print_grey(msg: str, *args, **kwargs):
1071
+ print(f"{bcolors.LIGHTGREY}{msg}{bcolors.ENDC}", *args, **kwargs)
1072
+
1073
+
1074
+ def print_warning(msg: str, *args, **kwargs):
1075
+ if "file" not in kwargs:
1076
+ kwargs["file"] = stderr
1077
+ print(f"{bcolors.WARNING}{msg}{bcolors.ENDC}", *args, **kwargs)
1078
+
1079
+
1080
+ def print_success(msg: str, *args, **kwargs):
1081
+ print(f"{bcolors.OKGREEN}{msg}{bcolors.ENDC}", *args, **kwargs)
1082
+
1083
+
1084
+ def print_help():
1085
+ print()
1086
+ print("Usage: tinytaskmanager [OPTIONS] COMMAND")
1087
+ print()
1088
+ print("Tiny task manager for Linux, MacOS and Unix-like systems")
1089
+ print()
1090
+ print("Options:")
1091
+
1092
+ print(
1093
+ f" --cache-dir Use a custom cache directory instead of the default {CACHE_DIR}"
1094
+ )
1095
+ print(" -h, --help Display help")
1096
+ print(" --version Display program version")
1097
+ print()
1098
+ print("Commands:")
1099
+ print(" logs Display logs of a task")
1100
+ print(" ls List tasks")
1101
+ print(" rm Remove tasks")
1102
+ print(" run Run a new task")
1103
+ print(" start Start tasks")
1104
+ print(" stop Stop tasks")
1105
+
1106
+
1107
+ def print_help_logs():
1108
+ print()
1109
+ print("Usage: tinytaskmanager logs TASK")
1110
+ print()
1111
+ print("Display logs of a task.")
1112
+ print("TASK can be a task ID or a task name.")
1113
+ print()
1114
+
1115
+
1116
+ def print_help_ls():
1117
+ print()
1118
+ print("Usage: tinytaskmanager ls [OPTIONS]")
1119
+ print()
1120
+ print("List tasks")
1121
+ print()
1122
+ print("Options:")
1123
+ print(" -a, --all List all tasks, including stopped")
1124
+ print()
1125
+
1126
+
1127
+ def print_help_rm():
1128
+ print()
1129
+ print("Usage: tinytaskmanager rm [OPTIONS] [TASK]")
1130
+ print()
1131
+ print("Remove tasks.")
1132
+ print("TASK can be a task ID or a task name.")
1133
+ print()
1134
+ print("Options:")
1135
+ print(" -a, --all Remove all tasks")
1136
+ print()
1137
+
1138
+
1139
+ def print_help_run():
1140
+ print()
1141
+ print("Usage: tinytaskmanager run [OPTIONS] COMMAND")
1142
+ print()
1143
+ print("Run a new task")
1144
+ print()
1145
+ print("Options:")
1146
+ print(" -s, --shell Run COMMAND in a shell")
1147
+ print()
1148
+ print("Examples:")
1149
+ print(" tinytaskmanager run /path/to/my/program arg1 arg2")
1150
+ print(" tinytaskmanager run -s 'echo test >> myfile'")
1151
+ print()
1152
+
1153
+
1154
+ def print_help_start():
1155
+ print()
1156
+ print("Usage: tinytaskmanager start TASK")
1157
+ print()
1158
+ print("Start a task")
1159
+ print()
1160
+ print("Examples:")
1161
+ print(" tinytaskmanager start 3")
1162
+ print(" tinytaskmanager start mytaskname")
1163
+ print()
1164
+
1165
+
1166
+ def print_help_stop():
1167
+ print()
1168
+ print("Usage: tinytaskmanager stop TASK")
1169
+ print()
1170
+ print("Stop a task")
1171
+ print()
1172
+ print("Options:")
1173
+ print(" -k, --kill Send a SIGKILL signal instead of the default SIGTERM.")
1174
+ print(" Equivalent to -9.")
1175
+ print(" -SIG Send a specific signal instead of the default SIGTERM")
1176
+ print()
1177
+ print("Examples:")
1178
+ print(" tinytaskmanager stop 3")
1179
+ print(" tinytaskmanager stop mytaskname")
1180
+ print(" tinytaskmanager stop --kill my_hanged_task")
1181
+ print()
1182
+ print(" # Send a SIGINT signal to a task")
1183
+ print(" tinytaskmanager stop -2 my_interruptable_task")
1184
+ print()
1185
+ print(" # Send a SIGTERM signal to a task")
1186
+ print(" tinytaskmanager stop -9 my_hanged_task")
1187
+ print()
1188
+
1189
+
1190
+ def main():
1191
+ for sig in [SIGINT, SIGTERM]:
1192
+ signal(sig, signal_handler)
1193
+
1194
+ try:
1195
+ if len(argv) == 1:
1196
+ print_error("No option provided. Use -h for help.")
1197
+ exit(1)
1198
+ global_args, option, option_args, command = parse_args(argv)
1199
+
1200
+ init_cache_dir(global_args.get("cache-dir"))
1201
+
1202
+ if option is None:
1203
+ if global_args.get("version"):
1204
+ print(VERSION)
1205
+ return
1206
+ if global_args.get("h") or global_args.get("help"):
1207
+ print_help()
1208
+ return
1209
+
1210
+ elif option == "logs":
1211
+ if option_args.get("h") or option_args.get("help"):
1212
+ print_help_logs()
1213
+ return
1214
+ if command is None:
1215
+ raise TtmException("Task ID or name must be provided")
1216
+ if len(command) > 1:
1217
+ raise TtmException(
1218
+ "A single task ID or name must be provided to 'logs'"
1219
+ )
1220
+ follow = option_args.get("f") or option_args.get("follow") or False
1221
+ head = option_args.get("head") or False
1222
+ logs(command[0], follow=follow, head=head)
1223
+
1224
+ elif option == "ls":
1225
+ if option_args.get("h") or option_args.get("help"):
1226
+ print_help_ls()
1227
+ return
1228
+ ls_all = bool(option_args.get("a") or option_args.get("all"))
1229
+ if command:
1230
+ if ls_all:
1231
+ raise TtmException(
1232
+ "-a/--all is not allowed when specific tasks are provided"
1233
+ )
1234
+ ls(ls_all=ls_all, command=command)
1235
+
1236
+ elif option == "rm":
1237
+ if option_args.get("h") or option_args.get("help"):
1238
+ print_help_rm()
1239
+ return
1240
+ rm_all = option_args.get("a") or option_args.get("all")
1241
+ if rm_all is True:
1242
+ rm(None, rm_all=rm_all)
1243
+ else:
1244
+ if command is None:
1245
+ raise TtmException("Task ID or name must be provided")
1246
+ pool = ThreadPool(len(command))
1247
+ results = pool.map(rm, command)
1248
+ if not all(results):
1249
+ exit(1)
1250
+
1251
+ elif option == "run":
1252
+ if option_args.get("h") or option_args.get("help"):
1253
+ print_help_run()
1254
+ return
1255
+ name = option_args.get("n") or option_args.get("name") or None
1256
+ if name is not None:
1257
+ if not re.match(r"^[a-zA-Z_]+$", name):
1258
+ raise TtmException(
1259
+ "Only letters and underscore are allowed in task name"
1260
+ )
1261
+ shell = bool(option_args.get("s") or option_args.get("shell"))
1262
+ if command is None:
1263
+ raise TtmException("A command must be provided")
1264
+ run(command, name=name, shell=shell)
1265
+
1266
+ elif option == "start":
1267
+ if option_args.get("h") or option_args.get("help"):
1268
+ print_help_start()
1269
+ return
1270
+ if command is None:
1271
+ raise TtmException("Task ID or name must be provided")
1272
+ pool = ThreadPool(len(command))
1273
+ results = pool.map(start, command)
1274
+ if not all(results):
1275
+ exit(1)
1276
+
1277
+ elif option == "stop":
1278
+ if option_args.get("h") or option_args.get("help"):
1279
+ print_help_stop()
1280
+ return
1281
+ if command is None:
1282
+ raise TtmException("Task ID or name must be provided")
1283
+ stop_sig = None
1284
+ for sig in signals_list():
1285
+ if option_args.get(sig):
1286
+ if stop_sig is not None:
1287
+ raise TtmException("Only one signal can be provided")
1288
+ stop_sig = int(sig)
1289
+ if option_args.get("k") or option_args.get("kill"):
1290
+ if stop_sig is not None:
1291
+ raise TtmException(
1292
+ "-k/--kill cannot be used when a signal is provided"
1293
+ )
1294
+ stop_sig = SIGKILL
1295
+ if stop_sig is None:
1296
+ stop_sig = SIGTERM
1297
+ pool = ThreadPool(len(command))
1298
+ arg_list = []
1299
+ for c in command:
1300
+ arg_list.append((c, stop_sig))
1301
+ results = pool.starmap(stop, arg_list)
1302
+ if not all(results):
1303
+ exit(1)
1304
+
1305
+ except TtmException as e:
1306
+ print_error(str(e))
1307
+ exit(1)
1308
+
1309
+
1310
+ if __name__ == "__main__":
1311
+ main()