primitive 0.1.57__py3-none-any.whl → 0.1.59__py3-none-any.whl

primitive/__about__.py

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

primitive/agent/actions.py

@@ -1,9 +1,12 @@
-import sys
 import shutil
+import sys
 from time import sleep
-from primitive.utils.actions import BaseAction
+
 from loguru import logger
+
 from primitive.__about__ import __version__
+from primitive.utils.actions import BaseAction
+
 from ..utils.cache import get_sources_cache
 from .runner import AgentRunner
 from .uploader import Uploader
@@ -73,12 +76,12 @@ class Agent(BaseAction):
                     sleep(sleep_amount)
                     continue
 
-                job_runs_data = self.primitive.jobs.get_job_runs(
+                job_runs_result = self.primitive.jobs.get_job_runs(
                     status="pending", first=1, reservation_id=active_reservation_id
                 )
 
                 pending_job_runs = [
-                    edge["node"] for edge in job_runs_data["jobRuns"]["edges"]
+                    edge["node"] for edge in job_runs_result.data["jobRuns"]["edges"]
                 ]
 
                 if not pending_job_runs:
@@ -119,18 +122,23 @@ class Agent(BaseAction):
                         job_run["jobSettings"]["rootDirectory"]
                     )
 
-                    runner = AgentRunner(
-                        primitive=self.primitive,
-                        source_dir=source_dir,
-                        job_id=job_run["id"],
-                        job_slug=job_run["job"]["slug"],
-                    )
-
-                    # Execute job
-                    runner.execute()
-
-                    # Clean up
-                    shutil.rmtree(path=downloaded_git_repository_dir)
+                    try:
+                        # Initialize Runner
+                        runner = AgentRunner(
+                            primitive=self.primitive,
+                            source_dir=source_dir,
+                            job_id=job_run["id"],
+                            job_slug=job_run["job"]["slug"],
+                        )
+                    except Exception as e:
+                        # Log Error
+                        logger.error(f"Error initializing agent runner: {e}")
+                    else:
+                        # Execute job
+                        runner.execute()
+                    finally:
+                        # Clean up
+                        shutil.rmtree(path=downloaded_git_repository_dir)
 
                 sleep(5)
         except KeyboardInterrupt:

primitive/agent/runner.py

@@ -1,16 +1,17 @@
-import yaml
-import sys
-import typing
 import os
 import threading
-from time import sleep
-from typing import TypedDict, Iterable, List, Optional, Dict
+import typing
 from pathlib import Path, PurePath
+from time import sleep
+from typing import Dict, Iterable, List, Optional, TypedDict
+
+import yaml
 from loguru import logger
-from .process import Process
-from .provision import ProvisionPython
+
 from ..utils.cache import get_artifacts_cache
 from ..utils.files import find_files_for_extension
+from .process import Process
+from .provision import ProvisionPython
 
 try:
     from yaml import CLoader as Loader
@@ -72,7 +73,8 @@ class AgentRunner:
             logger.error(
                 f"Found two job descriptions with the same slug: {self.job_slug}"
             )
-            sys.exit(1)
+            self.conclude(conclusion="failure")
+            raise FileExistsError
 
         if yaml_file.exists():
             self.job = yaml.load(open(yaml_file, "r"), Loader=Loader)
@@ -82,7 +84,8 @@ class AgentRunner:
             logger.error(
                 f"No job description with matching slug '{self.job_slug}' found"
             )
-            sys.exit(1)
+            self.conclude(conclusion="failure")
+            raise FileNotFoundError
 
         logger.info(f"Found job description for {self.job_slug}")
 

primitive/auth/actions.py

@@ -1,23 +1,16 @@
 from gql import gql
 
-from ..utils.config import read_config_file, update_config_file
-from ..utils.auth import guard
-
 from primitive.utils.actions import BaseAction
 
+from ..utils.auth import guard
+from ..utils.config import read_config_file, update_config_file
+from .graphql.queries import whoami_query
+
 
 class Auth(BaseAction):
     @guard
     def whoami(self):
-        query = gql(
-            """
-            query whoami {
-                whoami {
-                    username
-                }
-            }
-        """
-        )
+        query = gql(whoami_query)
 
         result = self.primitive.session.execute(query, get_execution_result=True)
 

primitive/auth/graphql/queries.py

@@ -0,0 +1,13 @@
+whoami_query = """
+query whoami {
+    whoami {
+        username
+        defaultOrganization {
+            id
+            pk
+            name
+            slug
+        }
+    }
+}
+"""

primitive/cli.py

@@ -1,18 +1,22 @@
 import os
 import sys
+
 import click
+
 from .__about__ import __version__
-from .client import Primitive
+from .agent.commands import cli as agent_commands
 from .auth.commands import config_command, whoami_command
+from .client import Primitive
+from .daemons.commands import cli as daemons_commands
+from .exec.commands import cli as exec_commands
 from .files.commands import cli as file_commands
-from .hardware.commands import cli as hardware_commands
-from .lint.commands import cli as lint_commands
-from .agent.commands import cli as agent_commands
 from .git.commands import cli as git_commands
-from .daemons.commands import cli as daemons_commands
+from .hardware.commands import cli as hardware_commands
 from .jobs.commands import cli as jobs_commands
+from .lint.commands import cli as lint_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 .sim.commands import cli as sim_commands
 
 
@@ -69,6 +73,8 @@ cli.add_command(jobs_commands, "jobs")
 cli.add_command(organizations_commands, "organizations")
 cli.add_command(projects_commands, "projects")
 cli.add_command(sim_commands, "sim")
+cli.add_command(reservations_commands, "reservations")
+cli.add_command(exec_commands, "exec")
 
 if __name__ == "__main__":
     cli(obj={})

primitive/client.py

@@ -11,6 +11,8 @@ from .git.actions import Git
 from .daemons.actions import Daemons
 from .jobs.actions import Jobs
 from .organizations.actions import Organizations
+from .exec.actions import Exec
+from .reservations.actions import Reservations
 
 from loguru import logger
 
@@ -58,11 +60,13 @@ class Primitive:
         self.jobs: Jobs = Jobs(self)
         self.files: Files = Files(self)
         self.sim: Sim = Sim(self)
+        self.reservations: Reservations = Reservations(self)
         self.hardware: Hardware = Hardware(self)
         self.lint: Lint = Lint(self)
         self.agent: Agent = Agent(self)
         self.git: Git = Git(self)
         self.daemons: Daemons = Daemons(self)
+        self.exec: Exec = Exec(self)
 
     def get_host_config(self):
         self.full_config = read_config_file()

primitive/exec/actions.py

@@ -0,0 +1,50 @@
+import typing
+
+if typing.TYPE_CHECKING:
+    pass
+
+
+from primitive.utils.actions import BaseAction
+
+
+class Exec(BaseAction):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def execute_command(self, hardware_identifier: str, command: str) -> None:
+        hardware = self.primitive.hardware.get_hardware_from_slug_or_id(
+            hardware_identifier=hardware_identifier
+        )
+
+        # since we found hardware, we need to check that the user:
+        # - has a valid reservation on it
+        # - OR if the device is free we can reserve it
+
+        # if we create a reservation on behalf of the user, we need to release it after
+        created_reservation_on_behalf_of_user = False
+
+        if active_reservation := hardware["activeReservation"]:
+            active_reservation_id = active_reservation["id"]
+            reservation_result = self.primitive.reservations.get_reservation(
+                reservation_id=active_reservation_id
+            )
+            reservation = reservation_result.data["reservation"]
+        else:
+            reservation_result = self.primitive.reservations.create_reservation(
+                requested_hardware_ids=[hardware["id"]],
+                reason="Executing command from Primitive CLI",
+            )
+            reservation = reservation_result.data["reservationCreate"]
+            created_reservation_on_behalf_of_user = True
+
+        reservation = self.primitive.reservations.wait_for_reservation_status(
+            reservation_id=reservation["id"], desired_status="in_progress"
+        )
+
+        print(f"Executing command: {command} on {hardware['name']}")
+
+        if created_reservation_on_behalf_of_user:
+            print("Cleaning up reservation.")
+            self.primitive.reservations.release_reservation(
+                reservation_or_hardware_identifier=reservation["id"]
+            )

primitive/exec/commands.py

@@ -0,0 +1,22 @@
+import typing
+
+import click
+
+if typing.TYPE_CHECKING:
+    from ..client import Primitive
+
+
+@click.command("exec")
+@click.pass_context
+@click.argument(
+    "hardware_identifier",
+    type=str,
+    required=True,
+)
+@click.argument("command", nargs=-1, required=True)
+def cli(context, hardware_identifier: str, command: str) -> None:
+    """Exec"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    primitive.exec.execute_command(
+        hardware_identifier=hardware_identifier, command=command
+    )

primitive/files/actions.py

@@ -1,10 +1,13 @@
 from pathlib import Path
+
 from gql import gql
-from primitive.graphql.sdk import create_requests_session
-from ..utils.auth import guard
 
+from primitive.graphql.sdk import create_requests_session
 from primitive.utils.actions import BaseAction
 
+from ..utils.auth import guard
+from .graphql.mutations import create_trace_mutation
+
 
 class Files(BaseAction):
     @guard
@@ -17,19 +20,7 @@ class Files(BaseAction):
         is_vector: bool,
         size: int,
     ):
-        mutation = gql(
-            """
-            mutation createTrace($input: TraceCreateInput!) {
-                traceCreate(input: $input) {
-                    ... on Trace {
-                        id
-                        signalId
-                        signalName
-                    }
-                }
-            }
-        """
-        )
+        mutation = gql(create_trace_mutation)
         input = {
             "fileId": file_id,
             "signalId": signal_id,
@@ -39,7 +30,9 @@ class Files(BaseAction):
             "size": size,
         }
         variables = {"input": input}
-        result = self.primitive.session.execute(mutation, variable_values=variables)
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
         return result
 
     @guard

primitive/files/graphql/mutations.py

@@ -0,0 +1,11 @@
+create_trace_mutation = """
+mutation createTrace($input: TraceCreateInput!) {
+    traceCreate(input: $input) {
+        ... on Trace {
+            id
+            signalId
+            signalName
+        }
+    }
+}
+"""

primitive/git/actions.py

@@ -1,26 +1,26 @@
+import os
 from pathlib import Path
+
+from gql import gql
+from loguru import logger
+
 from primitive.utils.actions import BaseAction
+
 from ..utils.auth import guard
-from loguru import logger
-import os
+from .graphql.queries import github_app_token_query
 
 
 class Git(BaseAction):
     @guard
     def get_github_access_token(self) -> str:
-        query = """
-query githubAppToken{
-    githubAppToken {
-        token
-    }
-}
-"""
-
+        query = gql(github_app_token_query)
         filters = {}
         variables = {
             "filters": filters,
         }
-        result = self.primitive.session.execute(query, variable_values=variables)
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
         return result
 
     def download_git_repository_at_ref(

primitive/git/commands.py

@@ -1,8 +1,10 @@
-import click
-from pathlib import Path
+import os
 import typing
+from pathlib import Path
+
+import click
+
 from ..utils.printer import print_result
-import os
 
 if typing.TYPE_CHECKING:
     from ..client import Primitive
@@ -39,11 +41,10 @@ def download_ref_command(
     destination: Path = Path.cwd(),
 ):
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    result, message = primitive.git.download_git_repository_at_ref(
+    path = primitive.git.download_git_repository_at_ref(
         git_repo_full_name=git_repo_full_name,
         git_ref=git_ref,
         github_access_token=github_access_token,
         destination=destination,
     )
-    fg = "red" if not result else "green"
-    print_result(message, context=context, fg=fg)
+    print_result(message=path, context=context)

primitive/git/graphql/queries.py

@@ -0,0 +1,7 @@
+github_app_token_query = """
+query githubAppToken{
+    githubAppToken {
+        token
+    }
+}
+"""

primitive/graphql/relay.py

@@ -0,0 +1,32 @@
+import base64
+from typing import Tuple
+
+
+def from_base64(value: str) -> Tuple[str, str]:
+    """
+    FROM:
+    https://github.com/strawberry-graphql/strawberry/blob/main/strawberry/relay/utils.py#L16C1-L40C1
+
+    Parse the base64 encoded relay value.
+
+    Args:
+        value:
+            The value to be parsed
+
+    Returns:
+        A tuple of (TypeName, NodeID).
+
+    Raises:
+        ValueError:
+            If the value is not in the expected format
+
+    """
+    try:
+        res = base64.b64decode(value.encode()).decode().split(":", 1)
+    except Exception as e:
+        raise ValueError(str(e)) from e
+
+    if len(res) != 2:
+        raise ValueError(f"{res} expected to contain only 2 items")
+
+    return res[0], res[1]

primitive/graphql/utility_fragments.py

@@ -0,0 +1,19 @@
+operation_info_fragment = """
+fragment OperationInfoFragment on OperationInfo {
+      messages {
+        kind
+        message
+        field
+        code
+    }
+}
+"""
+
+page_info_fragment = """
+fragment PageInfoFragment on PageInfo {
+    hasNextPage
+    hasPreviousPage
+    startCursor
+    endCursor
+}
+"""