213 lines
6.6 KiB
Python
213 lines
6.6 KiB
Python
# Crown Owned Copyright (C) Dstl 2023. DEFCON 703. Shared in confidence.
|
|
"""Provides a CLI using Typer as an entry point."""
|
|
import logging
|
|
import os
|
|
import shutil
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import pkg_resources
|
|
import typer
|
|
import yaml
|
|
from platformdirs import PlatformDirs
|
|
from typing_extensions import Annotated
|
|
|
|
from primaite.data_viz import PlotlyTemplate
|
|
|
|
app = typer.Typer()
|
|
|
|
|
|
@app.command()
|
|
def build_dirs():
|
|
"""Build the PrimAITE app directories."""
|
|
from primaite.setup import setup_app_dirs
|
|
|
|
setup_app_dirs.run()
|
|
|
|
|
|
@app.command()
|
|
def reset_notebooks(overwrite: bool = True):
|
|
"""
|
|
Force a reset of the demo notebooks in the users notebooks directory.
|
|
|
|
:param overwrite: If True, will overwrite existing demo notebooks.
|
|
"""
|
|
from primaite.setup import reset_demo_notebooks
|
|
|
|
reset_demo_notebooks.run(overwrite)
|
|
|
|
|
|
@app.command()
|
|
def logs(last_n: Annotated[int, typer.Option("-n")]):
|
|
"""
|
|
Print the PrimAITE log file.
|
|
|
|
:param last_n: The number of lines to print. Default value is 10.
|
|
"""
|
|
import re
|
|
|
|
from primaite import LOG_PATH
|
|
|
|
if os.path.isfile(LOG_PATH):
|
|
with open(LOG_PATH) as file:
|
|
lines = file.readlines()
|
|
for line in lines[-last_n:]:
|
|
print(re.sub(r"\n*", "", line))
|
|
|
|
|
|
_LogLevel = Enum("LogLevel", {k: k for k in logging._levelToName.values()}) # noqa
|
|
|
|
|
|
@app.command()
|
|
def log_level(level: Annotated[Optional[_LogLevel], typer.Argument()] = None):
|
|
"""
|
|
View or set the PrimAITE Log Level.
|
|
|
|
To View, simply call: primaite log-level
|
|
|
|
To set, call: primaite log-level <desired log level>
|
|
|
|
For example, to set the to debug, call: primaite log-level DEBUG
|
|
"""
|
|
app_dirs = PlatformDirs(appname="primaite")
|
|
app_dirs.user_config_path.mkdir(exist_ok=True, parents=True)
|
|
user_config_path = app_dirs.user_config_path / "primaite_config.yaml"
|
|
if user_config_path.exists():
|
|
with open(user_config_path, "r") as file:
|
|
primaite_config = yaml.safe_load(file)
|
|
|
|
if level:
|
|
primaite_config["logging"]["log_level"] = level.value
|
|
with open(user_config_path, "w") as file:
|
|
yaml.dump(primaite_config, file)
|
|
print(f"PrimAITE Log Level: {level}")
|
|
else:
|
|
level = primaite_config["logging"]["log_level"]
|
|
print(f"PrimAITE Log Level: {level}")
|
|
|
|
|
|
@app.command()
|
|
def notebooks():
|
|
"""Start Jupyter Lab in the users PrimAITE notebooks directory."""
|
|
from primaite.notebooks import start_jupyter_session
|
|
|
|
start_jupyter_session()
|
|
|
|
|
|
@app.command()
|
|
def version():
|
|
"""Get the installed PrimAITE version number."""
|
|
import primaite
|
|
|
|
print(primaite.__version__)
|
|
|
|
|
|
@app.command()
|
|
def clean_up():
|
|
"""Cleans up left over files from previous version installations."""
|
|
from primaite.setup import old_installation_clean_up
|
|
|
|
old_installation_clean_up.run()
|
|
|
|
|
|
@app.command()
|
|
def setup(overwrite_existing: bool = True):
|
|
"""
|
|
Perform the PrimAITE first-time setup.
|
|
|
|
WARNING: All user-data will be lost.
|
|
"""
|
|
# Does this way to avoid using PrimAITE package before config is loaded
|
|
app_dirs = PlatformDirs(appname="primaite")
|
|
app_dirs.user_config_path.mkdir(exist_ok=True, parents=True)
|
|
user_config_path = app_dirs.user_config_path / "primaite_config.yaml"
|
|
pkg_config_path = Path(pkg_resources.resource_filename("primaite", "setup/_package_data/primaite_config.yaml"))
|
|
|
|
shutil.copy2(pkg_config_path, user_config_path)
|
|
|
|
from primaite import getLogger
|
|
from primaite.setup import old_installation_clean_up, reset_demo_notebooks, reset_example_configs, setup_app_dirs
|
|
|
|
_LOGGER = getLogger(__name__)
|
|
|
|
_LOGGER.info("Performing the PrimAITE first-time setup...")
|
|
|
|
_LOGGER.info("Building primaite_config.yaml...")
|
|
|
|
_LOGGER.info("Building the PrimAITE app directories...")
|
|
setup_app_dirs.run()
|
|
|
|
_LOGGER.info("Rebuilding the demo notebooks...")
|
|
reset_demo_notebooks.run(overwrite_existing=True)
|
|
|
|
_LOGGER.info("Rebuilding the example notebooks...")
|
|
reset_example_configs.run(overwrite_existing=True)
|
|
|
|
_LOGGER.info("Performing a clean-up of previous PrimAITE installations...")
|
|
old_installation_clean_up.run()
|
|
|
|
_LOGGER.info("PrimAITE setup complete!")
|
|
|
|
|
|
@app.command()
|
|
def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[str] = None):
|
|
"""
|
|
Run a PrimAITE session.
|
|
|
|
tc: The training config filepath. Optional. If no value is passed then
|
|
example default training config is used from:
|
|
~/primaite/config/example_config/training/training_config_main.yaml.
|
|
|
|
ldc: The lay down config file path. Optional. If no value is passed then
|
|
example default lay down config is used from:
|
|
~/primaite/config/example_config/lay_down/lay_down_config_3_doc_very_basic.yaml.
|
|
|
|
load: The directory of a previous session. Optional. If no value is passed, then the session
|
|
will use the default training config and laydown config. Inversely, if a training config and laydown config
|
|
is passed while a session directory is passed, PrimAITE will load the session and ignore the training config
|
|
and laydown config.
|
|
"""
|
|
from primaite.config.lay_down_config import dos_very_basic_config_path
|
|
from primaite.config.training_config import main_training_config_path
|
|
from primaite.main import run
|
|
|
|
if load is not None:
|
|
run(session_path=load)
|
|
|
|
if not tc:
|
|
tc = main_training_config_path()
|
|
|
|
if not ldc:
|
|
ldc = dos_very_basic_config_path()
|
|
|
|
run(training_config_path=tc, lay_down_config_path=ldc)
|
|
|
|
|
|
@app.command()
|
|
def plotly_template(template: Annotated[Optional[PlotlyTemplate], typer.Argument()] = None):
|
|
"""
|
|
View or set the plotly template for Session plots.
|
|
|
|
To View, simply call: primaite plotly-template
|
|
|
|
To set, call: primaite plotly-template <desired template>
|
|
|
|
For example, to set as plotly_dark, call: primaite plotly-template PLOTLY_DARK
|
|
"""
|
|
app_dirs = PlatformDirs(appname="primaite")
|
|
app_dirs.user_config_path.mkdir(exist_ok=True, parents=True)
|
|
user_config_path = app_dirs.user_config_path / "primaite_config.yaml"
|
|
if user_config_path.exists():
|
|
with open(user_config_path, "r") as file:
|
|
primaite_config = yaml.safe_load(file)
|
|
|
|
if template:
|
|
primaite_config["session"]["outputs"]["plots"]["template"] = template.value
|
|
with open(user_config_path, "w") as file:
|
|
yaml.dump(primaite_config, file)
|
|
print(f"PrimAITE plotly template: {template.value}")
|
|
else:
|
|
template = primaite_config["session"]["outputs"]["plots"]["template"]
|
|
print(f"PrimAITE plotly template: {template}")
|