clan_cli: Add TimeTable class to cmd.py. Fix bugs in Machine rewrite
All checks were successful
checks / check-links (pull_request) Successful in 22s
checks / checks-impure (pull_request) Successful in 2m1s
checks / checks (pull_request) Successful in 3m51s

This commit is contained in:
Luis Hebendanz 2024-03-08 22:01:54 +07:00
parent 068f89e453
commit e4896814f2
3 changed files with 57 additions and 7 deletions

View File

@ -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(

View File

@ -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] = {}

View File

@ -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