primitive 0.2.39__py3-none-any.whl → 0.2.42__py3-none-any.whl

primitive/__about__.py

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

primitive/agent/actions.py

@@ -7,13 +7,11 @@ from loguru import logger
 from primitive.__about__ import __version__
 from primitive.agent.runner import Runner
 from primitive.agent.uploader import Uploader
-from primitive.db.models import JobRun
-from primitive.db.sqlite import wait_for_db
 from primitive.utils.actions import BaseAction
 
 
 class Agent(BaseAction):
-    def execute(self, job_run_id: Optional[str] = None):
+    def start(self, job_run_id: Optional[str] = None):
         logger.remove()
         logger.add(
             sink=sys.stderr,
@@ -31,9 +29,6 @@ class Agent(BaseAction):
             logger.info("Running in container...")
             RUNNING_IN_CONTAINER = True
 
-        # Wait for monitor to make database
-        wait_for_db()
-
         # Create uploader
         uploader = Uploader(primitive=self.primitive)
 
@@ -42,37 +37,79 @@ class Agent(BaseAction):
                 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)
+                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
 
-                api_job_run_data = self.primitive.jobs.get_job_run(
-                    id=db_job_run.job_run_id,
+                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"
                 )
 
-                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"
+                while job_run_data["status"] != "in_progress":
+                    logger.info(
+                        f"Waiting for JobRun {job_run_data.get('name')} to be in_progress"
                     )
-                    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')}")
+                    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=api_job_run,
-                    # max_log_size=500 * 1024,
+                    job_run=job_run_data,
                 )
 
                 try:
@@ -82,19 +119,18 @@ class Agent(BaseAction):
                         f"Exception while initializing runner: {exception}"
                     )
                     self.primitive.jobs.job_run_update(
-                        id=api_job_run.get("id"),
+                        id=job_run_id,
                         status="request_completed",
                         conclusion="failure",
                     )
-                    JobRun.objects.filter_by(job_run_id=api_job_run.get("id")).delete()
                     continue
 
                 try:
-                    runner.execute()
+                    runner.execute_job_run()
                 except Exception as exception:
                     logger.exception(f"Exception while executing job: {exception}")
                     self.primitive.jobs.job_run_update(
-                        id=api_job_run.get("id"),
+                        id=job_run_id,
                         status="request_completed",
                         conclusion="failure",
                     )

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/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/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/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/jobs/graphql/fragments.py

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

primitive/monitor/actions.py

@@ -2,14 +2,14 @@ import sys
 from time import sleep
 from typing import Optional
 
-import psutil
 from loguru import logger
 
 from primitive.__about__ import __version__
-from primitive.db.models import JobRun
-from primitive.db.sqlite import init, wait_for_db
 from primitive.utils.actions import BaseAction
-from primitive.utils.exceptions import P_CLI_100
+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):
@@ -44,194 +44,143 @@ class Monitor(BaseAction):
                 logger.exception(f"Error checking in hardware: {exception}")
                 sys.exit(1)
 
-        # Initialize the database
-        init()
-        wait_for_db()
+        # 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:
-            if job_run_id is not None:
-                if not JobRun.objects.filter_by(job_run_id=job_run_id).exists():
-                    logger.debug(f"Job run {job_run_id} does not exist in database.")
-                    logger.debug(f"Creating job run in database: {job_run_id}")
-                    JobRun.objects.create(job_run_id=job_run_id, pid=None)
-
+            active_reservation_data = None
+            previous_reservation_id = None
             active_reservation_id = None
-            active_reservation_pk = None
 
             while True:
-                # FIRST, check for jobs in the database that are running
-                db_job_runs = JobRun.objects.all()
-                for job_run in db_job_runs:
-                    # first check if the job run completed already
-                    status = self.primitive.jobs.get_job_status(job_run.job_run_id)
-                    if status is None or status.data is None:
-                        logger.error(
-                            f"Error fetching status of <JobRun {job_run.job_run_id}>."
+                # 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
                         )
-                        continue
+                        self.primitive.hardware._sync_children(hardware=hardware)
 
-                    status_value = status.data["jobRun"]["status"]
-                    if status_value == "completed":
-                        logger.debug(
-                            f"Job run {job_run.job_run_id} is completed. Removing from database."
+                        sleep_amount = 5
+                        logger.info(
+                            f"No active reservation found... [sleeping {sleep_amount} seconds]"
                         )
-                        JobRun.objects.filter_by(job_run_id=job_run.job_run_id).delete()
+                        sleep(sleep_amount)
                         continue
 
-                    if job_run.pid is None:
-                        pid_sleep_amount = 0.1
-                        logger.debug(
-                            f"Job run {job_run.job_run_id} has no PID. Agent has not started."
+                    # 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}..."
                         )
-                        logger.debug(
-                            f"Sleeping {pid_sleep_amount} seconds before checking again..."
+                        self.primitive.provisioning.remove_reservation_authorized_keys(
+                            reservation_id=previous_reservation_id
                         )
-                        sleep(pid_sleep_amount)
-                        pid_sleep_amount += 0.1
-                        continue
-
-                    logger.debug(
-                        f"Checking process PID {job_run.pid} for JobRun {job_run.job_run_id}..."
-                    )
-
-                    status = self.primitive.jobs.get_job_status(job_run.job_run_id)
-                    if status is None or status.data is None:
-                        logger.error(
-                            f"Error fetching status of <JobRun {job_run.job_run_id}>."
+                        job_run_data = (
+                            self.primitive.reservations.get_job_run_for_reservation_id(
+                                reservation_id=previous_reservation_id
+                            )
                         )
-                        continue
-
-                    status_value = status.data["jobRun"]["status"]
-                    conclusion_value = status.data["jobRun"]["conclusion"]
-
-                    logger.debug(f"- Status: {status_value}")
-                    logger.debug(f"- Conclusion: {conclusion_value}")
-
-                    try:
-                        parent = psutil.Process(job_run.pid)
-                    except psutil.NoSuchProcess:
-                        logger.debug("Process not found")
-                        continue
-
-                    children = parent.children(recursive=True)
-
-                    if status_value == "completed" and conclusion_value == "cancelled":
-                        logger.warning("Job cancelled by user")
-                        for child in children:
-                            logger.debug(f"Killing child process {child.pid}...")
-                            child.kill()
-
-                        logger.debug(f"Killing parent process {parent.pid}...")
-                        parent.kill()
-
-                    if status != "completed":
-                        sleep(1)
-                        continue
-
-                if RUNNING_IN_CONTAINER:
-                    if len(db_job_runs) == 0:
-                        # if we get here and we're running in a container,
-                        # it means the job run is complete and there is nothing left in the database
-                        # so we can exit
-                        logger.debug("Running in container, initial job complete.")
-                        sys.exit(0)
-                    else:
-                        continue
-
-                # Second, check for active reservations
-                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
+                        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)
                     ):
-                        logger.info("New reservation for this hardware.")
-                        active_reservation_id = hardware["activeReservation"]["id"]
-                        active_reservation_pk = hardware["activeReservation"]["pk"]
-                        reservation_number = hardware["activeReservation"][
-                            "reservationNumber"
-                        ]
-                        logger.debug("Active Reservation:")
-                        logger.info(f"Reservation Number: {reservation_number}")
-                        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
                         )
 
-                if not active_reservation_id:
-                    self.primitive.hardware.check_in_http(
-                        is_available=True, is_online=True
-                    )
-                    logger.debug("Syncing children...")
-                    self.primitive.hardware._sync_children(hardware=hardware)
+                    # 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")
 
-                    sleep_amount = 5
-                    logger.debug(
-                        f"No active reservation found... [sleeping {sleep_amount} seconds]"
-                    )
-                    sleep(sleep_amount)
-                    continue
-                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."
+                    # 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
                         )
-                        self.primitive.provisioning.remove_reservation_authorized_keys(
-                            reservation_id=active_reservation_id
+                        sleep_amount = 5
+                        logger.info(
+                            f"Waiting for Job Runs... [sleeping {sleep_amount} seconds]"
                         )
-                        active_reservation_id = None
-                        active_reservation_pk = None
-
-                # Third, see if the active reservation has any pending job runs
-                job_runs_for_reservation = self.primitive.jobs.get_job_runs(
-                    status="pending", first=1, reservation_id=active_reservation_id
-                )
-
-                if (
-                    job_runs_for_reservation is None
-                    or job_runs_for_reservation.data is None
-                ):
-                    logger.error("Error fetching job runs.")
+                        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
-                    logger.debug(
-                        f"Error fetching job runs... [sleeping {sleep_amount} seconds]"
-                    )
-                    sleep(sleep_amount)
-                    continue
 
-                pending_job_runs = [
-                    edge["node"]
-                    for edge in job_runs_for_reservation.data["jobRuns"]["edges"]
-                ]
+                    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 pending_job_runs:
-                    self.primitive.hardware.check_in_http(
-                        is_available=False, is_online=True
-                    )
-                    sleep_amount = 5
-                    logger.debug(
-                        f"Waiting for Job Runs... [sleeping {sleep_amount} seconds]"
-                    )
-                    sleep(sleep_amount)
-                    continue
-
-                # If we did find a pending job run, check if it exists in the database
-                # and create it if it doesn't.
-                # This will trigger the agent to start the job run.
-                job_run = pending_job_runs[0]
-                if not JobRun.objects.filter_by(job_run_id=job_run["id"]).exists():
-                    JobRun.objects.create(job_run_id=job_run["id"], pid=None)
-                    logger.debug(f"Creating job run in database: {job_run['id']}")
+                    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...")

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/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/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

{primitive-0.2.39.dist-info → primitive-0.2.42.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.39
+Version: 0.2.42
 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
@@ -26,7 +26,6 @@ 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.39.dist-info → primitive-0.2.42.dist-info}/RECORD

@@ -1,11 +1,11 @@
-primitive/__about__.py,sha256=ZQkflfute36QGfeP2TGex_AxLvyIYgASuKv02D9bnfM,130
+primitive/__about__.py,sha256=Vl72jTU-bFf1jFAZpJ-xozZlZLbpzZ7x5f0FzuwRdFc,130
 primitive/__init__.py,sha256=bwKdgggKNVssJFVPfKSxqFMz4IxSr54WWbmiZqTMPNI,106
 primitive/cli.py,sha256=g7EtHI9MATAB0qQu5w-WzbXtxz_8zu8z5E7sETmMkKU,2509
-primitive/client.py,sha256=h8WZVnQylVe0vbpuyC8YZHl2JyITSPC-1HbUcmrE5pc,3623
+primitive/client.py,sha256=RMF46F89oK82gfZH6Bf0WZrhXPUu01pbieSO_Vcuoc4,3624
 primitive/agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-primitive/agent/actions.py,sha256=vq4_CZy7Mn37jUg-JPuJlMLtcx4sMXp6tUNvXGZOOyQ,4128
-primitive/agent/commands.py,sha256=kqa-PGqmzS-APd4BSMAkX4l8SdK5N1PRQwq9S8jBjsw,390
-primitive/agent/runner.py,sha256=8q9xjQgcpXvQiUDdW9O6aS08lrJt9M_0_r9TZUAOdgw,15483
+primitive/agent/actions.py,sha256=zQuxP7oGZ9PRIJLlJ5dhza-Vcg29gffRreCylDfYTtg,6365
+primitive/agent/commands.py,sha256=o847pK7v7EWQGG67tky6a33qtwoutX6LZrP2FIS_NOk,388
+primitive/agent/runner.py,sha256=wkI7ANC9ob99Fy5UWJAj__Ah7f7gwMgrBJ_LgpWZf4A,15638
 primitive/agent/uploader.py,sha256=ZzrzsajNBogwEC7mT6Ejy0h2Jd9axMYGzt9pbCvVMlk,3171
 primitive/auth/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 primitive/auth/actions.py,sha256=9NIEXJ1BNJutJs6AMMSjMN_ziONUAUhY_xHwojYJCLA,942
@@ -18,9 +18,6 @@ primitive/daemons/commands.py,sha256=Xt4qFymNrDLdHJhRnEH_4Re-2xX6w1OT-chV9k7dFCs
 primitive/daemons/launch_agents.py,sha256=cCxsvCBFmGlKOiBINsKi_NxcozLC8yPw3w6pxqnz-qI,7803
 primitive/daemons/launch_service.py,sha256=iuklHeuEqadlf8U1n9xFg4ZG1EKdK2jyaPI-VTlpJ4I,7907
 primitive/daemons/ui.py,sha256=Af3OJWJ0jdGlb1nfA5yaGYdhBEqqpM8zP2U2vUQdCbw,1236
-primitive/db/base.py,sha256=mH7f2d_jiyxJSSx9Gk53QBXRa3LiKBsBjkFgvmtH1WA,83
-primitive/db/models.py,sha256=GfnJdAq4Tb68CI4BKAuJDZVqioGavveaAHbCPeLNngw,2840
-primitive/db/sqlite.py,sha256=ZtSi8Z0wiA5NEvgnP7udSwZnoUMLp0qMWfNZNC3htuI,2116
 primitive/exec/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 primitive/exec/actions.py,sha256=4d_TCjNDcVFoZ9Zw7ZuBa6hKMv2Xzm7_UX_8wcX1aSk,4124
 primitive/exec/commands.py,sha256=66LO2kkJC-ynNZQpUCXv4Ol15QoacdSZAHblePDcmLo,510
@@ -51,13 +48,13 @@ primitive/hardware/graphql/fragments.py,sha256=R1J75AQZNr6FSaaRzOn_IyCFIEaSb_YaU
 primitive/hardware/graphql/mutations.py,sha256=_4Hkbfik9Ron4T-meulu6T-9FR_BZjyPNwn745MPksU,1484
 primitive/hardware/graphql/queries.py,sha256=I86uLuOSjHSph11Y5MVCYko5Js7hoiEZ-cEoPTc4J-k,1392
 primitive/jobs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-primitive/jobs/actions.py,sha256=Fx2cPc1x09nRasOVtjhPjNRJ-jNoi3RJhXqC3verD9s,5444
+primitive/jobs/actions.py,sha256=Fej-rpdfIiHjs0SEFY2ZcMn64AWHn6kD90RC5wtRM3Q,6531
 primitive/jobs/commands.py,sha256=MxPCkBEYW_eLNqgCRYeyj7ZcLOFAWfpVZlqDR2Y_S0o,830
 primitive/jobs/graphql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-primitive/jobs/graphql/fragments.py,sha256=flBi4REkfPX3sao22McmjclK-lFsqiXQa21YKmTbxu8,653
+primitive/jobs/graphql/fragments.py,sha256=Ldhsc-NFHZGY8_1FjO2uZIdOL5UUGZjsldHea_GLJ_g,667
 primitive/jobs/graphql/mutations.py,sha256=8ASvCmwQh7cMeeiykOdYaYVryG8FRIuVF6v_J8JJZuw,219
 primitive/jobs/graphql/queries.py,sha256=ZxNmm-WovytbggNuKRnwa0kc26T34_0yhqkoqx-2uj0,1736
-primitive/monitor/actions.py,sha256=l8J8h_GyCmFqmf6hRL34m-tc04cCW76_VYlsCCDk5IE,11120
+primitive/monitor/actions.py,sha256=7T205EU2ULvGQmiAJy99qG5vXx_7xVpjV_P6X6Osso8,9295
 primitive/monitor/commands.py,sha256=VDlEL_Qpm_ysHxug7VpI0cVAZ0ny6AS91Y58D7F1zkU,409
 primitive/organizations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 primitive/organizations/actions.py,sha256=kVHOhG1oS2sI5p8uldSo5L-RUZsnG36eaulVuKLyZ-M,1863
@@ -78,7 +75,7 @@ primitive/provisioning/actions.py,sha256=IYZYAbtomtZtlkqDaBxx4e7PFKGkRNqek_tABH6
 primitive/provisioning/graphql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 primitive/provisioning/graphql/queries.py,sha256=cBtuKa6shoatYZfKSnQoPJP6B8g8y3QhFqJ_pkvMcG0,134
 primitive/reservations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-primitive/reservations/actions.py,sha256=FiodRTVUgGgFfoksnN9W0XNdGTd2AxPJTfUrZbmQ0_g,6179
+primitive/reservations/actions.py,sha256=AYAWJmUItXTb7n1zyFt0Oov09rpBrLy6yIYHPpg69gM,7075
 primitive/reservations/commands.py,sha256=LFRoV59QGgWIjBdrGjJdffHugg8TLe0Fwlcyu_JaTkk,2369
 primitive/reservations/graphql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 primitive/reservations/graphql/fragments.py,sha256=h3edd0Es38SNldhHiRU3hAJNYwrUIJMWxiP8av3BfrI,431
@@ -91,14 +88,15 @@ primitive/utils/cache.py,sha256=FHGmVWYLJFQOazpXXcEwI0YJEZbdkgG39nOLdOv6VNk,1575
 primitive/utils/chunk_size.py,sha256=PAuVuirUTA9oRXyjo1c6MWxo31WVBRkWMuWw-AS58Bw,2914
 primitive/utils/config.py,sha256=DlFM5Nglo22WPtbpZSVtH7NX-PTMaKYlcrUE7GPRG4c,1058
 primitive/utils/daemons.py,sha256=mSoSHitiGfS4KYAEK9sKsiv_YcACHKgY3qISnDpUUIE,1086
-primitive/utils/exceptions.py,sha256=DrYHTcCAJGC7cCUwOx_FmdlVLWRdpzvDvpLb82heppE,311
+primitive/utils/exceptions.py,sha256=KSAR2zqWEjQhqb6OAtaMUvYejKPETdBVRRRIKzy90A8,554
 primitive/utils/logging.py,sha256=vpwu-hByZC1BgJfUi6iSfAxzCobP_zg9-99EUf80KtY,1132
 primitive/utils/memory_size.py,sha256=4xfha21kW82nFvOTtDFx9Jk2ZQoEhkfXii-PGNTpIUk,3058
 primitive/utils/printer.py,sha256=f1XUpqi5dkTL3GWvYRUGlSwtj2IxU1q745T4Fxo7Tn4,370
+primitive/utils/psutil.py,sha256=xa7ef435UL37jyjmUPbEqCO2ayQMpCs0HCrxVEvLcuM,763
 primitive/utils/shell.py,sha256=Z4zxmOaSyGCrS0D6I436iQci-ewHLt4UxVg1CD9Serc,2171
 primitive/utils/text.py,sha256=XiESMnlhjQ534xE2hMNf08WehE1SKaYFRNih0MmnK0k,829
-primitive-0.2.39.dist-info/METADATA,sha256=lJU7WTvdfIpoE4cpyEAtBRkQRw0tzhXSW6rjBNNbJEU,3569
-primitive-0.2.39.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-primitive-0.2.39.dist-info/entry_points.txt,sha256=p1K8DMCWka5FqLlqP1sPek5Uovy9jq8u51gUsP-z334,48
-primitive-0.2.39.dist-info/licenses/LICENSE.txt,sha256=B8kmQMJ2sxYygjCLBk770uacaMci4mPSoJJ8WoDBY_c,1098
-primitive-0.2.39.dist-info/RECORD,,
+primitive-0.2.42.dist-info/METADATA,sha256=9VdknXdkOAgcn7ckBzk3VYumF1wkmO7iE1nExBibq8w,3535
+primitive-0.2.42.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+primitive-0.2.42.dist-info/entry_points.txt,sha256=p1K8DMCWka5FqLlqP1sPek5Uovy9jq8u51gUsP-z334,48
+primitive-0.2.42.dist-info/licenses/LICENSE.txt,sha256=B8kmQMJ2sxYygjCLBk770uacaMci4mPSoJJ8WoDBY_c,1098
+primitive-0.2.42.dist-info/RECORD,,

primitive/db/base.py

@@ -1,5 +0,0 @@
-from sqlalchemy.orm import DeclarativeBase
-
-
-class Base(DeclarativeBase):
-    pass

primitive/db/models.py

@@ -1,88 +0,0 @@
-from typing import Any, Callable, Dict, Generic, List, Optional, Type, TypeVar, Union
-
-from sqlalchemy import Column, Integer, String
-from sqlalchemy.orm import Mapped, Query, mapped_column
-
-from .base import Base
-from .sqlite import Session
-
-T = TypeVar("T", bound="Base")
-
-
-class Manager(Generic[T]):
-    def __init__(self, model_cls_lambda: Callable[[], Type[T]]) -> None:
-        self.model_cls_lambda = model_cls_lambda
-        self.filters: Dict[str, Any] = {}
-
-    def create(self, **kwargs) -> T:
-        with Session() as session:
-            model = self.model_cls_lambda()
-            obj = model(**kwargs)
-            session.add(obj)
-            session.commit()
-            session.refresh(obj)
-            return obj
-
-    def filter_by(self, **kwargs) -> "Manager[T]":
-        self.filters = kwargs
-        return self
-
-    def exists(self) -> bool:
-        with Session() as session:
-            model = self.model_cls_lambda()
-            query = session.query(model)
-            query.filter_by(**self.filters)
-            self.filters.clear()
-            return query.count() > 0
-
-    def all(self) -> List[T]:
-        with Session() as session:
-            model = self.model_cls_lambda()
-            query = session.query(model)
-            query.filter_by(**self.filters)
-            self.filters.clear()
-            return query.all()
-
-    def first(self) -> Union[T, None]:
-        with Session() as session:
-            model = self.model_cls_lambda()
-            query = session.query(model)
-            query.filter_by(**self.filters)
-            self.filters.clear()
-            return query.first()
-
-    def update(self, update: Dict[Any, Any]) -> Query[T]:
-        with Session() as session:
-            model = self.model_cls_lambda()
-            query = session.query(model).filter_by(**self.filters)
-
-            if query.count() > 0:
-                query.update(update)
-                session.commit()
-                return query
-            else:
-                raise ValueError(f"Update failed, {model.__name__} not found")
-
-    def delete(self) -> None:
-        with Session() as session:
-            model = self.model_cls_lambda()
-            query = session.query(model).filter_by(**self.filters)
-
-            if query.count() > 0:
-                query.delete()
-                session.commit()
-            else:
-                raise ValueError(f"Delete failed, {model.__name__} not found")
-
-
-class JobRun(Base):
-    __tablename__ = "JobRun"
-
-    id = Column(Integer, primary_key=True)
-    job_run_id: Mapped[str] = mapped_column(String, nullable=False)
-    pid: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
-
-    objects: Manager["JobRun"] = Manager(lambda: JobRun)
-
-    def __repr__(self):
-        return f"<JobRun(id={self.id} job_run_id={self.job_run_id}, pid={self.pid})>"

primitive/db/sqlite.py

@@ -1,70 +0,0 @@
-from pathlib import Path
-from time import sleep
-
-from loguru import logger
-from sqlalchemy import Engine, create_engine, inspect
-from sqlalchemy.orm import Session as SQLAlchemySession
-
-from primitive.db.base import Base
-from primitive.utils.cache import get_cache_dir
-
-
-def init() -> None:
-    db_path: Path = get_cache_dir() / "primitive.sqlite3"
-
-    # Drop DB existing database if it exists
-    # if db_path.exists():
-    #     logger.warning(f"Deleting existing SQLite database at {db_path}")
-    #     db_path.unlink()
-    if db_path.exists():
-        return
-
-    logger.info(f"Initializing SQLite database at {db_path}")
-    engine = create_engine(f"sqlite:///{db_path}", echo=False)
-    Base.metadata.create_all(engine)
-
-
-def engine() -> Engine:
-    db_path: Path = get_cache_dir() / "primitive.sqlite3"
-    return create_engine(f"sqlite:///{db_path}", echo=False)
-
-
-def wait_for_db() -> None:
-    # Wait for the database to be created
-    db_path: Path = get_cache_dir() / "primitive.sqlite3"
-
-    max_tries = 60
-    current_try = 1
-    while not db_path.exists() and current_try <= max_tries:
-        logger.debug(
-            f"Waiting for SQLite database to be created... [{current_try} / {max_tries}]"
-        )
-        sleep(1)
-        current_try += 1
-        continue
-    if current_try > max_tries:
-        message = f"SQLite database was not created after {max_tries} tries. Exiting..."
-        logger.error(message)
-        raise RuntimeError(message)
-
-    # check if table exists
-    max_tries = 60
-    current_try = 1
-    while not inspect(engine()).has_table("JobRun") and current_try <= max_tries:
-        logger.error("SQLite database is not ready.")
-        sleep(1)
-        current_try += 1
-        continue
-    if current_try > max_tries:
-        message = (
-            f"Database tables were not created after {max_tries} tries. Exiting..."
-        )
-        logger.error(message)
-        raise RuntimeError(message)
-    logger.debug("SQLite database is ready.")
-
-
-def Session() -> SQLAlchemySession:
-    from sqlalchemy.orm import sessionmaker
-
-    return sessionmaker(bind=engine())()