primitive 0.2.10__tar.gz → 0.2.11__tar.gz

primitive-0.2.11/.github/workflows/pyright.yml

@@ -0,0 +1,20 @@
+name: Pyright
+
+on:
+  workflow_dispatch:
+  push:
+    branches: ["*"]
+  pull_request:
+    branches: ["*"]
+
+jobs:
+  build:
+    name: Pyright
+    runs-on: ubuntu-latest
+    strategy:
+      max-parallel: 4
+      matrix:
+        python-version: [3.13]
+
+    steps:
+      - uses: jakebailey/pyright-action@v1

primitive-0.2.11/.vscode/extensions.json

@@ -0,0 +1,13 @@
+{
+  "recommendations": [
+    "charliermarsh.ruff",
+    "github.copilot-chat",
+    "ms-python.debugpy",
+    "ms-python.python",
+    "ms-python.vscode-pylance",
+    "ms-vscode.remote-explorer",
+    "ms-vscode-remote.remote-ssh",
+    "ms-vsliveshare.vsliveshare",
+    "streetsidesoftware.code-spell-checker"
+  ]
+}

{primitive-0.2.10 → primitive-0.2.11}/.vscode/settings.json

@@ -11,8 +11,11 @@
     "Untaring",
     "VERIBLE",
     "verilog",
-    "workdir"
+    "workdir",
+    "plutil",
+    "procs"
   ],
+  "python.languageServer": "Pylance",
   "[python]": {
     "editor.defaultFormatter": "charliermarsh.ruff",
     "editor.formatOnSave": true,

{primitive-0.2.10 → primitive-0.2.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.10
+Version: 0.2.11
 Project-URL: Documentation, https://github.com//primitivecorp/primitive-cli#readme
 Project-URL: Issues, https://github.com//primitivecorp/primitive-cli/issues
 Project-URL: Source, https://github.com//primitivecorp/primitive-cli
@@ -23,9 +23,11 @@ Requires-Dist: gql[all]
 Requires-Dist: loguru
 Requires-Dist: paramiko[invoke]
 Requires-Dist: primitive-pal==0.1.4
+Requires-Dist: psutil>=7.0.0
 Requires-Dist: pyyaml
 Requires-Dist: rich>=13.9.4
 Requires-Dist: speedtest-cli
+Requires-Dist: sqlalchemy>=2.0.40
 Description-Content-Type: text/markdown
 
 # primitive

{primitive-0.2.10 → primitive-0.2.11}/pyproject.toml

@@ -35,6 +35,8 @@ dependencies = [
   "paramiko[invoke]",
   "speedtest-cli",
   "rich>=13.9.4",
+  "sqlalchemy>=2.0.40",
+  "psutil>=7.0.0",
 ]
 
 [tool.uv]
@@ -79,3 +81,17 @@ tests = ["tests", "*/primitive/tests"]
 
 [tool.coverage.report]
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
+
+[tool.pyright]
+include = ["."]
+exclude = ["**/__pycache__"]
+defineConstant = { DEBUG = true }
+stubPath = "stubs"
+
+reportMissingImports = true
+reportMissingTypeStubs = false
+
+pythonVersion = "3.13"
+pythonPlatform = "Linux"
+
+executionEnvironments = [{ root = "src" }, { root = "tests" }]

{primitive-0.2.10 → primitive-0.2.11}/src/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.11"

{primitive-0.2.10 → primitive-0.2.11}/src/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-0.2.10 → primitive-0.2.11}/src/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",
@@ -249,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(
@@ -355,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-0.2.10 → primitive-0.2.11}/src/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-0.2.10 → primitive-0.2.11}/src/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-0.2.11/src/primitive/daemons/actions.py

@@ -0,0 +1,61 @@
+import platform
+import typing
+from typing import Dict, Optional, List
+
+if typing.TYPE_CHECKING:
+    from ..client import Primitive
+
+from .launch_agents import LaunchAgent
+from .launch_service import LaunchService
+from ..utils.daemons import Daemon
+
+
+class Daemons:
+    def __init__(self, primitive) -> None:
+        self.primitive: Primitive = primitive
+        self.os_family = platform.system()
+
+        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-0.2.11/src/primitive/daemons/commands.py

@@ -0,0 +1,111 @@
+import click
+
+import typing
+from typing import Optional
+from .ui import render_daemon_list
+
+from loguru import logger
+
+if typing.TYPE_CHECKING:
+    from ..client import Primitive
+
+
+@click.group()
+@click.pass_context
+def cli(context):
+    """Daemon"""
+    pass
+
+
+@cli.command("install")
+@click.pass_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")
+    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
+@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")
+    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
+@click.argument(
+    "name",
+    type=str,
+    required=False,
+)
+def stop_daemon_command(context, name: Optional[str]):
+    """Stop primitive Daemon"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    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
+@click.argument(
+    "name",
+    type=str,
+    required=False,
+)
+def start_daemon_command(context, name: Optional[str]):
+    """Start primitive Daemon"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    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
+@click.argument(
+    "name",
+    type=str,
+    required=True,
+)
+def log_daemon_command(context, name: str):
+    """Logs from primitive Daemon"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    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)