primitive 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

primitive/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2024-present Dylan Stein <dylan@primitive.tech>
 #
 # SPDX-License-Identifier: MIT
-__version__ = "0.2.9"
+__version__ = "0.2.11"

primitive/agent/actions.py

@@ -6,25 +6,29 @@ from loguru import logger
 from primitive.__about__ import __version__
 from primitive.utils.actions import BaseAction
 
-from ..utils.exceptions import P_CLI_100
 from .runner import Runner
 from .uploader import Uploader
+from ..db import sqlite
+from ..db.models import JobRun
 
 
 class Agent(BaseAction):
     def execute(
         self,
     ):
-        logger.enable("primitive")
         logger.remove()
         logger.add(
             sink=sys.stderr,
-            # catch=True,
+            format="<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <level>{message}</level>",
             backtrace=True,
             diagnose=True,
+            level="DEBUG" if self.primitive.DEBUG else "INFO",
         )
-        logger.info(" [*] primitive")
-        logger.info(f" [*] Version: {__version__}")
+        logger.info("[*] primitive agent")
+        logger.info(f"[*] Version: {__version__}")
+
+        # Initialize the database
+        sqlite.init()
 
         # Create uploader
         uploader = Uploader(primitive=self.primitive)
@@ -49,9 +53,6 @@ class Agent(BaseAction):
                 logger.debug("Scanning for files to upload...")
                 uploader.scan()
 
-                logger.debug("Syncing children...")
-                self.primitive.hardware._sync_children()
-
                 hardware = self.primitive.hardware.get_own_hardware_details()
 
                 if hardware["activeReservation"]:
@@ -108,6 +109,7 @@ class Agent(BaseAction):
                 ]
 
                 if not pending_job_runs:
+                    self.primitive.hardware.check_in_http(is_online=True)
                     sleep_amount = 5
                     logger.debug(
                         f"Waiting for Job Runs... [sleeping {sleep_amount} seconds]"
@@ -120,6 +122,11 @@ class Agent(BaseAction):
                     logger.debug(f"Job Run ID: {job_run['id']}")
                     logger.debug(f"Job Name: {job_run['job']['name']}")
 
+                    JobRun.objects.create(
+                        job_run_id=job_run["id"],
+                        pid=None,
+                    )
+
                     runner = Runner(
                         primitive=self.primitive,
                         job_run=job_run,
@@ -137,6 +144,7 @@ class Agent(BaseAction):
                             status="request_completed",
                             conclusion="failure",
                         )
+                        JobRun.objects.filter_by(job_run_id=job_run["id"]).delete()
                         continue
 
                     try:
@@ -152,17 +160,11 @@ class Agent(BaseAction):
                         runner.cleanup()
 
                         # NOTE: also run scan here to force upload of artifacts
-                        # This should probably eventuall be another daemon?
+                        # This should probably eventually be another daemon?
                         uploader.scan()
 
+                        JobRun.objects.filter_by(job_run_id=job_run["id"]).delete()
+
                 sleep(5)
         except KeyboardInterrupt:
-            logger.info(" [*] Stopping primitive...")
-            try:
-                self.primitive.hardware.check_in_http(
-                    is_available=False, is_online=False, stopping_agent=True
-                )
-            except P_CLI_100 as exception:
-                logger.error(" [*] Error stopping primitive.")
-                logger.error(str(exception))
-            sys.exit()
+            logger.info("[*] Stopping primitive agent...")

primitive/agent/runner.py

@@ -2,12 +2,12 @@ import asyncio
 import os
 import re
 import shutil
-import time
 import typing
 from abc import abstractmethod
 from enum import Enum, IntEnum
 from pathlib import Path, PurePath
 from typing import Dict, List, TypedDict
+from ..db.models import JobRun
 
 import yaml
 from loguru import logger
@@ -80,8 +80,6 @@ class Runner:
         self.modified_env = {}
         self.file_logger = None
 
-        logger.enable("primitive")
-
         # If max_log_size set to <= 0, disable file logging
         if max_log_size > 0:
             log_name = f"{self.job['slug']}_{self.job_run['jobRunNumber']}_{{time}}.primitive.log"
@@ -158,22 +156,40 @@ class Runner:
         self.modified_env = {**self.initial_env}
 
         task_failed = False
+        cancelled = False
         conclusion = "success"
         for task in self.config["executes"]:
+            status = self.primitive.jobs.get_job_status(self.job_run["id"])
+            status_value = status.data["jobRun"]["status"]
+            conclusion_value = status.data["jobRun"]["conclusion"]
+
+            if status_value == "completed" and conclusion_value == "cancelled":
+                cancelled = True
+                break
+
             with logger.contextualize(label=task["label"]):
                 with asyncio.Runner() as async_runner:
                     if task_failed := async_runner.run(self.run_task(task)):
                         break
 
+        number_of_files_produced = self.get_number_of_files_produced()
+        logger.info(
+            f"Produced {number_of_files_produced} files for {self.job['slug']} job"
+        )
+
+        if cancelled:
+            logger.warning("Job cancelled by user")
+            self.primitive.jobs.job_run_update(
+                self.job_run["id"],
+                number_of_files_produced=number_of_files_produced,
+            )
+            return
+
         if task_failed:
             conclusion = "failure"
         else:
             logger.success(f"Completed {self.job['slug']} job")
 
-        number_of_files_produced = self.get_number_of_files_produced()
-        logger.info(
-            f"Produced {number_of_files_produced} files for {self.job['slug']} job"
-        )
         self.primitive.jobs.job_run_update(
             self.job_run["id"],
             status="request_completed",
@@ -187,11 +203,19 @@ class Runner:
 
         # Logs can be produced even if no artifact stores are created for the job run.
         job_run_logs_cache = get_logs_cache(self.job_run["id"])
-        log_files = [
-            file
-            for _, _, current_path_files in job_run_logs_cache.walk()
-            for file in current_path_files
-        ]
+        has_walk = getattr(job_run_logs_cache, "walk", None)
+        if has_walk:
+            log_files = [
+                file
+                for _, _, current_path_files in job_run_logs_cache.walk()
+                for file in current_path_files
+            ]
+        else:
+            log_files = [
+                file
+                for _, _, current_path_files in os.walk(job_run_logs_cache)
+                for file in current_path_files
+            ]
 
         number_of_files_produced += len(log_files)
 
@@ -199,11 +223,19 @@ class Runner:
             return number_of_files_produced
 
         job_run_artifacts_cache = get_artifacts_cache(self.job_run["id"])
-        artifact_files = [
-            file
-            for _, _, current_path_files in job_run_artifacts_cache.walk()
-            for file in current_path_files
-        ]
+        has_walk = getattr(job_run_artifacts_cache, "walk", None)
+        if has_walk:
+            artifact_files = [
+                file
+                for _, _, current_path_files in job_run_artifacts_cache.walk()
+                for file in current_path_files
+            ]
+        else:
+            artifact_files = [
+                file
+                for _, _, current_path_files in os.walk(job_run_artifacts_cache)
+                for file in current_path_files
+            ]
 
         number_of_files_produced += len(artifact_files)
 
@@ -233,24 +265,24 @@ class Runner:
                 stderr=asyncio.subprocess.PIPE,
             )
 
-            loop = asyncio.get_running_loop()
-            monitor_task = loop.run_in_executor(None, self.monitor_cmd, process)
+            JobRun.objects.filter_by(job_run_id=self.job_run["id"]).update(
+                {"pid": process.pid}
+            )
 
-            stdout_failed, stderr_failed, cancelled = await asyncio.gather(
+            stdout_failed, stderr_failed = await asyncio.gather(
                 self.log_cmd(
                     process=process, stream=process.stdout, tags=task.get("tags", {})
                 ),
                 self.log_cmd(
                     process=process, stream=process.stderr, tags=task.get("tags", {})
                 ),
-                monitor_task,
             )
 
             returncode = await process.wait()
 
-            if cancelled:
-                logger.warning("Job cancelled by user")
-                return True
+            JobRun.objects.filter_by(job_run_id=self.job_run["id"]).update(
+                {"pid": None}
+            )
 
             if returncode > 0:
                 logger.error(
@@ -339,25 +371,6 @@ class Runner:
 
         return [line for line in lines if len(line) > 0]
 
-    def monitor_cmd(self, process) -> bool:
-        while process.returncode is None:
-            status = self.primitive.jobs.get_job_status(self.job_run["id"])
-
-            status_value = status.data["jobRun"]["status"]
-            conclusion_value = status.data["jobRun"]["conclusion"]
-
-            if status_value == "completed" and conclusion_value == "cancelled":
-                try:
-                    process.terminate()
-                except ProcessLookupError:
-                    pass
-
-                return True
-
-            time.sleep(10)
-
-        return False
-
     def cleanup(self) -> None:
         logger.remove(self.file_logger)
 

primitive/cli.py

@@ -16,6 +16,7 @@ from .jobs.commands import cli as jobs_commands
 from .organizations.commands import cli as organizations_commands
 from .projects.commands import cli as projects_commands
 from .reservations.commands import cli as reservations_commands
+from .monitor.commands import cli as monitor_commands
 
 
 @click.group()
@@ -71,6 +72,7 @@ cli.add_command(organizations_commands, "organizations")
 cli.add_command(projects_commands, "projects")
 cli.add_command(reservations_commands, "reservations")
 cli.add_command(exec_commands, "exec")
+cli.add_command(monitor_commands, "monitor")
 
 if __name__ == "__main__":
     cli(obj={})

primitive/client.py

@@ -1,7 +1,8 @@
-import sys
-
 from gql import Client
 from loguru import logger
+from rich.logging import RichHandler
+from rich.traceback import install
+from typing import Optional
 
 from .agent.actions import Agent
 from .auth.actions import Auth
@@ -15,10 +16,9 @@ from .organizations.actions import Organizations
 from .projects.actions import Projects
 from .provisioning.actions import Provisioning
 from .reservations.actions import Reservations
+from .monitor.actions import Monitor
 from .utils.config import read_config_file
 
-logger.disable("primitive")
-
 
 class Primitive:
     def __init__(
@@ -26,24 +26,48 @@ class Primitive:
         host: str = "api.primitive.tech",
         DEBUG: bool = False,
         JSON: bool = False,
-        token: str = None,
-        transport: str = None,
+        token: Optional[str] = None,
+        transport: Optional[str] = None,
     ) -> None:
         self.host: str = host
-        self.session: Client = None
+        self.session: Optional[Client] = None
         self.DEBUG: bool = DEBUG
         self.JSON: bool = JSON
 
+        # Enable tracebacks with local variables
         if self.DEBUG:
-            logger.enable("primitive")
-            logger.remove()
-            logger.add(
-                sink=sys.stderr,
-                serialize=self.JSON,
-                catch=True,
-                backtrace=True,
-                diagnose=True,
-            )
+            install(show_locals=True)
+
+        # Configure rich logging handler
+        rich_handler = RichHandler(
+            rich_tracebacks=self.DEBUG,  # Pretty tracebacks
+            markup=True,  # Allow Rich markup tags
+            show_time=self.DEBUG,  # Show timestamps
+            show_level=self.DEBUG,  # Show log levels
+            show_path=self.DEBUG,  # Hide source path (optional)
+        )
+
+        def formatter(record) -> str:
+            match record["level"].name:
+                case "ERROR":
+                    return "[bold red]Error>[/bold red] {name}:{function}:{line} - {message}"
+                case "CRITICAL":
+                    return "[italic bold red]Critical>[/italic bold red] {name}:{function}:{line} - {message}"
+                case "WARNING":
+                    return "[bold yellow]Warning>[/bold yellow] {message}"
+                case _:
+                    return "[#666666]>[/#666666] {message}"
+
+        logger.remove()
+        logger.add(
+            sink=rich_handler,
+            format="{message}" if self.DEBUG else formatter,
+            level="DEBUG" if self.DEBUG else "INFO",
+            backtrace=self.DEBUG,
+        )
+
+        # Nothing will print here if DEBUG is false
+        logger.debug("Debug mode enabled")
 
         # Generate full or partial host config
         if not token and not transport:
@@ -67,6 +91,7 @@ class Primitive:
         self.daemons: Daemons = Daemons(self)
         self.exec: Exec = Exec(self)
         self.provisioning: Provisioning = Provisioning(self)
+        self.monitor: Monitor = Monitor(self)
 
     def get_host_config(self):
         self.full_config = read_config_file()

primitive/daemons/actions.py

@@ -1,24 +1,13 @@
 import platform
 import typing
+from typing import Dict, Optional, List
 
 if typing.TYPE_CHECKING:
     from ..client import Primitive
 
-from .launch_agents import (
-    full_launch_agent_install,
-    full_launch_agent_uninstall,
-    start_launch_agent,
-    stop_launch_agent,
-    view_launch_agent_logs,
-)
-
-from .launch_service import (
-    full_service_install,
-    full_service_uninstall,
-    start_service,
-    stop_service,
-    view_service_logs,
-)
+from .launch_agents import LaunchAgent
+from .launch_service import LaunchService
+from ..utils.daemons import Daemon
 
 
 class Daemons:
@@ -26,50 +15,47 @@ class Daemons:
         self.primitive: Primitive = primitive
         self.os_family = platform.system()
 
-    def install(self):
-        result = True
-        if self.os_family == "Darwin":
-            full_launch_agent_install()
-        elif self.os_family == "Linux":
-            full_service_install()
-        elif self.os_family == "Windows":
-            print("Not Implemented")
-        return result
-
-    def uninstall(self):
-        result = True
-        if self.os_family == "Darwin":
-            full_launch_agent_uninstall()
-        elif self.os_family == "Linux":
-            full_service_uninstall()
-        elif self.os_family == "Windows":
-            print("Not Implemented")
-        return result
-
-    def stop(self) -> bool:
-        result = True
-        if self.os_family == "Darwin":
-            result = stop_launch_agent()
-        elif self.os_family == "Linux":
-            stop_service()
-        elif self.os_family == "Windows":
-            print("Not Implemented")
-        return result
-
-    def start(self) -> bool:
-        result = True
-        if self.os_family == "Darwin":
-            result = start_launch_agent()
-        elif self.os_family == "Linux":
-            start_service()
-        elif self.os_family == "Windows":
-            print("Not Implemented")
-        return result
-
-    def logs(self):
-        if self.os_family == "Darwin":
-            view_launch_agent_logs()
-        elif self.os_family == "Linux":
-            view_service_logs()
-        elif self.os_family == "Windows":
-            print("Not Implemented")
+        match self.os_family:
+            case "Darwin":
+                self.daemons: Dict[str, Daemon] = {
+                    "agent": LaunchAgent("tech.primitive.agent"),
+                    "monitor": LaunchAgent("tech.primitive.monitor"),
+                }
+            case "Linux":
+                self.daemons: Dict[str, Daemon] = {
+                    "agent": LaunchService("tech.primitive.agent"),
+                    "monitor": LaunchService("tech.primitive.monitor"),
+                }
+            case _:
+                raise NotImplementedError(f"{self.os_family} is not supported.")
+
+    def install(self, name: Optional[str]) -> bool:
+        if name:
+            return self.daemons[name].install()
+        else:
+            return all([daemon.install() for daemon in self.daemons.values()])
+
+    def uninstall(self, name: Optional[str]) -> bool:
+        if name:
+            return self.daemons[name].uninstall()
+        else:
+            return all([daemon.uninstall() for daemon in self.daemons.values()])
+
+    def stop(self, name: Optional[str]) -> bool:
+        if name:
+            return self.daemons[name].stop()
+        else:
+            return all([daemon.stop() for daemon in self.daemons.values()])
+
+    def start(self, name: Optional[str]) -> bool:
+        if name:
+            return self.daemons[name].start()
+        else:
+            return all([daemon.start() for daemon in self.daemons.values()])
+
+    def list(self) -> List[Daemon]:
+        """List all daemons"""
+        return list(self.daemons.values())
+
+    def logs(self, name: str) -> None:
+        self.daemons[name].view_logs()

primitive/daemons/commands.py

@@ -1,7 +1,10 @@
 import click
 
-from ..utils.printer import print_result
 import typing
+from typing import Optional
+from .ui import render_daemon_list
+
+from loguru import logger
 
 if typing.TYPE_CHECKING:
     from ..client import Primitive
@@ -16,50 +19,93 @@ def cli(context):
 
 @cli.command("install")
 @click.pass_context
-def install_daemon_command(context):
+@click.argument(
+    "name",
+    type=str,
+    required=False,
+)
+def install_daemon_command(context, name: Optional[str]):
     """Install the full primitive daemon"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result = primitive.daemons.install()
-    print_result(message=result, context=context)
+    installed = primitive.daemons.install(name=name)
+
+    if installed:
+        logger.info(":white_check_mark: daemon(s) installed successfully!")
+    else:
+        logger.error("Unable to install daemon(s).")
 
 
 @cli.command("uninstall")
 @click.pass_context
-def uninstall_daemon_command(context):
+@click.argument(
+    "name",
+    type=str,
+    required=False,
+)
+def uninstall_daemon_command(context, name: Optional[str]):
     """Uninstall the full primitive Daemon"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result = primitive.daemons.uninstall()
-    print_result(message=result, context=context)
+    uninstalled = primitive.daemons.uninstall(name=name)
+
+    if uninstalled:
+        logger.info(":white_check_mark: daemon(s) uninstalled successfully!")
+    else:
+        logger.error("Unable to uninstall daemon(s).")
 
 
 @cli.command("stop")
 @click.pass_context
-def stop_daemon_command(context):
+@click.argument(
+    "name",
+    type=str,
+    required=False,
+)
+def stop_daemon_command(context, name: Optional[str]):
     """Stop primitive Daemon"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result = primitive.daemons.stop()
-    message = "stopping primitive daemon"
-    if context.obj["JSON"]:
-        message = result
-    print_result(message=message, context=context)
+    stopped = primitive.daemons.stop(name=name)
+
+    if stopped:
+        logger.info(":white_check_mark: daemon(s) stopped successfully!")
+    else:
+        logger.error("Unable to stop daemon(s).")
 
 
 @cli.command("start")
 @click.pass_context
-def start_daemon_command(context):
+@click.argument(
+    "name",
+    type=str,
+    required=False,
+)
+def start_daemon_command(context, name: Optional[str]):
     """Start primitive Daemon"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result = primitive.daemons.start()
-    message = "starting primitive daemon"
-    if context.obj["JSON"]:
-        message = result
-    print_result(message=message, context=context)
+    started = primitive.daemons.start(name=name)
+
+    if started:
+        logger.info(":white_check_mark: daemon(s) started successfully!")
+    else:
+        logger.error("Unable to start daemon(s).")
 
 
 @cli.command("logs")
 @click.pass_context
-def log_daemon_command(context):
+@click.argument(
+    "name",
+    type=str,
+    required=True,
+)
+def log_daemon_command(context, name: str):
     """Logs from primitive Daemon"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result = primitive.daemons.logs()
-    print_result(message=result, context=context)
+    primitive.daemons.logs(name=name)
+
+
+@cli.command("list")
+@click.pass_context
+def list_daemon_command(context):
+    """List all daemons"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    daemon_list = primitive.daemons.list()
+    render_daemon_list(daemons=daemon_list)