primitive 0.2.10__py3-none-any.whl → 0.2.12__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.10"
+__version__ = "0.2.12"

primitive/agent/actions.py

@@ -6,7 +6,8 @@ from loguru import logger
 from primitive.__about__ import __version__
 from primitive.utils.actions import BaseAction
 
-from ..utils.exceptions import P_CLI_100
+from ..db import sqlite
+from ..db.models import JobRun
 from .runner import Runner
 from .uploader import Uploader
 
@@ -15,154 +16,95 @@ 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)
 
-        # self.primitive.hardware.update_hardware_system_info()
-        try:
-            # hey stupid:
-            # do not set is_available to True here, it will mess up the reservation logic
-            # only set is_available after we've checked that no active reservation is present
-            # setting is_available of the parent also effects the children,
-            # which may have active reservations as well
-            self.primitive.hardware.check_in_http(is_online=True)
-        except Exception as exception:
-            logger.exception(f"Error checking in hardware: {exception}")
-            sys.exit(1)
-
         try:
-            active_reservation_id = None
-            active_reservation_pk = None
-
             while True:
                 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"]:
-                    if (
-                        hardware["activeReservation"]["id"] != active_reservation_id
-                        or hardware["activeReservation"]["pk"] != active_reservation_pk
-                    ):
-                        logger.warning("New reservation for this hardware.")
-                        active_reservation_id = hardware["activeReservation"]["id"]
-                        active_reservation_pk = hardware["activeReservation"]["pk"]
-                        logger.debug("Active Reservation:")
-                        logger.debug(f"Node ID: {active_reservation_id}")
-                        logger.debug(f"PK: {active_reservation_pk}")
-
-                        logger.debug("Running pre provisioning steps for reservation.")
-                        self.primitive.provisioning.add_reservation_authorized_keys(
-                            reservation_id=active_reservation_id
-                        )
-                else:
-                    if (
-                        hardware["activeReservation"] is None
-                        and active_reservation_id is not None
-                        # and hardware["isAvailable"] NOTE: this condition was causing the CLI to get into a loop searching for job runs
-                    ):
-                        logger.debug("Previous Reservation is Complete:")
-                        logger.debug(f"Node ID: {active_reservation_id}")
-                        logger.debug(f"PK: {active_reservation_pk}")
-                        logger.debug(
-                            "Running cleanup provisioning steps for reservation."
-                        )
-                        self.primitive.provisioning.remove_reservation_authorized_keys(
-                            reservation_id=active_reservation_id
-                        )
-                        active_reservation_id = None
-                        active_reservation_pk = None
-
-                if not active_reservation_id:
-                    self.primitive.hardware.check_in_http(
-                        is_available=True, is_online=True
-                    )
+                db_job_run = JobRun.objects.first()
+
+                if not db_job_run:
                     sleep_amount = 5
                     logger.debug(
-                        f"No active reservation found... [sleeping {sleep_amount} seconds]"
+                        f"No pending job runs... [sleeping {sleep_amount} seconds]"
                     )
                     sleep(sleep_amount)
                     continue
 
-                job_runs_result = self.primitive.jobs.get_job_runs(
-                    status="pending", first=1, reservation_id=active_reservation_id
+                api_job_run_data = self.primitive.jobs.get_job_run(
+                    id=db_job_run.job_run_id,
                 )
 
-                pending_job_runs = [
-                    edge["node"] for edge in job_runs_result.data["jobRuns"]["edges"]
-                ]
-
-                if not pending_job_runs:
-                    sleep_amount = 5
-                    logger.debug(
-                        f"Waiting for Job Runs... [sleeping {sleep_amount} seconds]"
+                if not api_job_run_data or not api_job_run_data.data:
+                    logger.error(
+                        f"Job Run {db_job_run.job_run_id} not found in API, deleting from DB"
                     )
-                    sleep(sleep_amount)
+                    JobRun.objects.filter_by(job_run_id=db_job_run.job_run_id).delete()
                     continue
 
-                for job_run in pending_job_runs:
-                    logger.debug("Found pending Job Run")
-                    logger.debug(f"Job Run ID: {job_run['id']}")
-                    logger.debug(f"Job Name: {job_run['job']['name']}")
+                api_job_run = api_job_run_data.data["jobRun"]
 
-                    runner = Runner(
-                        primitive=self.primitive,
-                        job_run=job_run,
-                        max_log_size=500 * 1024,
+                logger.debug("Found pending Job Run")
+                logger.debug(f"Job Run ID: {api_job_run.get('id')}")
+                logger.debug(f"Job Name: {api_job_run.get('name')}")
+
+                runner = Runner(
+                    primitive=self.primitive,
+                    job_run=api_job_run,
+                    # max_log_size=500 * 1024,
+                )
+
+                try:
+                    runner.setup()
+                except Exception as exception:
+                    logger.exception(
+                        f"Exception while initializing runner: {exception}"
                     )
+                    self.primitive.jobs.job_run_update(
+                        id=api_job_run.get("id"),
+                        status="request_completed",
+                        conclusion="failure",
+                    )
+                    JobRun.objects.filter_by(job_run_id=api_job_run.get("id")).delete()
+                    continue
 
-                    try:
-                        runner.setup()
-                    except Exception as exception:
-                        logger.exception(
-                            f"Exception while initializing runner: {exception}"
-                        )
-                        self.primitive.jobs.job_run_update(
-                            id=job_run["id"],
-                            status="request_completed",
-                            conclusion="failure",
-                        )
-                        continue
-
-                    try:
-                        runner.execute()
-                    except Exception as exception:
-                        logger.exception(f"Exception while executing job: {exception}")
-                        self.primitive.jobs.job_run_update(
-                            id=job_run["id"],
-                            status="request_completed",
-                            conclusion="failure",
-                        )
-                    finally:
-                        runner.cleanup()
-
-                        # NOTE: also run scan here to force upload of artifacts
-                        # This should probably eventuall be another daemon?
-                        uploader.scan()
+                try:
+                    runner.execute()
+                except Exception as exception:
+                    logger.exception(f"Exception while executing job: {exception}")
+                    self.primitive.jobs.job_run_update(
+                        id=api_job_run.get("id"),
+                        status="request_completed",
+                        conclusion="failure",
+                    )
+                finally:
+                    runner.cleanup()
+
+                    # NOTE: also run scan here to force upload of artifacts
+                    # This should probably eventually be another daemon?
+                    uploader.scan()
+
+                    JobRun.objects.filter_by(
+                        job_run_id=api_job_run.get("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/commands.py

@@ -1,6 +1,7 @@
-import click
 import typing
 
+import click
+
 if typing.TYPE_CHECKING:
     from ..client import Primitive
 

primitive/agent/runner.py

@@ -2,7 +2,6 @@ import asyncio
 import os
 import re
 import shutil
-import time
 import typing
 from abc import abstractmethod
 from enum import Enum, IntEnum
@@ -12,6 +11,7 @@ from typing import Dict, List, TypedDict
 import yaml
 from loguru import logger
 
+from ..db.models import JobRun
 from ..utils.cache import get_artifacts_cache, get_logs_cache, get_sources_cache
 from ..utils.shell import env_to_dict
 
@@ -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,52 @@ class Runner:
         self.modified_env = {**self.initial_env}
 
         task_failed = False
-        conclusion = "success"
+        cancelled = False
+
         for task in self.config["executes"]:
+            # the get status check here is to ensure that if cancel is called
+            # while one task is running, we do not run any OTHER laebeled tasks
+            # THIS is required for MULTI STEP JOBS
+            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"
+        )
+
+        # FOR NONE MULTI STEP JOBS
+        # we still have to check that the job was cancelled here as well
+        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
+
+        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
+
+        conclusion = "success"
         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",
@@ -249,24 +277,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(
@@ -355,25 +383,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/agent/uploader.py

@@ -50,8 +50,8 @@ class Uploader:
                     path=file,
                     key_prefix=str(PurePath(file).relative_to(cache.parent).parent),
                 )
-            except Exception as e:
-                if "is empty" in str(e):
+            except Exception as exception:
+                if "is empty" in str(exception):
                     logger.warning(f"{file} is empty, skipping upload")
                 continue
 

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