primitive 0.2.10__tar.gz → 0.2.12__tar.gz

primitive-0.2.12/.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.12/.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.12}/.vscode/settings.json

@@ -6,6 +6,9 @@
     "Mbit",
     "mbps",
     "pkey",
+    "plutil",
+    "procs",
+    "sessionmaker",
     "testsuites",
     "untar",
     "Untaring",
@@ -13,6 +16,7 @@
     "verilog",
     "workdir"
   ],
+  "python.languageServer": "Pylance",
   "[python]": {
     "editor.defaultFormatter": "charliermarsh.ruff",
     "editor.formatOnSave": true,

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.10
+Version: 0.2.12
 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.12}/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.12}/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.12"

primitive-0.2.12/src/primitive/agent/actions.py

@@ -0,0 +1,110 @@
+import sys
+from time import sleep
+
+from loguru import logger
+
+from primitive.__about__ import __version__
+from primitive.utils.actions import BaseAction
+
+from ..db import sqlite
+from ..db.models import JobRun
+from .runner import Runner
+from .uploader import Uploader
+
+
+class Agent(BaseAction):
+    def execute(
+        self,
+    ):
+        logger.remove()
+        logger.add(
+            sink=sys.stderr,
+            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 agent")
+        logger.info(f"[*] Version: {__version__}")
+
+        # Initialize the database
+        sqlite.init()
+
+        # Create uploader
+        uploader = Uploader(primitive=self.primitive)
+
+        try:
+            while True:
+                logger.debug("Scanning for files to upload...")
+                uploader.scan()
+
+                db_job_run = JobRun.objects.first()
+
+                if not db_job_run:
+                    sleep_amount = 5
+                    logger.debug(
+                        f"No pending job runs... [sleeping {sleep_amount} seconds]"
+                    )
+                    sleep(sleep_amount)
+                    continue
+
+                api_job_run_data = self.primitive.jobs.get_job_run(
+                    id=db_job_run.job_run_id,
+                )
+
+                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"
+                    )
+                    JobRun.objects.filter_by(job_run_id=db_job_run.job_run_id).delete()
+                    continue
+
+                api_job_run = api_job_run_data.data["jobRun"]
+
+                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.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 agent...")

{primitive-0.2.10 → primitive-0.2.12}/src/primitive/agent/commands.py

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

{primitive-0.2.10 → primitive-0.2.12}/src/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-0.2.10 → primitive-0.2.12}/src/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-0.2.10 → primitive-0.2.12}/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.12}/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.12/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.12/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)