primitive 0.2.39__tar.gz → 0.2.40__tar.gz

{primitive-0.2.39 → primitive-0.2.40}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.39
+Version: 0.2.40
 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

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/__about__.py

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

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

@@ -0,0 +1,150 @@
+import sys
+from time import sleep
+from typing import Optional
+
+from loguru import logger
+
+from primitive.__about__ import __version__
+from primitive.agent.runner import Runner
+from primitive.agent.uploader import Uploader
+from primitive.utils.actions import BaseAction
+
+
+class Agent(BaseAction):
+    def start(self, job_run_id: Optional[str] = None):
+        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__}")
+
+        # TODO: tighten logic for determining if we're running in a container
+        RUNNING_IN_CONTAINER = False
+        if job_run_id is not None:
+            logger.info("Running in container...")
+            RUNNING_IN_CONTAINER = True
+
+        # Create uploader
+        uploader = Uploader(primitive=self.primitive)
+
+        try:
+            while True:
+                logger.debug("Scanning for files to upload...")
+                uploader.scan()
+
+                logger.debug("Checking for pending job runs for this device...")
+
+                # From Dylan June 30th:
+                # If passed an explicit job_run_id:
+                # - check if the JobRun exists in the API
+                # - if it does, set it to request_in_progress
+                # - if it does not, log an error and stop execution
+                # If no job_run_id is passed:
+                # - verify that this is a Node with an active Reservation
+                # - if the Reservation is active AND it has a JobRun associated with it,
+                #   then query for that JobRun
+                # - if no JobRuns are found in the API, wait for another active reservation
+                # - if a JobRun is found, set it to request_in_progress
+                # - then wait for the JobRun to be in_progress from the API
+
+                active_reservation_id = None
+                job_run_data: dict = {}
+
+                if RUNNING_IN_CONTAINER and job_run_id:
+                    job_run_result = self.primitive.jobs.get_job_run(id=job_run_id)
+                    if job_run_result.data:
+                        job_run_data = job_run_result.data.get("jobRun", {})
+                else:
+                    hardware = self.primitive.hardware.get_own_hardware_details()
+                    # fetch the latest hardware and activeReservation details
+                    if active_reservation_data := hardware["activeReservation"]:
+                        active_reservation_id = active_reservation_data.get("id", None)
+
+                    if active_reservation_id is not None:
+                        job_run_data = (
+                            self.primitive.reservations.get_job_run_for_reservation_id(
+                                reservation_id=active_reservation_id
+                            )
+                        )
+                        job_run_id = job_run_data.get("id", None)
+
+                if (
+                    len(job_run_data.keys()) == 0
+                    or not job_run_data.get("id")
+                    or job_run_id is None
+                ):
+                    if RUNNING_IN_CONTAINER:
+                        logger.info("Running in container, exiting due to no JobRun.")
+                        break
+                    logger.debug("No pending Job Run found, sleeping...")
+                    sleep(5)
+                    continue
+
+                logger.debug("Found pending Job Run")
+                logger.debug(f"Job Run ID: {job_run_data.get('id')}")
+                logger.debug(f"Job Name: {job_run_data.get('job').get('name')}")
+
+                logger.info(
+                    f"Setting JobRun {job_run_data.get('job').get('name')} to request_in_progress"
+                )
+                # we are setting to request_in_progress here which puts a started_at time on the JobRun in the API's database
+                # any time spent pulling Git repositories, setting up, etc, counts as compute time
+                job_run_result = self.primitive.jobs.job_run_update(
+                    id=job_run_id, status="request_in_progress"
+                )
+
+                while job_run_data["status"] != "in_progress":
+                    logger.info(
+                        f"Waiting for JobRun {job_run_data.get('name')} to be in_progress"
+                    )
+                    sleep(1)
+                    job_run_result = self.primitive.jobs.get_job_run(id=job_run_id)
+                    if job_run_result.data is not None:
+                        job_run_data = job_run_result.data.get("jobRun", {})
+
+                runner = Runner(
+                    primitive=self.primitive,
+                    job_run=job_run_data,
+                )
+
+                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_job_run()
+                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 eventually be another daemon?
+                    uploader.scan()
+
+                if RUNNING_IN_CONTAINER:
+                    logger.info("Running in container, exiting after job run")
+                    break
+
+                sleep(5)
+        except KeyboardInterrupt:
+            logger.info("Stopping primitive agent...")

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/agent/commands.py

@@ -12,4 +12,4 @@ if typing.TYPE_CHECKING:
 def cli(context, job_run_id: typing.Optional[str] = None):
     """agent"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    primitive.agent.execute(job_run_id=job_run_id)
+    primitive.agent.start(job_run_id=job_run_id)

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/agent/runner.py

@@ -10,10 +10,10 @@ 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.logging import fmt, log_context
-from ..utils.shell import env_to_dict
+from primitive.utils.cache import get_artifacts_cache, get_logs_cache, get_sources_cache
+from primitive.utils.logging import fmt, log_context
+from primitive.utils.psutil import kill_process_and_children
+from primitive.utils.shell import env_to_dict
 
 try:
     from yaml import CLoader as Loader
@@ -109,27 +109,21 @@ class Runner:
             self.job_settings["rootDirectory"]
         )
 
-        db_config = self.job_settings.get("config", None)
-        if db_config:
-            logger.info(f"Using job config from database for {self.job['slug']}")
-            self.config = db_config
-        else:
-            # Attempt to parse the job yaml file
-            job_filename = self.job_settings["repositoryFilename"]
-            logger.info(f"Scanning directory for job file {job_filename}")
+        job_filename = self.job_settings["repositoryFilename"]
+        logger.info(f"Scanning directory for job file {job_filename}")
 
-            job_config_file = Path(self.source_dir / ".primitive" / job_filename)
+        job_config_file = Path(self.source_dir / ".primitive" / job_filename)
 
-            if job_config_file.exists():
-                logger.info(
-                    f"Found job description for {self.job['slug']} at {job_config_file}"
-                )
-                self.config = yaml.load(open(job_config_file, "r"), Loader=Loader)
-            else:
-                logger.error(
-                    f"No job description with matching filename '{job_filename}' found"
-                )
-                raise FileNotFoundError
+        if job_config_file.exists():
+            logger.info(
+                f"Found job description for {self.job['slug']} at {job_config_file}"
+            )
+            self.config = yaml.load(open(job_config_file, "r"), Loader=Loader)
+        else:
+            logger.error(
+                f"No job description with matching filename '{job_filename}' found"
+            )
+            raise FileNotFoundError
 
         # Setup initial process environment
         self.initial_env = os.environ
@@ -147,15 +141,11 @@ class Runner:
         )
 
     @log_context(label="execute")
-    def execute(self) -> None:
-        logger.info(f"Executing {self.job['slug']} job")
-        self.primitive.jobs.job_run_update(
-            self.job_run["id"], status="request_in_progress"
-        )
-
+    def execute_job_run(self) -> None:
         self.modified_env = {**self.initial_env}
         task_failed = False
         cancelled = False
+        timed_out = False
 
         for task in self.config["executes"]:
             # Everything inside this loop should be contextualized with the task label
@@ -171,6 +161,9 @@ class Runner:
                 if status_value == "completed" and conclusion_value == "cancelled":
                     cancelled = True
                     break
+                if status_value == "completed" and conclusion_value == "timed_out":
+                    timed_out = True
+                    break
 
                 # Everything within this block should be contextualized as user logs
                 with logger.contextualize(type="user"):
@@ -186,11 +179,17 @@ class Runner:
             conclusion_value = status.data["jobRun"]["conclusion"]
             if status_value == "completed" and conclusion_value == "cancelled":
                 cancelled = True
+            if status_value == "completed" and conclusion_value == "timed_out":
+                timed_out = True
 
             if cancelled:
                 logger.warning("Job cancelled by user")
                 return
 
+            if timed_out:
+                logger.error("Job timed out")
+                return
+
             conclusion = "success"
             if task_failed:
                 conclusion = "failure"
@@ -270,9 +269,17 @@ class Runner:
                 stderr=asyncio.subprocess.PIPE,
             )
 
-            JobRun.objects.filter_by(job_run_id=self.job_run["id"]).update(
-                {"pid": process.pid}
-            )
+            try:
+                await self.primitive.jobs.ajob_run_update(
+                    self.job_run["id"],
+                    parent_pid=process.pid,
+                )
+            except ValueError:
+                logger.error(
+                    f"Failed to update job run {self.job_run['id']} with process PID {process.pid}"
+                )
+                kill_process_and_children(pid=process.pid)
+                return False
 
             stdout_failed, stderr_failed = await asyncio.gather(
                 self.log_cmd(
@@ -289,10 +296,6 @@ class Runner:
                 f"Finished executing command {i + 1}/{len(commands)}: {cmd} with return code {returncode}"
             )
 
-            JobRun.objects.filter_by(job_run_id=self.job_run["id"]).update(
-                {"pid": None}
-            )
-
             if returncode > 0:
                 logger.error(
                     f"Task {task['label']} failed on '{cmd}' with return code {returncode}"

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/client.py

@@ -1,8 +1,9 @@
+from typing import Optional
+
 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
@@ -12,11 +13,11 @@ from .files.actions import Files
 from .git.actions import Git
 from .hardware.actions import Hardware
 from .jobs.actions import Jobs
+from .monitor.actions import Monitor
 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
 
 

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/jobs/actions.py

@@ -115,6 +115,7 @@ class Jobs(BaseAction):
         conclusion: str = None,
         file_ids: Optional[List[str]] = [],
         number_of_files_produced: Optional[int] = None,
+        parent_pid: Optional[int] = None,
     ):
         mutation = gql(job_run_update_mutation)
         input = {"id": id}
@@ -126,12 +127,42 @@ class Jobs(BaseAction):
             input["files"] = file_ids
         if number_of_files_produced is not None:
             input["numberOfFilesProduced"] = number_of_files_produced
+        if parent_pid is not None:
+            input["parentPid"] = parent_pid
         variables = {"input": input}
         result = self.primitive.session.execute(
             mutation, variable_values=variables, get_execution_result=True
         )
         return result
 
+    @guard
+    async def ajob_run_update(
+        self,
+        id: str,
+        status: str = None,
+        conclusion: str = None,
+        file_ids: Optional[List[str]] = [],
+        number_of_files_produced: Optional[int] = None,
+        parent_pid: Optional[int] = None,
+    ):
+        mutation = gql(job_run_update_mutation)
+        input = {"id": id}
+        if status:
+            input["status"] = status
+        if conclusion:
+            input["conclusion"] = conclusion
+        if file_ids and len(file_ids) > 0:
+            input["files"] = file_ids
+        if number_of_files_produced is not None:
+            input["numberOfFilesProduced"] = number_of_files_produced
+        if parent_pid is not None:
+            input["parentPid"] = parent_pid
+        variables = {"input": input}
+        result = await self.primitive.session.execute_async(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result
+
     @guard
     def github_access_token_for_job_run(self, job_run_id: str):
         query = gql(github_app_token_for_job_run_query)

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/jobs/graphql/fragments.py

@@ -50,5 +50,6 @@ fragment JobRunStatusFragment on JobRun {
     id
     status
     conclusion
+    parentPid
 }
 """

primitive-0.2.40/src/primitive/monitor/actions.py

@@ -0,0 +1,196 @@
+import sys
+from time import sleep
+from typing import Optional
+
+from loguru import logger
+
+from primitive.__about__ import __version__
+from primitive.utils.actions import BaseAction
+from primitive.utils.exceptions import P_CLI_100, P_CLI_101
+from primitive.utils.psutil import kill_process_and_children
+
+MAX_GET_STATUS_TIMEOUT = 30
+
+
+class Monitor(BaseAction):
+    def start(self, job_run_id: Optional[str] = None):
+        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 monitor")
+        logger.info(f"Version: {__version__}")
+
+        # TODO: tighten logic for determining if we're running in a container
+        RUNNING_IN_CONTAINER = False
+        if job_run_id is not None:
+            logger.info("Running in container...")
+            RUNNING_IN_CONTAINER = True
+
+        # can't check if if it is a container
+        if not RUNNING_IN_CONTAINER:
+            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)
+
+        # From Dylan on June 30th:
+        # If passed an explicit job_run_id we know it is running in a container.
+        # If no job_run_id is passed, we need to check that this device has an active reservation.
+        # Fetch the active reservations. If it exists AND has a JobRun associated with it:
+        # - check if the JobRun exists in the API
+        #   - if it does exist, check if it is already running
+        #   - if it is in status [pending, request_in_progress, in_progress] where we should wait for the PID
+        #   - if it is in status [request_completed, completed] and there is a PID, kill it
+        # Finally, if running in a container, kill the process.
+        # Else, wait for a new active reservation to be created.
+
+        try:
+            active_reservation_data = None
+            previous_reservation_id = None
+            active_reservation_id = None
+
+            while True:
+                # this block determines if there is a reservation at all
+                # handles cleanup of old reservations
+                # obtains an active JobRun's ID
+                if not RUNNING_IN_CONTAINER:
+                    hardware = self.primitive.hardware.get_own_hardware_details()
+                    # fetch the latest hardware and activeReservation details
+                    if active_reservation_data := hardware["activeReservation"]:
+                        active_reservation_id = active_reservation_data.get("id", None)
+                        if previous_reservation_id is None:
+                            previous_reservation_id = active_reservation_id
+                    else:
+                        active_reservation_data = None
+                        active_reservation_id = None
+
+                    # if there is no activeReservation or previous reservation, sync + sleep
+                    if (
+                        active_reservation_data is None
+                        and active_reservation_id is None
+                        and previous_reservation_id is None
+                    ):
+                        self.primitive.hardware.check_in_http(
+                            is_available=True, is_online=True
+                        )
+                        self.primitive.hardware._sync_children(hardware=hardware)
+
+                        sleep_amount = 5
+                        logger.info(
+                            f"No active reservation found... [sleeping {sleep_amount} seconds]"
+                        )
+                        sleep(sleep_amount)
+                        continue
+
+                    # if there is a previous_reservation_id but no activeReservation, cleanup
+                    elif active_reservation_data is None and previous_reservation_id:
+                        logger.info(
+                            f"Cleaning up previous reservation {previous_reservation_id}..."
+                        )
+                        self.primitive.provisioning.remove_reservation_authorized_keys(
+                            reservation_id=previous_reservation_id
+                        )
+                        job_run_data = (
+                            self.primitive.reservations.get_job_run_for_reservation_id(
+                                reservation_id=previous_reservation_id
+                            )
+                        )
+                        job_run_id = job_run_data.get("id")
+                        previous_reservation_id = None
+
+                    # if we are on the new reservation
+                    elif (
+                        (previous_reservation_id is not None)
+                        and (active_reservation_id is not None)
+                        and (previous_reservation_id == active_reservation_id)
+                    ):
+                        self.primitive.provisioning.add_reservation_authorized_keys(
+                            reservation_id=active_reservation_id
+                        )
+
+                    # we have an active reservation, check if we have JobRuns attached to it
+                    if active_reservation_id is not None:
+                        logger.info(f"Active Reservation ID: {active_reservation_id}")
+                        job_run_data = (
+                            self.primitive.reservations.get_job_run_for_reservation_id(
+                                reservation_id=active_reservation_id
+                            )
+                        )
+                        job_run_id = job_run_data.get("id")
+
+                    # Golden state for normal reservation
+                    if not job_run_id and active_reservation_id:
+                        self.primitive.hardware.check_in_http(
+                            is_available=False, is_online=True
+                        )
+                        sleep_amount = 5
+                        logger.info(
+                            f"Waiting for Job Runs... [sleeping {sleep_amount} seconds]"
+                        )
+                        sleep(sleep_amount)
+                        continue
+
+                # job_run_data can come from 3 places:
+                # 1. an explicitly passed job_run_id
+                # 2. the previous reservation has an job_run_id (kill old PIDs)
+                # 3. the active reservation has an job_run_id (check status)
+                while job_run_id:
+                    status_result = self.primitive.jobs.get_job_status(id=job_run_id)
+                    get_status_timeout = 0
+                    sleep_amount = 5
+
+                    while get_status_timeout < MAX_GET_STATUS_TIMEOUT:
+                        if not status_result or not status_result.data:
+                            logger.error(
+                                f"Error fetching job status for Job Run {job_run_id}. Retrying... [sleeping {sleep_amount} seconds]"
+                            )
+                            get_status_timeout += sleep_amount
+                            sleep(sleep_amount)
+                            continue
+                        else:
+                            break
+
+                    if not status_result or not status_result.data:
+                        raise P_CLI_101()
+
+                    status_value = status_result.data["jobRun"]["status"]
+                    parent_pid = status_result.data["jobRun"]["parentPid"]
+
+                    if status_value == "completed":
+                        logger.info(
+                            f"Job run {job_run_id} is completed. Killing children if they exist."
+                        )
+                        if parent_pid is not None:
+                            kill_process_and_children(pid=parent_pid)
+                        status_value = None
+                        job_run_id = None
+                    else:
+                        logger.info(
+                            f"Job Run {job_run_id} with Status {status_value} with PID {parent_pid}. [sleeping {sleep_amount} seconds]"
+                        )
+                        sleep(sleep_amount)
+                        continue
+
+        except KeyboardInterrupt:
+            logger.info("Stopping primitive monitor...")
+            try:
+                if not RUNNING_IN_CONTAINER:
+                    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 monitor.")
+                logger.error(str(exception))
+            sys.exit()

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/reservations/actions.py

@@ -176,3 +176,27 @@ class Reservations(BaseAction):
             )
 
         return reservation_result
+
+    @guard
+    def get_job_run_for_reservation_id(self, reservation_id: str) -> dict:
+        if not reservation_id:
+            logger.error("No reservation ID provided.")
+            return {}
+
+        job_runs_for_reservation = self.primitive.jobs.get_job_runs(
+            first=1,
+            reservation_id=reservation_id,
+        )
+
+        while job_runs_for_reservation is None or job_runs_for_reservation.data is None:
+            logger.error("Error fetching job runs.")
+            sleep_amount = 5
+            logger.info(f"Error fetching job runs... [sleeping {sleep_amount} seconds]")
+            sleep(sleep_amount)
+            continue
+
+        if not job_runs_for_reservation.data["jobRuns"]["edges"]:
+            logger.error("No job runs found for the given reservation ID.")
+            return {}
+
+        return job_runs_for_reservation.data["jobRuns"]["edges"][0]["node"]

{primitive-0.2.39 → primitive-0.2.40}/src/primitive/utils/exceptions.py

@@ -10,3 +10,14 @@ class P_CLI_100(Exception):
 
     def __str__(self):
         return f"{self.codename}: {self.message}"
+
+
+@dataclass
+class P_CLI_101(Exception):
+    """Could Not Get Status for JobRun"""
+
+    codename: str = "P_CLI_101"
+    message: str = "Could Not Get Status for JobRun"
+
+    def __str__(self):
+        return f"{self.codename}: {self.message}"

primitive-0.2.40/src/primitive/utils/psutil.py

@@ -0,0 +1,26 @@
+import psutil
+from loguru import logger
+
+
+def kill_process_and_children(pid: int) -> bool:
+    """Kill a process and all its children."""
+    try:
+        try:
+            parent = psutil.Process(pid)
+            logger.info(f"Process PID {parent.pid} found.")
+        except psutil.NoSuchProcess:
+            logger.info("Process not found")
+            return False
+
+        children = parent.children(recursive=True)
+
+        for child in children:
+            logger.info(f"Killing child process {child.pid}...")
+            child.kill()
+
+        logger.info(f"Killing parent process {parent.pid}...")
+        parent.kill()
+        return True
+    except psutil.NoSuchProcess:
+        logger.warning(f"Process with PID {pid} not found.")
+        return False