You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
from tempfile import NamedTemporaryFile
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(PROJECT_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from common.nn_profile_schedule import ( # noqa: E402
|
|
DEFAULT_PROFILE,
|
|
format_active_profile_env,
|
|
get_requested_active_profile,
|
|
load_simple_env_file,
|
|
normalize_profile_name,
|
|
resolve_active_profile,
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Switch DroneDetector NN model profile by schedule")
|
|
parser.add_argument(
|
|
"--project-root",
|
|
default=str(PROJECT_ROOT),
|
|
help="Path to repository root (default: script parent project)",
|
|
)
|
|
parser.add_argument(
|
|
"--no-restart",
|
|
action="store_true",
|
|
help="Update runtime profile file without restarting dronedetector-nn-server",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
project_root = Path(args.project_root).resolve()
|
|
env_path = project_root / ".env"
|
|
runtime_path = project_root / "runtime" / "nn_active_profile.env"
|
|
compose_file = project_root / "deploy" / "docker" / "docker-compose.yml"
|
|
|
|
if not env_path.exists():
|
|
raise SystemExit(f"Missing root .env: {env_path}")
|
|
if not compose_file.exists():
|
|
raise SystemExit(f"Missing docker-compose file: {compose_file}")
|
|
|
|
config = load_simple_env_file(env_path)
|
|
requested_profile = get_requested_active_profile(config)
|
|
desired_profile = resolve_active_profile(config)
|
|
runtime_exists = runtime_path.exists()
|
|
current_profile = read_current_profile(runtime_path)
|
|
|
|
if requested_profile != desired_profile:
|
|
print(
|
|
f"[nn_profile_switcher] requested profile {requested_profile} is not configured, "
|
|
f"falling back to {desired_profile}"
|
|
)
|
|
|
|
if current_profile == desired_profile and runtime_exists:
|
|
print(f"[nn_profile_switcher] profile already active: {desired_profile}")
|
|
return 0
|
|
|
|
previous_contents = runtime_path.read_text(encoding="utf-8") if runtime_path.exists() else None
|
|
write_runtime_profile(runtime_path, desired_profile)
|
|
if runtime_exists:
|
|
print(
|
|
f"[nn_profile_switcher] updated runtime profile {current_profile or DEFAULT_PROFILE} -> {desired_profile}"
|
|
)
|
|
else:
|
|
print(f"[nn_profile_switcher] initialized runtime profile: {desired_profile}")
|
|
|
|
if current_profile == desired_profile:
|
|
return 0
|
|
|
|
if args.no_restart:
|
|
print("[nn_profile_switcher] --no-restart set, skipping container restart")
|
|
return 0
|
|
|
|
try:
|
|
restart_nn_server(project_root, compose_file)
|
|
except Exception:
|
|
restore_runtime_profile(runtime_path, previous_contents)
|
|
raise
|
|
|
|
print("[nn_profile_switcher] restarted dronedetector-nn-server")
|
|
return 0
|
|
|
|
|
|
def read_current_profile(runtime_path: Path) -> str:
|
|
runtime_values = load_simple_env_file(runtime_path)
|
|
return normalize_profile_name(runtime_values.get("NN_ACTIVE_PROFILE"))
|
|
|
|
|
|
def write_runtime_profile(runtime_path: Path, profile: str) -> None:
|
|
runtime_path.parent.mkdir(parents=True, exist_ok=True)
|
|
contents = format_active_profile_env(profile)
|
|
with NamedTemporaryFile("w", encoding="utf-8", dir=runtime_path.parent, delete=False) as handle:
|
|
handle.write(contents)
|
|
temp_path = Path(handle.name)
|
|
temp_path.replace(runtime_path)
|
|
|
|
|
|
def restore_runtime_profile(runtime_path: Path, previous_contents: str | None) -> None:
|
|
if previous_contents is None:
|
|
runtime_path.unlink(missing_ok=True)
|
|
return
|
|
|
|
with NamedTemporaryFile("w", encoding="utf-8", dir=runtime_path.parent, delete=False) as handle:
|
|
handle.write(previous_contents)
|
|
temp_path = Path(handle.name)
|
|
temp_path.replace(runtime_path)
|
|
|
|
|
|
def restart_nn_server(project_root: Path, compose_file: Path) -> None:
|
|
subprocess.run(
|
|
[
|
|
"docker",
|
|
"compose",
|
|
"-f",
|
|
str(compose_file),
|
|
"restart",
|
|
"dronedetector-nn-server",
|
|
],
|
|
cwd=project_root,
|
|
check=True,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|