business-assistant/scripts/memory_service.py

"""Render (never install) a Linux systemd memory-isolated API service."""

import argparse
import hashlib
from pathlib import Path
import re
import sys


def unit_name(root):
    return 'business-assistant-' + hashlib.sha256(str(root).encode()).hexdigest()[:12] + '.service'


def render(root, user, memory_mb, *, offline=False):
    root = str(root)
    # Keep the generated systemd command unambiguous; no shell evaluation or
    # systemd specifiers from user-controlled paths.
    if not re.fullmatch(r'/[A-Za-z0-9_./-]+', root) or '..' in root.split('/'):
        raise ValueError('service_path_requires_simple_absolute_linux_path')
    if not re.fullmatch(r'[a-z_][a-z0-9_-]{0,31}', user) or user == 'root':
        raise ValueError('service_requires_non_root_user')
    if type(memory_mb) is not int or not 2048 <= memory_mb <= 1048576:
        raise ValueError('invalid_memory_budget')
    soft = memory_mb * 3 // 4
    deploy_script = 'deploy-offline.sh' if offline else 'deploy.sh'
    return f'''[Unit]
Description=Business Assistant memory-isolated API
After=network-online.target
Wants=network-online.target
StartLimitIntervalSec=300
StartLimitBurst=3

[Service]
Type=simple
User={user}
WorkingDirectory={root}
ExecStart=/bin/sh {root}/{deploy_script} serve
Environment=API_REQUIRE_MEMORY_CGROUP=1
Environment=API_MEMORY_HARD_LIMIT_MB={memory_mb}
Environment=API_MEMORY_BUDGET_MB={soft}
MemoryAccounting=yes
MemoryHigh={memory_mb * 7 // 8}M
MemoryMax={memory_mb}M
MemorySwapMax=0
OOMPolicy=kill
KillMode=control-group
Restart=on-failure
RestartSec=10
TimeoutStopSec=45
NoNewPrivileges=yes

[Install]
WantedBy=multi-user.target
'''


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--user', required=True)
    parser.add_argument('--memory-max-mb', type=int, required=True)
    parser.add_argument('--offline', action='store_true')
    args = parser.parse_args()
    root = Path(__file__).resolve().parents[1]
    content = render(root, args.user, args.memory_max_mb, offline=args.offline)
    output = root / 'output/deploy' / unit_name(root)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(content, encoding='utf-8')
    print(output)


if __name__ == '__main__':
    main()