clan_cli: Add TimeTable class to cmd.py. Fix bugs in Machine rewrite
This commit is contained in:
@@ -4,7 +4,8 @@ import select
|
|||||||
import shlex
|
import shlex
|
||||||
import subprocess
|
import subprocess
|
||||||
import sys
|
import sys
|
||||||
from datetime import datetime
|
import weakref
|
||||||
|
from datetime import datetime, timedelta
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import IO, Any
|
from typing import IO, Any
|
||||||
@@ -58,6 +59,45 @@ def handle_output(process: subprocess.Popen, log: Log) -> tuple[str, str]:
|
|||||||
return stdout_buf.decode("utf-8"), stderr_buf.decode("utf-8")
|
return stdout_buf.decode("utf-8"), stderr_buf.decode("utf-8")
|
||||||
|
|
||||||
|
|
||||||
|
class TimeTable:
|
||||||
|
"""
|
||||||
|
This class is used to store the time taken by each command
|
||||||
|
and print it at the end of the program if env PERF=1 is set.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self) -> None:
|
||||||
|
self.table: dict[str, timedelta] = {}
|
||||||
|
weakref.finalize(self, self.table_print)
|
||||||
|
|
||||||
|
def table_print(self) -> None:
|
||||||
|
if os.getenv("PERF") != "1":
|
||||||
|
return
|
||||||
|
print("======== CMD TIMETABLE ========")
|
||||||
|
|
||||||
|
# Sort the table by time in descending order
|
||||||
|
sorted_table = sorted(
|
||||||
|
self.table.items(), key=lambda item: item[1], reverse=True
|
||||||
|
)
|
||||||
|
|
||||||
|
for k, v in sorted_table:
|
||||||
|
# Check if timedelta is greater than 1 second
|
||||||
|
if v.total_seconds() > 1:
|
||||||
|
# Print in red
|
||||||
|
print(f"\033[91mTook {v}s\033[0m for command: '{k}'")
|
||||||
|
else:
|
||||||
|
# Print in default color
|
||||||
|
print(f"Took {v} for command: '{k}'")
|
||||||
|
|
||||||
|
def add(self, cmd: str, time: timedelta) -> None:
|
||||||
|
if cmd in self.table:
|
||||||
|
self.table[cmd] += time
|
||||||
|
else:
|
||||||
|
self.table[cmd] = time
|
||||||
|
|
||||||
|
|
||||||
|
TIME_TABLE = TimeTable()
|
||||||
|
|
||||||
|
|
||||||
def run(
|
def run(
|
||||||
cmd: list[str],
|
cmd: list[str],
|
||||||
*,
|
*,
|
||||||
@@ -83,7 +123,8 @@ def run(
|
|||||||
rc = process.wait()
|
rc = process.wait()
|
||||||
tend = datetime.now()
|
tend = datetime.now()
|
||||||
|
|
||||||
glog.debug(f"Command took {tend - tstart}s to run")
|
global TIME_TABLE
|
||||||
|
TIME_TABLE.add(shlex.join(cmd), tend - tstart)
|
||||||
|
|
||||||
# Wait for the subprocess to finish
|
# Wait for the subprocess to finish
|
||||||
cmd_out = CmdOut(
|
cmd_out = CmdOut(
|
||||||
|
|||||||
@@ -41,6 +41,15 @@ class QMPWrapper:
|
|||||||
|
|
||||||
|
|
||||||
class Machine:
|
class Machine:
|
||||||
|
flake: str | Path
|
||||||
|
name: str
|
||||||
|
data: MachineData
|
||||||
|
eval_cache: dict[str, str]
|
||||||
|
build_cache: dict[str, Path]
|
||||||
|
_flake_path: Path | None
|
||||||
|
_deployment_info: None | dict[str, str]
|
||||||
|
vm: QMPWrapper
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
name: str,
|
name: str,
|
||||||
@@ -57,11 +66,11 @@ class Machine:
|
|||||||
if machine is None:
|
if machine is None:
|
||||||
uri = ClanURI.from_str(str(flake), name)
|
uri = ClanURI.from_str(str(flake), name)
|
||||||
machine = uri.machine
|
machine = uri.machine
|
||||||
self.flake = str(machine.url.value)
|
self.flake: str | Path = machine.url.value
|
||||||
self.name = machine.name
|
self.name: str = machine.name
|
||||||
self.data = machine
|
self.data: MachineData = machine
|
||||||
else:
|
else:
|
||||||
self.data = machine
|
self.data: MachineData = machine
|
||||||
|
|
||||||
self.eval_cache: dict[str, str] = {}
|
self.eval_cache: dict[str, str] = {}
|
||||||
self.build_cache: dict[str, Path] = {}
|
self.build_cache: dict[str, Path] = {}
|
||||||
|
|||||||
@@ -22,7 +22,7 @@ class VmConfig:
|
|||||||
|
|
||||||
def inspect_vm(machine: Machine) -> VmConfig:
|
def inspect_vm(machine: Machine) -> VmConfig:
|
||||||
data = json.loads(machine.eval_nix("config.clanCore.vm.inspect"))
|
data = json.loads(machine.eval_nix("config.clanCore.vm.inspect"))
|
||||||
return VmConfig(flake_url=machine.flake, **data)
|
return VmConfig(flake_url=str(machine.flake), **data)
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|||||||
Reference in New Issue
Block a user