watch_fma_download.py
2.06 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#!/usr/bin/env python3
"""Ensure the FMA archive download keeps running; restart if stalled or dead."""
from __future__ import annotations
import argparse
import json
import subprocess
import time
from pathlib import Path
PYTHON = "/usr/local/miniconda3/bin/python"
INSPECT = [PYTHON, "scripts/prepare_fma_archive.py", "inspect"]
BG = [PYTHON, "scripts/prepare_fma_archive.py", "bg-download"]
DEFAULT_LOG = Path("/tmp/fma_modelscope_watch.log")
def inspect() -> dict:
out = subprocess.check_output(INSPECT, text=True)
return json.loads(out)
def bg_download() -> dict:
out = subprocess.check_output(BG, text=True)
return json.loads(out)
def has_live_curl() -> bool:
proc = subprocess.run(
["bash", "-lc", "ps -ef | grep 'fma_small.zip' | grep -v grep >/dev/null"],
capture_output=True,
text=True,
)
return proc.returncode == 0
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--interval", type=float, default=5.0)
parser.add_argument("--cycles", type=int, default=3)
parser.add_argument("--log-path", default=str(DEFAULT_LOG))
args = parser.parse_args()
log_path = Path(args.log_path)
log_path.parent.mkdir(parents=True, exist_ok=True)
events = []
previous_size = None
for _ in range(args.cycles):
snapshot = inspect()
size = int(snapshot["archive_size"])
live = has_live_curl()
restarted = None
if (previous_size is not None and size <= previous_size) or not live:
restarted = bg_download()
time.sleep(2)
snapshot = inspect()
size = int(snapshot["archive_size"])
live = has_live_curl()
event = {
"snapshot": snapshot,
"live_curl": live,
"restarted": restarted,
}
events.append(event)
previous_size = size
time.sleep(args.interval)
text = json.dumps({"status": "ok", "events": events}, indent=2, ensure_ascii=False)
log_path.write_text(text)
print(text)
if __name__ == "__main__":
main()