primitive 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

primitive/hardware/actions.py

@@ -0,0 +1,447 @@
+import csv
+import io
+import json
+import platform
+from shutil import which
+import subprocess
+from typing import Dict, List
+import click
+from loguru import logger
+from primitive.utils.memory_size import MemorySize
+from gql import gql
+from aiohttp import client_exceptions
+from ..utils.config import update_config_file
+
+import typing
+
+if typing.TYPE_CHECKING:
+    pass
+
+
+from primitive.utils.actions import BaseAction
+
+
+class Hardware(BaseAction):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.previous_state = {
+            "isHealthy": False,
+            "isQuarantined": False,
+            "isAvailable": False,
+            "isOnline": False,
+        }
+
+    def _get_darwin_system_profiler_values(self) -> Dict[str, str]:
+        system_profiler_hardware_data_type = subprocess.check_output(
+            ["system_profiler", "SPHardwareDataType", "-json"]
+        )
+        system_profiler_hardware_data = json.loads(system_profiler_hardware_data_type)
+        data_type = system_profiler_hardware_data.get("SPHardwareDataType")[0]
+        return {
+            "apple_model_name": data_type.get("machine_model"),
+            "apple_model_identifier": data_type.get("machine_name"),
+            "apple_model_number": data_type.get("model_number"),
+            "physical_memory": data_type.get("physical_memory"),
+            "apple_serial_number": data_type.get("serial_number"),
+        }
+
+    def _get_supported_metal_device(self) -> int | None:
+        """
+        Checks if metal hardware is supported. If so, the index
+        of the supported metal device is returned
+        """
+        supported_metal_device = None
+        is_system_profiler_available = bool(which("system_profiler"))
+        if is_system_profiler_available:
+            system_profiler_display_data_type_command = (
+                "system_profiler SPDisplaysDataType -json"
+            )
+            try:
+                system_profiler_display_data_type_output = subprocess.check_output(
+                    system_profiler_display_data_type_command.split(" ")
+                )
+            except subprocess.CalledProcessError as exception:
+                message = f"Error running system_profiler: {exception}"
+                logger.error(message)
+                return supported_metal_device
+
+            try:
+                system_profiler_display_data_type_json = json.loads(
+                    system_profiler_display_data_type_output
+                )
+            except json.JSONDecodeError as exception:
+                message = f"Error decoding JSON: {exception}"
+                logger.error(message)
+                return supported_metal_device
+
+            # Checks if any attached displays have metal support
+            # Note, other devices here could be AMD GPUs or unconfigured Nvidia GPUs
+            for index, display in enumerate(
+                system_profiler_display_data_type_json["SPDisplaysDataType"]
+            ):
+                if "spdisplays_mtlgpufamilysupport" in display:
+                    supported_metal_device = index
+                    return supported_metal_device
+
+        return supported_metal_device
+
+    def _get_gpu_config(self) -> List:
+        """
+        For Nvidia based systems, nvidia-smi will be used to profile the gpu/s.
+        For Metal based systems, we will gather information from SPDisplaysDataType.
+        """
+        gpu_config = []
+
+        # Check nvidia gpu availability
+        is_nvidia_smi_available = bool(which("nvidia-smi"))
+        if is_nvidia_smi_available:
+            nvidia_smi_query_gpu_csv_command = "nvidia-smi --query-gpu=gpu_name,driver_version,memory.total --format=csv"  # noqa
+            try:
+                nvidia_smi_query_gpu_csv_output = subprocess.check_output(
+                    nvidia_smi_query_gpu_csv_command.split(" "),
+                )
+            except subprocess.CalledProcessError as exception:
+                message = f"Command {nvidia_smi_query_gpu_csv_command} failed with exception: {exception}"  # noqa
+                logger.error(message)
+                raise exception
+
+            try:
+                nvidia_smi_query_gpu_csv_decoded = (
+                    nvidia_smi_query_gpu_csv_output.decode("utf-8")
+                    .replace("\r", "")
+                    .replace(", ", ",")
+                    .lstrip("\n")
+                )
+            except UnicodeDecodeError as exception:
+                message = f"Error decoding: {exception}"
+                logger.error(message)
+                raise exception
+
+            nvidia_smi_query_gpu_csv_dict_reader = csv.DictReader(
+                io.StringIO(nvidia_smi_query_gpu_csv_decoded)
+            )
+
+            for gpu_info in nvidia_smi_query_gpu_csv_dict_reader:
+                # Refactor key into B
+                memory_total_in_mebibytes = gpu_info.pop("memory.total [MiB]")
+                memory_size = MemorySize(memory_total_in_mebibytes)
+                gpu_info["memory_total"] = memory_size.to_bytes()
+
+                gpu_config.append(gpu_info)
+
+        if platform.system() == "Darwin":
+            # Check Metal gpu availability
+            supported_metal_device = self._get_supported_metal_device()
+            if supported_metal_device is not None:
+                # Since Apple's SoC contains Metal,
+                # we query the system itself for total memory
+                system_profiler_hardware_data_type_command = (
+                    "system_profiler SPHardwareDataType -json"
+                )
+
+                try:
+                    system_profiler_hardware_data_type_output = subprocess.check_output(
+                        system_profiler_hardware_data_type_command.split(" ")
+                    )
+                except subprocess.CalledProcessError as exception:
+                    message = f"Error running {system_profiler_hardware_data_type_command}: {exception}"  # noqa
+                    logger.error(message)
+                    raise exception
+
+                try:
+                    system_profiler_hardware_data_type_json = json.loads(
+                        system_profiler_hardware_data_type_output
+                    )
+                except json.JSONDecodeError as exception:
+                    message = f"Error decoding JSON: {exception}"  # noqa
+                    logger.error(message)
+                    raise exception
+
+                metal_device_json = system_profiler_hardware_data_type_json[
+                    "SPHardwareDataType"
+                ][supported_metal_device]
+
+                gpu_info = {}
+                gpu_info["name"] = metal_device_json.get("chip_type")
+
+                # Refactor key into B
+                physical_memory = metal_device_json.get("physical_memory")
+                memory_size = MemorySize(physical_memory)
+                gpu_info["memory_total"] = memory_size.to_bytes()
+
+                gpu_config.append(gpu_info)
+
+        # Raise an error if there is no valid gpu config
+        if not gpu_config:
+            message = "No valid gpu configuration"
+            logger.error(message)
+            raise NotImplementedError(message)
+
+        return gpu_config
+
+    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"
+        )
+        windows_computer_service_product_csv_output = subprocess.check_output(
+            windows_computer_service_product_csv_command.split(" "),
+            stderr=subprocess.DEVNULL,
+        )
+        windows_computer_service_product_csv_decoded = (
+            windows_computer_service_product_csv_output.decode("utf-8")
+            .replace("\r", "")
+            .lstrip("\n")
+        )
+        windows_computer_service_product_dict = csv.DictReader(
+            io.StringIO(windows_computer_service_product_csv_decoded)
+        )
+        csp_info = list(windows_computer_service_product_dict)[0]
+        return {
+            "windows_model_name": csp_info.get("Name", ""),
+            "windows_model_vendor": csp_info.get("Vendor", ""),
+            "windows_model_version": csp_info.get("Version", ""),
+            "windows_model_uuid": csp_info.get("UUID", ""),
+        }
+
+    def _get_windows_cpu_values(self) -> Dict[str, str]:
+        windows_cpu_csv_command = (
+            "cmd.exe /C wmic cpu get Name, MaxClockSpeed /format:csv"  # noqa
+        )
+        windows_cpu_csv_output = subprocess.check_output(
+            windows_cpu_csv_command.split(" "),
+            stderr=subprocess.DEVNULL,
+        )
+        windows_cpu_csv_decoded = (
+            windows_cpu_csv_output.decode("utf-8").replace("\r", "").lstrip("\n")
+        )
+        windows_cpu_dict = csv.DictReader(io.StringIO(windows_cpu_csv_decoded))
+        cpu_info = list(windows_cpu_dict)[0]
+        return {
+            "cpu_brand": cpu_info.get("Name", "").strip(),
+            "cpu_max_clock_speed": cpu_info.get("MaxClockSpeed", ""),
+        }
+
+    def _get_ubuntu_values(self) -> Dict[str, str]:
+        get_machine_id_command = "cat /etc/machine-id"
+        machine_id = subprocess.check_output(
+            get_machine_id_command.split(" "),
+            stderr=subprocess.DEVNULL,
+        ).decode("utf-8")
+        if machine_id:
+            return {"linux_machine_id": machine_id}
+        return {}
+
+    def get_system_info(self):
+        os_family = platform.system()
+        system_info = {}
+        if os_family == "Darwin":
+            system_info = {**system_info, **self._get_darwin_system_profiler_values()}
+            system_info["cpu_brand"] = (
+                subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
+                .strip()
+                .decode("utf-8")
+            )
+            system_info["apple_mac_os_version"] = platform.mac_ver()[0]
+        elif os_family == "Linux":
+            # Support for Linux-based VMs in Windows
+            if "WSL2" in platform.platform():
+                system_info = {
+                    **system_info,
+                    **self._get_windows_computer_service_product_values(),
+                    **self._get_windows_cpu_values(),
+                }
+            else:
+                system_info = {**system_info, **self._get_ubuntu_values()}
+        elif os_family == "Windows":
+            system_info = {
+                **system_info,
+                **self._get_windows_computer_service_product_values(),
+                **self._get_windows_cpu_values(),
+            }
+
+        system_info["name"] = platform.node()
+        system_info["os_family"] = os_family
+        system_info["os_release"] = platform.release()
+        system_info["os_version"] = platform.version()
+        system_info["platform"] = platform.platform()
+        system_info["processor"] = platform.processor()
+        system_info["machine"] = platform.machine()
+        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()
+        return system_info
+
+    def register(self):
+        system_info = self.get_system_info()
+        mutation = gql(
+            """
+            mutation registerHardware($input: RegisterHardwareInput!) {
+                registerHardware(input: $input) {
+                    ... on Hardware {
+                        fingerprint
+                    }
+                    ... on OperationInfo {
+                          messages {
+                            kind
+                            message
+                            field
+                            code
+                        }
+                    }
+                }
+            }
+        """
+        )
+        input = {"systemInfo": system_info}
+        variables = {"input": input}
+        result = self.primitive.session.execute(mutation, variable_values=variables)
+        if messages := result.get("registerHardware").get("messages"):
+            for message in messages:
+                logger.enable("primitive")
+                if message.get("kind") == "ERROR":
+                    logger.error(message.get("message"))
+                else:
+                    logger.debug(message.get("message"))
+            return False
+
+        fingerprint = result.get("registerHardware").get("fingerprint")
+
+        self.primitive.host_config["fingerprint"] = fingerprint
+        self.primitive.full_config[self.primitive.host] = self.primitive.host_config
+        update_config_file(new_config=self.primitive.full_config)
+
+        # then check in that the hardware, validate that it is saved correctly
+        # and headers are set correctly
+        self.primitive.get_host_config()
+        self.check_in_http(is_healthy=True)
+        return True
+
+    def update_hardware_system_info(self):
+        """
+        Updates hardware system information and returns the GraphQL response.
+
+        Returns:
+            dict: GraphQL response
+        Raises:
+            Exception: If no fingerprint is found or an error occurs
+        """
+
+        fingerprint = self.primitive.host_config.get("fingerprint", None)
+        if not fingerprint:
+            message = (
+                "No fingerprint found. Please register: primitive hardware register"
+            )
+            raise Exception(message)
+
+        system_info = self.get_system_info()
+        new_state = {
+            "systemInfo": system_info,
+        }
+
+        mutation = gql(
+            """
+            mutation hardwareUpdate($input: HardwareUpdateInput!) {
+                hardwareUpdate(input: $input) {
+                    ... on Hardware {
+                        systemInfo
+                    }
+                    ... on OperationInfo {
+                          messages {
+                            kind
+                            message
+                            field
+                            code
+                        }
+                    }
+                }
+            }
+            """
+        )
+
+        input = new_state
+        variables = {"input": input}
+        try:
+            result = self.primitive.session.execute(mutation, variable_values=variables)
+        except client_exceptions.ClientConnectorError as exception:
+            message = " [*] Failed to update hardware system info! "
+            logger.error(message)
+            raise exception
+
+        message = " [*] Updated hardware system info successfully! "
+        logger.info(message)
+
+        return result
+
+    def check_in_http(
+        self,
+        is_healthy: bool = True,
+        is_quarantined: bool = False,
+        is_available: bool = False,
+        is_online: bool = True,
+    ):
+        fingerprint = self.primitive.host_config.get("fingerprint", None)
+        if not fingerprint:
+            message = (
+                "No fingerprint found. Please register: primitive hardware register"
+            )
+            raise Exception(message)
+
+        new_state = {
+            "isHealthy": is_healthy,
+            "isQuarantined": is_quarantined,
+            "isAvailable": is_available,
+            "isOnline": is_online,
+        }
+
+        mutation = gql(
+            """
+            mutation checkIn($input: CheckInInput!) {
+                checkIn(input: $input) {
+                    ... on Hardware {
+                        createdAt
+                        updatedAt
+                        lastCheckIn
+                    }
+                    ... on OperationInfo {
+                          messages {
+                            kind
+                            message
+                            field
+                            code
+                        }
+                    }
+                }
+            }
+        """  # noqa
+        )
+        input = new_state
+        variables = {"input": input}
+        try:
+            result = self.primitive.session.execute(mutation, variable_values=variables)
+            previous_state = self.previous_state
+            self.previous_state = new_state.copy()
+
+            message = " [*] Checked in successfully: "
+            for key, value in new_state.items():
+                if value != previous_state.get(key, None):
+                    if value is True:
+                        message = (
+                            message
+                            + click.style(f"{key}: ")
+                            + click.style("💤")
+                            + click.style(" ==> ✅ ", fg="green")
+                        )
+                    else:
+                        message = (
+                            message
+                            + click.style(f"{key}: ")
+                            + click.style("✅")
+                            + click.style(" ==> 💤 ", fg="yellow")
+                        )
+            logger.info(message)
+
+            return result
+        except client_exceptions.ClientConnectorError as exception:
+            message = " [*] Failed to check in! "
+            logger.error(message)
+            raise exception

primitive/hardware/commands.py

@@ -0,0 +1,53 @@
+import click
+
+from ..utils.printer import print_result
+
+import typing
+
+if typing.TYPE_CHECKING:
+    from ..client import Primitive
+
+
+@click.group()
+@click.pass_context
+def cli(context):
+    """Hardware"""
+    pass
+
+
+@cli.command("systeminfo")
+@click.pass_context
+def systeminfo_command(context):
+    """Get System Info"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    message = primitive.hardware.get_system_info()
+    print_result(message=message, context=context)
+
+
+@cli.command("register")
+@click.pass_context
+def register_command(context):
+    """Register Hardware with Primitive"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    result = primitive.hardware.register()
+    color = "green" if result else "red"
+    if result:
+        message = "Hardware registered successfully"
+    else:
+        message = (
+            "There was an error registering this device. Please review the above logs."
+        )
+    print_result(message=message, context=context, fg=color)
+
+
+@cli.command("checkin")
+@click.pass_context
+def checkin_command(context):
+    """Checkin Hardware with Primitive"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    result = primitive.hardware.check_in_http()
+    if messages := result.get("checkIn").get("messages"):
+        print_result(message=messages, context=context, fg="yellow")
+    else:
+        message = "Hardware checked in successfully"
+        print_result(message=message, context=context, fg="green")

primitive/lint/actions.py

@@ -0,0 +1,6 @@
+from primitive.utils.actions import BaseAction
+
+
+class Lint(BaseAction):
+    def run_lint(self):
+        print("lint 100%")

primitive/lint/commands.py

@@ -0,0 +1,15 @@
+import click
+
+
+import typing
+
+if typing.TYPE_CHECKING:
+    from ..client import Primitive
+
+
+@click.command("lint")
+@click.pass_context
+def cli(context):
+    """Lint"""
+    primitive: Primitive = context.obj.get("PRIMITIVE")
+    primitive.lint.run_lint()

primitive/projects/actions.py

@@ -0,0 +1,55 @@
+from typing import List, Optional
+from gql import gql
+
+
+from primitive.utils.actions import BaseAction
+
+
+class Projects(BaseAction):
+    def get_job_run(self, id: str):
+        query = gql(
+            """
+            query jobRun($id: GlobalID!) {
+                jobRun(id: $id) {
+                    id
+                    organization {
+                        id
+                    }
+                }
+            }
+            """
+        )
+        variables = {"id": id}
+        result = self.primitive.session.execute(query, variable_values=variables)
+        return result
+
+    def job_run_update(
+        self,
+        id: str,
+        status: str = None,
+        conclusion: str = None,
+        file_ids: Optional[List[str]] = [],
+    ):
+        mutation = gql(
+            """
+            mutation jobRunUpdate($input: JobRunUpdateInput!) {
+                jobRunUpdate(input: $input) {
+                    ... on JobRun {
+                        id
+                        status
+                        conclusion
+                    }
+                }
+            }
+        """
+        )
+        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
+        variables = {"input": input}
+        result = self.primitive.session.execute(mutation, variable_values=variables)
+        return result

primitive/simulations/actions.py

@@ -0,0 +1,48 @@
+from gql import gql
+
+
+from primitive.utils.actions import BaseAction
+
+
+class Simulations(BaseAction):
+    def trace_create(
+        self,
+        id_code: str,
+        module: str,
+        var_type: str,
+        var_size: int,
+        reference: str,
+        bit_index: str,
+        timescale_unit: str,
+        timescale_magnitude: int,
+        organization: str,
+        file: str,
+        job_run: str,
+    ):
+        mutation = gql(
+            """
+            mutation createTrace($input: TraceCreateInput!) {
+                traceCreate(input: $input) {
+                    ... on Trace {
+                        id
+                    }
+                }
+            }
+        """
+        )
+        input = {
+            "idCode": id_code,
+            "module": module,
+            "varType": var_type,
+            "varSize": var_size,
+            "reference": reference,
+            "bitIndex": bit_index,
+            "timescaleUnit": timescale_unit,
+            "timescaleMagnitude": timescale_magnitude,
+            "organization": organization,
+            "file": file,
+            "jobRun": job_run,
+        }
+        variables = {"input": input}
+        result = self.primitive.session.execute(mutation, variable_values=variables)
+        return result

primitive/utils/actions.py

@@ -0,0 +1,9 @@
+import typing
+
+if typing.TYPE_CHECKING:
+    import primitive.client
+
+
+class BaseAction:
+    def __init__(self, primitive: "primitive.client.Primitive") -> None:
+        self.primitive = primitive

primitive/utils/config.py

@@ -0,0 +1,34 @@
+import json
+from pathlib import Path
+from typing import Dict
+
+HOME_DIRECTORY = Path.home()
+PRIMITIVE_CREDENTIALS_FILEPATH = Path(
+    HOME_DIRECTORY / ".config" / "primitive" / "credentials.json"
+)
+
+
+def create_directory(filepath: Path = PRIMITIVE_CREDENTIALS_FILEPATH):
+    filepath.mkdir(parents=True, exist_ok=True)
+
+
+def create_config_file(filepath: Path = PRIMITIVE_CREDENTIALS_FILEPATH):
+    filepath.parent.mkdir(parents=True, exist_ok=True)
+    filepath.touch()
+    with filepath.open("w") as json_file:
+        json.dump({}, json_file)
+
+
+def update_config_file(
+    filepath: Path = PRIMITIVE_CREDENTIALS_FILEPATH, new_config: Dict = {}
+):
+    existing_config = read_config_file(filepath=filepath)
+    merged_config = {**existing_config, **new_config}
+    with filepath.open("w") as json_file:
+        json.dump(merged_config, json_file, indent=2)
+
+
+def read_config_file(filepath: Path = PRIMITIVE_CREDENTIALS_FILEPATH) -> Dict:
+    if not filepath.exists():
+        create_config_file(filepath=filepath)
+    return json.loads(filepath.read_text())

primitive/utils/git.py

@@ -0,0 +1,15 @@
+import os
+from loguru import logger
+
+
+def download_source(github_access_token, git_repository, git_ref) -> None:
+    # Download code to current directory
+    logger.debug(f"Downloading source code from {git_repository} {git_ref}")
+    url = f"https://api.github.com/repos/{git_repository}/tarball/{git_ref}"
+    # TODO: switch to supbrocess.run or subprocess.Popen
+    result = os.system(
+        f"curl -s -L -H 'Accept: application/vnd.github+json' -H 'Authorization: Bearer {github_access_token}' -H 'X-GitHub-Api-Version: 2022-11-28' {url} | tar zx --strip-components 1 -C ."
+    )
+
+    if result != 0:
+        raise Exception("Failed to import repository.")