primitive 0.2.63__py3-none-any.whl → 0.2.64__py3-none-any.whl

primitive/__about__.py

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

primitive/agent/actions.py

@@ -53,7 +53,9 @@ class Agent(BaseAction):
                 # - then wait for the JobRun to be in_progress from the API
 
                 active_reservation_id = None
+                hardware = None
                 job_run_data: dict = {}
+                active_reservation_data: dict = {}
 
                 if RUNNING_IN_CONTAINER and job_run_id:
                     job_run_result = self.primitive.jobs.get_job_run(id=job_run_id)
@@ -81,16 +83,43 @@ class Agent(BaseAction):
                     if RUNNING_IN_CONTAINER:
                         logger.info("Running in container, exiting due to no JobRun.")
                         break
-                    logger.debug("No pending Job Run found, sleeping...")
+                    logger.debug("No pending Job Run found. [sleeping 5 seconds]")
                     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.debug(f"Job Run ID: {job_run_data.get('id', 'No Job ID Found')}")
+                logger.debug(
+                    f"Job Name: {job_run_data.get('job', {}).get('name', 'No Job Name Found')}"
+                )
 
                 job_run_status = job_run_data.get("status", None)
 
+                hardware_id = hardware.get("id", None) if hardware else None
+                execution_hardware_id = (
+                    job_run_data.get("executionHardware", {}).get("id", None)
+                    if job_run_data
+                    else None
+                )
+                target_hardware_id = None
+
+                if (
+                    hardware_id is not None
+                    and execution_hardware_id is not None
+                    and (hardware_id != execution_hardware_id)
+                ):
+                    logger.info(
+                        f"Job Run {job_run_id} is being executed by the controller. Agent may stop. [sleeping 5 seconds]"
+                    )
+                    sleep(5)
+                    continue
+
+                if active_reservation_data and active_reservation_id:
+                    for hardware in active_reservation_data.get("hardware", []):
+                        if hardware.get("id", None) != execution_hardware_id:
+                            target_hardware_id = hardware.get("id", None)
+                            break
+
                 while True:
                     if job_run_status == "pending":
                         # we are setting to request_in_progress here which puts a started_at time on the JobRun in the API's database
@@ -135,6 +164,7 @@ class Agent(BaseAction):
                         runner = Runner(
                             primitive=self.primitive,
                             job_run=job_run_data,
+                            target_hardware_id=target_hardware_id,
                         )
                         runner.setup()
                     except Exception as exception:

primitive/agent/runner.py

@@ -5,6 +5,7 @@ import typing
 from enum import Enum
 from pathlib import Path, PurePath
 from typing import Dict, List, TypedDict
+from datetime import datetime, timezone
 
 from loguru import logger
 
@@ -42,12 +43,14 @@ class Runner:
         self,
         primitive: "primitive.client.Primitive",
         job_run: Dict,
+        target_hardware_id: str | None = None,
         max_log_size: int = 10 * 1024 * 1024,  # 10MB
     ) -> None:
         self.primitive = primitive
         self.job = job_run["job"]
         self.job_run = job_run
         self.job_settings = job_run["jobSettings"]
+        self.target_hardware_id = target_hardware_id
         self.config = job_run["jobSettings"]["config"]
         self.initial_env = {}
         self.modified_env = {}
@@ -64,6 +67,14 @@ class Runner:
                 backtrace=True,
             )
 
+        if self.target_hardware_id is not None:
+            target_hardware_secret = self.primitive.hardware.get_hardware_secret(
+                hardware_id=self.target_hardware_id
+            )
+            self.target_hardware_secret = {
+                k: v for k, v in target_hardware_secret.items() if v is not None
+            }
+
     @log_context(label="setup")
     def setup(self) -> None:
         # Attempt to download the job source code
@@ -217,11 +228,96 @@ class Runner:
     async def run_task(self, task: Task) -> bool:
         logger.info(f"Running step '{task['label']}'")
         commands = task["cmd"].strip().split("\n")
+
         for i, cmd in enumerate(commands):
-            # Adding an additional echo and utilizing stdbuf to force line buffering
-            # This ensures that the environment variables and starting delimiter are
-            # always in a new chunk, vastly simplifying our parsing logic
+            if cmd == "oobpowercycle":
+                from primitive.network.redfish import RedfishClient
+
+                bmc_host = self.target_hardware_secret.get("bmcHostname", None)
+                bmc_username = self.target_hardware_secret.get("bmcUsername", None)
+                bmc_password = self.target_hardware_secret.get("bmcPassword", "")
+
+                if bmc_host is None:
+                    logger.error(
+                        "No BMC host found in target hardware secret for out-of-band power cycle"
+                    )
+                    return True
+                if bmc_username is None:
+                    logger.error(
+                        "No BMC username found in target hardware secret for out-of-band power cycle"
+                    )
+                    return True
+
+                redfish = RedfishClient(
+                    host=bmc_host, username=bmc_username, password=bmc_password
+                )
+                redfish.compute_system_reset(system_id="1", reset_type="ForceRestart")
+                if self.target_hardware_id:
+                    self.primitive.hardware.update_hardware(
+                        hardware_id=self.target_hardware_id,
+                        is_online=False,
+                        is_rebooting=True,
+                        start_rebooting_at=datetime.now(timezone.utc),
+                    )
+                continue
+
+            if cmd == "pxeboot":
+                from primitive.network.redfish import RedfishClient
+
+                bmc_host = self.target_hardware_secret.get("bmcHostname", None)
+                bmc_username = self.target_hardware_secret.get("bmcUsername", None)
+                bmc_password = self.target_hardware_secret.get("bmcPassword", "")
+
+                if bmc_host is None:
+                    logger.error(
+                        "No BMC host found in target hardware secret for out-of-band power cycle"
+                    )
+                    return True
+                if bmc_username is None:
+                    logger.error(
+                        "No BMC username found in target hardware secret for out-of-band power cycle"
+                    )
+                    return True
+
+                redfish = RedfishClient(
+                    host=bmc_host, username=bmc_username, password=bmc_password
+                )
+                redfish.update_boot_options(
+                    system_id="1",
+                    boot_source_override_target="Pxe",
+                    boot_source_override_enabled="Once",
+                    boot_source_override_mode="UEFI",
+                )
+                redfish.compute_system_reset(system_id="1", reset_type="ForceRestart")
+                if self.target_hardware_id:
+                    self.primitive.hardware.update_hardware(
+                        hardware_id=self.target_hardware_id,
+                        is_online=False,
+                        is_rebooting=True,
+                        start_rebooting_at=datetime.now(timezone.utc),
+                    )
+                continue
+
             args = ["/bin/bash", "-c", cmd]
+            if self.target_hardware_secret:
+                username = self.target_hardware_secret.get("username")
+                password = self.target_hardware_secret.get("password")
+                hostname = self.target_hardware_secret.get("hostname")
+                args = [
+                    "sshpass",
+                    "-p",
+                    password,
+                    "ssh",
+                    "-o",
+                    "StrictHostKeyChecking=no",
+                    "-o",
+                    "UserKnownHostsFile=/dev/null",
+                    "-o",
+                    "IdentitiesOnly=yes",
+                    f"{username}@{hostname}",
+                    "--",
+                    f'/bin/bash -c "cd {task.get("workdir", "~/")} && {cmd}"',
+                ]
 
             logger.info(
                 f"Executing command {i + 1}/{len(commands)}: {cmd} at {self.source_dir / task.get('workdir', '')}"

primitive/cli.py

@@ -17,6 +17,7 @@ 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
+from .network.commands import cli as network_commands
 
 
 @click.group()
@@ -73,6 +74,7 @@ cli.add_command(projects_commands, "projects")
 cli.add_command(reservations_commands, "reservations")
 cli.add_command(exec_commands, "exec")
 cli.add_command(monitor_commands, "monitor")
+cli.add_command(network_commands, "network")
 
 if __name__ == "__main__":
     cli(obj={})

primitive/client.py

@@ -16,6 +16,7 @@ from .git.actions import Git
 from .hardware.actions import Hardware
 from .jobs.actions import Jobs
 from .monitor.actions import Monitor
+from .network.actions import Network
 from .organizations.actions import Organizations
 from .projects.actions import Projects
 from .provisioning.actions import Provisioning
@@ -97,6 +98,7 @@ class Primitive:
         self.exec: Exec = Exec(self)
         self.provisioning: Provisioning = Provisioning(self)
         self.monitor: Monitor = Monitor(self)
+        self.network: Network = Network(self)
 
     def get_host_config(self):
         self.full_config = read_config_file()

primitive/hardware/actions.py

@@ -1,4 +1,5 @@
 import csv
+from datetime import datetime
 import io
 import json
 import platform
@@ -10,6 +11,7 @@ from typing import Dict, List, Optional
 
 import click
 import psutil
+from socket import AddressFamily
 from aiohttp import client_exceptions
 from gql import gql
 from loguru import logger
@@ -22,20 +24,22 @@ from ..utils.auth import guard
 from ..utils.config import update_config_file
 from ..utils.exceptions import P_CLI_100
 from .graphql.mutations import (
+    hardware_certificate_create_mutation,
     hardware_checkin_mutation,
     hardware_update_mutation,
+    hardware_update_system_info_mutation,
     register_child_hardware_mutation,
     register_hardware_mutation,
     unregister_hardware_mutation,
-    hardware_certificate_create_mutation,
 )
 from .graphql.queries import (
     hardware_details,
     hardware_list,
     nested_children_hardware_list,
+    hardware_secret,
+    hardware_with_parent_list,
 )
 
-
 if typing.TYPE_CHECKING:
     pass
 
@@ -199,6 +203,34 @@ class Hardware(BaseAction):
 
         return gpu_config
 
+    def _get_network_interfaces(self) -> Dict[str, str]:
+        os_family = platform.system()
+        network_interfaces = {}
+        for interface, addresses in psutil.net_if_addrs().items():
+            if (
+                interface in ["lo", "lo0", "awdl0", "llw0"]
+                or interface.startswith("utun")
+                or interface.startswith("bridge")
+            ):
+                continue
+
+            network_interfaces[interface] = {}
+            for address in addresses:
+                # get linux's mac address
+                if os_family == "Linux" and address.family == AddressFamily.AF_PACKET:
+                    mac_address = address.address
+                    if mac_address and mac_address != "00:00:00:00:00:00":
+                        network_interfaces[interface]["mac_address"] = mac_address
+                if os_family == "Darwin" and address.family == AddressFamily.AF_LINK:
+                    mac_address = address.address
+                    if mac_address and mac_address != "00:00:00:00:00:00":
+                        network_interfaces[interface]["mac_address"] = mac_address
+                elif address.family == AddressFamily.AF_INET:
+                    ip_address = address.address
+                    if ip_address and ip_address != "127.0.0.1":
+                        network_interfaces[interface]["ip_address"] = ip_address
+        return network_interfaces
+
     def _get_windows_computer_service_product_values(self) -> Dict[str, str]:
         windows_computer_service_product_csv_command = (
             "cmd.exe /C wmic csproduct get Name, Vendor, Version, UUID /format:csv"
@@ -289,6 +321,8 @@ class Hardware(BaseAction):
         system_info["architecture"] = platform.architecture()[0]
         system_info["cpu_cores"] = str(platform.os.cpu_count())  # type: ignore exits
         system_info["gpu_config"] = self._get_gpu_config()
+        system_info["network_interfaces"] = self._get_network_interfaces()
+
         return system_info
 
     @guard
@@ -410,7 +444,7 @@ class Hardware(BaseAction):
             "systemInfo": system_info,
         }
 
-        mutation = gql(hardware_update_mutation)
+        mutation = gql(hardware_update_system_info_mutation)
 
         input = new_state
         variables = {"input": input}
@@ -435,16 +469,26 @@ class Hardware(BaseAction):
         is_available: bool = False,
         is_online: bool = True,
         stopping_agent: Optional[bool] = False,
+        http: bool = False,
     ):
-        message = {
-            "is_healthy": is_healthy,
-            "is_quarantined": is_quarantined,
-            "is_available": is_available,
-            "is_online": is_online,
-        }
-        self.primitive.messaging.send_message(
-            message_type=MESSAGE_TYPES.CHECK_IN, message=message
-        )
+        if not http and self.primitive.messaging.ready:
+            message = {
+                "is_healthy": is_healthy,
+                "is_quarantined": is_quarantined,
+                "is_available": is_available,
+                "is_online": is_online,
+            }
+            self.primitive.messaging.send_message(
+                message_type=MESSAGE_TYPES.CHECK_IN, message=message
+            )
+        else:
+            self.check_in_http(
+                is_healthy=is_healthy,
+                is_quarantined=is_quarantined,
+                is_available=is_available,
+                is_online=is_online,
+                stopping_agent=stopping_agent,
+            )
 
     @guard
     def check_in_http(
@@ -544,8 +588,13 @@ class Hardware(BaseAction):
         id: Optional[str] = None,
         slug: Optional[str] = None,
         nested_children: Optional[bool] = False,
+        parent: Optional[bool] = False,
     ):
-        query = gql(nested_children_hardware_list if nested_children else hardware_list)
+        query = gql(hardware_list)
+        if parent:
+            query = gql(hardware_with_parent_list)
+        if nested_children:
+            query = gql(nested_children_hardware_list)
 
         filters = {
             "isRegistered": {"exact": True},
@@ -556,8 +605,8 @@ class Hardware(BaseAction):
             filters["slug"] = {"exact": slug}
         if id is not None:
             filters["id"] = {"exact": id}
-        if nested_children is True:
-            filters["hasParent"] = {"exact": False}
+        # if nested_children is True:
+        #     filters["hasParent"] = {"exact": False}
 
         variables = {
             "filters": filters,
@@ -565,6 +614,43 @@ class Hardware(BaseAction):
         result = self.primitive.session.execute(
             query, variable_values=variables, get_execution_result=True
         )
+
+        return result
+
+    @guard
+    async def aget_hardware_list(
+        self,
+        fingerprint: Optional[str] = None,
+        id: Optional[str] = None,
+        slug: Optional[str] = None,
+        nested_children: Optional[bool] = False,
+        parent: Optional[bool] = False,
+    ):
+        query = gql(hardware_list)
+        if parent:
+            query = gql(hardware_with_parent_list)
+        if nested_children:
+            query = gql(nested_children_hardware_list)
+
+        filters = {
+            "isRegistered": {"exact": True},
+        }
+        if fingerprint is not None:
+            filters["fingerprint"] = {"exact": fingerprint}
+        if slug is not None:
+            filters["slug"] = {"exact": slug}
+        if id is not None:
+            filters["id"] = {"exact": id}
+        # if nested_children is True:
+        #     filters["hasParent"] = {"exact": False}
+
+        variables = {
+            "filters": filters,
+        }
+        result = await self.primitive.session.execute_async(
+            query, variable_values=variables, get_execution_result=True
+        )
+
         return result
 
     @guard
@@ -601,11 +687,56 @@ class Hardware(BaseAction):
             fingerprint=self.primitive.host_config.get("fingerprint"),
             nested_children=True,
         )
+        if len(hardware_list_result.data.get("hardwareList").get("edges", [])) == 0:
+            raise Exception(
+                "No hardware found with fingerprint: "
+                f"{self.primitive.host_config.get('fingerprint')}. "
+                "Please register: primitive hardware register"
+            )
+
         hardware = (
             hardware_list_result.data.get("hardwareList").get("edges")[0].get("node")
         )
         return hardware
 
+    def get_parent_hardware_details(self):
+        hardware_list_result = self.get_hardware_list(
+            fingerprint=self.primitive.host_config.get("fingerprint"), parent=True
+        )
+        if len(hardware_list_result.data.get("hardwareList").get("edges", [])) == 0:
+            logger.warning(
+                "No hardware found with fingerprint: "
+                f"{self.primitive.host_config.get('fingerprint')}. "
+                "Please register: primitive hardware register"
+            )
+
+        hardware = (
+            hardware_list_result.data.get("hardwareList").get("edges")[0].get("node")
+        )
+        parent = hardware.get("parent", None)
+        if not parent:
+            logger.warning("No parent network device found.")
+        return parent
+
+    async def aget_parent_hardware_details(self):
+        hardware_list_result = await self.aget_hardware_list(
+            fingerprint=self.primitive.host_config.get("fingerprint"), parent=True
+        )
+        if len(hardware_list_result.data.get("hardwareList").get("edges", [])) == 0:
+            logger.warning(
+                "No hardware found with fingerprint: "
+                f"{self.primitive.host_config.get('fingerprint')}. "
+                "Please register: primitive hardware register"
+            )
+
+        hardware = (
+            hardware_list_result.data.get("hardwareList").get("edges")[0].get("node")
+        )
+        parent = hardware.get("parent", None)
+        if not parent:
+            logger.warning("No parent network device found.")
+        return parent
+
     def get_hardware_from_slug_or_id(self, hardware_identifier: str):
         is_id = False
         is_slug = False
@@ -634,6 +765,56 @@ class Hardware(BaseAction):
 
         return hardware
 
+    @guard
+    def get_hardware_secret(self, hardware_id: str):
+        query = gql(hardware_secret)
+        variables = {"hardwareId": hardware_id}
+        result = self.primitive.session.execute(
+            query, variable_values=variables, get_execution_result=True
+        )
+
+        secret = result.data.get("hardwareSecret", {})
+        return secret
+
+    @guard
+    async def aget_hardware_secret(self, hardware_id: str):
+        query = gql(hardware_secret)
+        variables = {"hardwareId": hardware_id}
+        result = await self.primitive.session.execute_async(
+            query, variable_values=variables, get_execution_result=True
+        )
+
+        secret = result.data.get("hardwareSecret", {})
+        return secret
+
+    @guard
+    def get_and_set_switch_info(self):
+        parent = self.get_parent_hardware_details()
+        if not parent:
+            return
+
+        parent_secret = self.get_hardware_secret(hardware_id=parent.get("id"))
+        self.primitive.network.switch_connection_info = {
+            "vendor": parent.get("manufacturer", {}).get("slug"),
+            "hostname": parent_secret.get("hostname"),
+            "username": parent_secret.get("username"),
+            "password": parent_secret.get("password"),
+        }
+
+    @guard
+    async def aget_and_set_switch_info(self):
+        parent = await self.aget_parent_hardware_details()
+        if not parent:
+            return
+
+        parent_secret = await self.aget_hardware_secret(hardware_id=parent.get("id"))
+        self.primitive.network.switch_connection_info = {
+            "vendor": parent.get("manufacturer", {}).get("slug"),
+            "hostname": parent_secret.get("hostname"),
+            "username": parent_secret.get("username"),
+            "password": parent_secret.get("password"),
+        }
+
     @guard
     def register_child(self, child: AndroidDevice):
         system_info = child.system_info
@@ -764,3 +945,39 @@ class Hardware(BaseAction):
             "disk_free": disk_usage.free,
         }
         return metrics
+
+    @guard
+    def update_hardware(
+        self,
+        hardware_id: str,
+        is_online: Optional[bool] = None,
+        is_rebooting: Optional[bool] = None,
+        start_rebooting_at: Optional[datetime] = None,
+    ):
+        new_state: dict = {
+            "id": hardware_id,
+        }
+        if is_online is not None:
+            new_state["isOnline"] = is_online
+        if is_rebooting is not None:
+            new_state["isRebooting"] = is_rebooting
+        if start_rebooting_at is not None:
+            new_state["startRebootingAt"] = start_rebooting_at
+
+        mutation = gql(hardware_update_mutation)
+
+        input = new_state
+        variables = {"input": input}
+        try:
+            result = self.primitive.session.execute(
+                mutation, variable_values=variables, get_execution_result=True
+            )
+        except client_exceptions.ClientConnectorError as exception:
+            message = "Failed to update hardware! "
+            logger.exception(message)
+            raise exception
+
+        message = "Updated hardware successfully! "
+        logger.info(message)
+
+        return result

primitive/hardware/commands.py

@@ -101,11 +101,17 @@ def unregister_command(context):
 
 
 @cli.command("checkin")
+@click.option(
+    "--http",
+    is_flag=True,
+    default=False,
+    help="Use HTTP instead of amqp for check-in",
+)
 @click.pass_context
-def checkin_command(context):
+def checkin_command(context, http: bool = False):
     """Checkin Hardware with Primitive"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    primitive.hardware.check_in()
+    primitive.hardware.check_in(http=http)
 
 
 @cli.command("list")

primitive/hardware/graphql/fragments.py

@@ -10,9 +10,16 @@ fragment HardwareFragment on Hardware {
   isOnline
   isQuarantined
   isHealthy
+  isController
   systemInfo
   hostname
   sshUsername
+  manufacturer {
+    id
+    pk
+    name
+    slug
+  }
   organization {
     id
     pk
@@ -25,6 +32,9 @@ fragment HardwareFragment on Hardware {
     status
     reservationNumber
     reason
+    hardware {
+      id
+    }
     createdBy {
       username
     }

primitive/hardware/graphql/mutations.py

@@ -1,4 +1,5 @@
 from primitive.graphql.utility_fragments import operation_info_fragment
+from primitive.hardware.graphql.fragments import hardware_fragment
 
 hardware_certificate_create_mutation = (
     operation_info_fragment
@@ -60,14 +61,28 @@ mutation unregisterHardware($input: UnregisterHardwareInput!) {
 """
 )
 
+hardware_update_system_info_mutation = (
+    operation_info_fragment
+    + """
+mutation hardwareUpdate($input: HardwareUpdateInput!) {
+    hardwareUpdate(input: $input) {
+        ... on Hardware {
+            systemInfo
+        }
+        ...OperationInfoFragment
+    }
+}
+"""
+)
 
 hardware_update_mutation = (
     operation_info_fragment
+    + hardware_fragment
     + """
 mutation hardwareUpdate($input: HardwareUpdateInput!) {
     hardwareUpdate(input: $input) {
         ... on Hardware {
-            systemInfo
+            ...HardwareFragment
         }
         ...OperationInfoFragment
     }