primitive 0.1.68__py3-none-any.whl → 0.1.69__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.68"
+__version__ = "0.1.69"

primitive/agent/process.py

@@ -41,7 +41,14 @@ class Process:
                 continue
 
             if key.fileobj is self.process.stdout:
-                logger.info(data.rstrip())
+                raw_data = data.rstrip()
+                if "error" in raw_data.lower():
+                    logger.error(raw_data)
+                    self._errors += 1
+                elif "warning" in raw_data.lower():
+                    logger.warning(raw_data)
+                else:
+                    logger.info(raw_data)
             elif key.fileobj is self.process.stderr:
                 logger.error(data.rstrip())
                 self._errors += 1

primitive/agent/runner.py

@@ -1,9 +1,10 @@
 import os
 import threading
 import typing
+import json
 from pathlib import Path, PurePath
 from time import sleep
-from typing import Dict, Iterable, List, Optional, TypedDict
+from typing import Dict, Iterable, List, Optional, TypedDict, Callable
 
 import yaml
 from loguru import logger
@@ -47,7 +48,7 @@ class AgentRunner:
         job_id: str,
         job_slug: str,
         max_log_size: int = 10 * 1024 * 1024,
-        job: Optional[JobDescription] = None,
+        log_to_file: bool = True,
     ) -> None:
         self.primitive = primitive
         self.source_dir = source_dir
@@ -55,22 +56,21 @@ class AgentRunner:
         self.job_id = job_id
         self.job_slug = job_slug
         self.max_log_size = max_log_size
-        self.artifacts_dir = get_artifacts_cache(self.job_id)
-        self.logs_dir = get_logs_cache(self.job_id)
-        self.logger_handle = None
+        self.log_to_file = log_to_file
 
+        # Enable and configure logger
         logger.enable("primitive")
-        self.swap_logs(label="init")
 
-        if job:
-            self.job = job
-        else:
-            self.load_job_from_file()
-    
-    def load_job_from_file(self) -> JobDescription:
-        logger.info(f"Scanning directory for job file {self.job_slug}")
+        if self.log_to_file:
+            log_name = f"{self.job_slug}_{self.job_id}_{{time}}.primitive.log"
+            logger.add(
+                Path(get_logs_cache(self.job_id) / log_name),
+                rotation=self.max_log_size,
+                format=AgentRunner.log_serializer(),
+            )
 
-        # Look for job based on slug
+        # Attempt to load job from file
+        logger.info(f"Scanning directory for job file {self.job_slug}")
         yaml_file = Path(self.source_dir / ".primitive" / f"{self.job_slug}.yaml")
         yml_file = Path(self.source_dir / ".primitive" / f"{self.job_slug}.yml")
 
@@ -78,7 +78,9 @@ class AgentRunner:
             logger.error(
                 f"Found two job descriptions with the same slug: {self.job_slug}"
             )
-            self.conclude(conclusion="failure")
+            self.primitive.jobs.job_run_update(
+                self.job_id, status="request_completed", conclusion="failure"
+            )
             raise FileExistsError
 
         if yaml_file.exists():
@@ -89,11 +91,35 @@ class AgentRunner:
             logger.error(
                 f"No job description with matching slug '{self.job_slug}' found"
             )
-            self.conclude(conclusion="failure")
+            self.primitive.jobs.job_run_update(
+                self.job_id, status="request_completed", conclusion="failure"
+            )
             raise FileNotFoundError
 
         logger.info(f"Found job description for {self.job_slug}")
 
+    @staticmethod
+    def log_serializer() -> Callable:
+        def fmt(record):
+            step = ""
+            if "step" in record["extra"]:
+                step = record["extra"]["step"]
+
+            log = {
+                "time": record["time"].strftime("%Y-%m-%d %H:%M:%S.%f"),
+                "utc": record["time"].strftime("%Y-%m-%d %H:%M:%S.%f%z"),
+                "level": record["level"].name,
+                "message": record["message"],
+                "name": record["name"],
+                "step": step,
+            }
+
+            record["extra"]["serialized"] = json.dumps(log)
+
+            return "{extra[serialized]}\n"
+
+        return fmt
+
     def name(self) -> str:
         return self.job["name"]
 
@@ -105,112 +131,85 @@ class AgentRunner:
         logger.info(f"Executing {self.job_slug} job")
         self.primitive.jobs.job_run_update(self.job_id, status="request_in_progress")
 
-        # Initial environment is the system env
+        # Initialize the environment with the system
         environment = os.environ
         if "provision" in self.job:
             logger.info(f"Provisioning for {self.job['provision']} environment")
             environment = self.provision()
 
             if not environment:
-                self.conclude("failure")
+                self.primitive.jobs.job_run_update(
+                    self.job_id, status="request_completed", conclusion="failure"
+                )
+                logger.error(f"{self.job_slug} concluded with error(s)")
                 return
 
-        conclusion = None
-        total_errors = 0
+        total_job_errors = 0
         for step in self.steps():
-            # Swap logger
-            self.swap_logs(label=step["name"])
-
             logger.info(f"Beginning step {step['name']}")
 
-            # Update workdir
-            if "workdir" in step:
-                self.workdir = step["workdir"]
-
-            # Define step proc
-            proc = Process(
-                cmd=step["cmd"],
-                workdir=Path(self.source_dir / self.workdir),
-                env=environment,
-            )
-
-            # Try to start
-            try:
-                proc.start()
-            except Exception as e:
-                logger.error(f"Error while attempting to run command {e}")
-                conclusion = "failure"
-                break
-
-            def status_check():
-                while proc.is_running():
-                    # Check job status
-                    status = self.primitive.jobs.get_job_status(self.job_id)
-                    status_value = status.data["jobRun"]["status"]
+            with logger.contextualize(step=step["name"]):
+                if "workdir" in step:
+                    self.workdir = step["workdir"]
+
+                proc = Process(
+                    cmd=step["cmd"],
+                    workdir=Path(self.source_dir / self.workdir),
+                    env=environment,
+                )
+
+                try:
+                    proc.start()
+                except Exception as e:
+                    logger.error(f"Error while attempting to run process {e}")
+                    self.primitive.jobs.job_run_update(
+                        self.job_id, status="request_completed", conclusion="failure"
+                    )
+                    logger.error(f"{self.job_slug} concluded with error(s)")
+                    return
+
+                def status_check():
+                    while proc.is_running():
+                        # Check job status
+                        status = self.primitive.jobs.get_job_status(self.job_id)
+                        status_value = status.data["jobRun"]["status"]
+
+                        # TODO: Should probably use request_cancelled or something
+                        # once we change it, we'll have to call conclude w/ cancelled status
+                        if status_value == "completed":
+                            logger.warning("Job cancelled by user")
+                            proc.terminate()
+                            return
+
+                        sleep(5)
+
+                status_thread = threading.Thread(target=status_check)
+                status_thread.start()
+
+                returncode = proc.wait()
+                total_job_errors += proc.errors
+                status_thread.join()
+
+            self.collect_artifacts(step)
 
-                    # TODO: Should probably use request_cancelled or something
-                    # once we change it, we'll have to call conclude w/ cancelled status
-                    if status_value == "completed":
-                        logger.warning("Job cancelled by user")
-                        proc.terminate()
-                        return
-
-                    sleep(5)
-
-            status_thread = threading.Thread(target=status_check)
-            status_thread.start()
-
-            # Wait for proc to finish
-            returncode = proc.wait()
-            total_errors += proc.errors
-
-            # Wait for status check
-            status_thread.join()
-
-            # Collect artifacts
-            if "artifacts" in step:
-                self.collect_artifacts(step)
-
-            # Check if we have a good result
             if returncode > 0:
-                conclusion = "failure"
-                break
-
-        if not conclusion and total_errors == 0:
-            conclusion = "success"
-        else:
-            logger.error(f"Job failed with {total_errors} errors.")
-            conclusion = "failure"
+                self.primitive.jobs.job_run_update(
+                    self.job_id, status="request_completed", conclusion="failure"
+                )
+                logger.error(f"{self.job_slug} concluded with error(s)")
+                return
 
-        self.conclude(conclusion)
+        if total_job_errors > 0:
+            self.primitive.jobs.job_run_update(
+                self.job_id, status="request_completed", conclusion="failure"
+            )
+            logger.error(f"{self.job_slug} concluded with error(s)")
+            return
 
-    def conclude(self, conclusion: str) -> None:
         self.primitive.jobs.job_run_update(
-            self.job_id, status="request_completed", conclusion=conclusion
-        )
-
-        logger.info(f"Completed {self.job_slug} job")
-        logger.remove(self.logger_handle)
-
-    def swap_logs(self, label: str):
-        # Remove Handle
-        if self.logger_handle:
-            logger.remove(self.logger_handle)
-        
-        # Custom format for UTC time
-        logger_format = (
-            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS!UTC}</green> | "
-            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
-            "<level>{level: <8}</level> | "
-            "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
-            "<level>{message}</level>"
-        )
-
-        self.logger_handle = logger.add(
-            Path(self.logs_dir / f"{label}_{{time}}.primitive.log"),
-            rotation=self.max_log_size,
-            format=logger_format
+            self.job_id, status="request_completed", conclusion="success"
         )
+        logger.success(f"Completed {self.job_slug} job")
 
     def provision(self) -> Optional[Dict]:
         match self.job["provision"]:
@@ -227,7 +226,9 @@ class AgentRunner:
                 return prov.create_env()
 
     def collect_artifacts(self, step: JobStep) -> None:
-        # Search each artifact type
+        if "artifacts" not in step:
+            return
+
         for artifact in step["artifacts"]:
             files = find_files_for_extension(self.source_dir, artifact["extension"])
 
@@ -236,7 +237,7 @@ class AgentRunner:
                 relative_path = PurePath(file).relative_to(self.source_dir)
 
                 # Construct destination to preserve directory structure
-                destination = Path(self.artifacts_dir / relative_path)
+                destination = Path(get_artifacts_cache(self.job_id) / relative_path)
 
                 # Create directories if they don't exist
                 destination.parent.mkdir(parents=True, exist_ok=True)

primitive/agent/uploader.py

@@ -1,8 +1,8 @@
-import typing
-from typing import Dict
-import shutil
 import os
+import shutil
+import typing
 from pathlib import Path, PurePath
+from typing import Dict
 
 from loguru import logger
 
@@ -45,12 +45,11 @@ class Uploader:
             )
 
         for file in files:
-            response = self.primitive.files.upload_file_via_api(
+            result = self.primitive.files.upload_file_direct(
                 path=file,
-                key_prefix=str(PurePath(file).relative_to(cache.parent).parent),
-                job_run_id=job_run_id,
+                key_prefix=str(PurePath(file).relative_to(cache.parent).parent)
             )
-            upload_id = response.json()["data"]["fileUpload"]["id"]
+            upload_id = result.data["fileUpdate"]["id"]
 
             if upload_id:
                 file_ids.append(upload_id)

primitive/cli.py

@@ -13,11 +13,9 @@ from .files.commands import cli as file_commands
 from .git.commands import cli as git_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
 
 
 @click.group()
@@ -65,14 +63,12 @@ cli.add_command(config_command, "config")
 cli.add_command(whoami_command, "whoami")
 cli.add_command(file_commands, "files")
 cli.add_command(hardware_commands, "hardware")
-cli.add_command(lint_commands, "lint")
 cli.add_command(agent_commands, "agent")
 cli.add_command(git_commands, "git")
 cli.add_command(daemons_commands, "daemons")
 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")
 

primitive/client.py

@@ -1,5 +1,6 @@
 import sys
 
+from gql import Client
 from loguru import logger
 
 from .agent.actions import Agent
@@ -10,12 +11,10 @@ from .files.actions import Files
 from .git.actions import Git
 from .hardware.actions import Hardware
 from .jobs.actions import Jobs
-from .lint.actions import Lint
 from .organizations.actions import Organizations
 from .projects.actions import Projects
 from .provisioning.actions import Provisioning
 from .reservations.actions import Reservations
-from .sim.actions import Sim
 from .utils.config import read_config_file
 
 logger.disable("primitive")
@@ -30,10 +29,10 @@ class Primitive:
         token: str = None,
         transport: str = None,
     ) -> None:
-        self.host = host
-        self.session = None
-        self.DEBUG = DEBUG
-        self.JSON = JSON
+        self.host: str = host
+        self.session: Client = None
+        self.DEBUG: bool = DEBUG
+        self.JSON: bool = JSON
 
         if self.DEBUG:
             logger.enable("primitive")
@@ -61,10 +60,8 @@ class Primitive:
         self.projects: Projects = Projects(self)
         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)