primitive 0.2.35__py3-none-any.whl → 0.2.37__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.2.35"
+__version__ = "0.2.37"

primitive/agent/actions.py

@@ -1,5 +1,6 @@
 import sys
 from time import sleep
+from typing import Optional
 
 from loguru import logger
 
@@ -12,9 +13,7 @@ from primitive.utils.actions import BaseAction
 
 
 class Agent(BaseAction):
-    def execute(
-        self,
-    ):
+    def execute(self, job_run_id: Optional[str] = None):
         logger.remove()
         logger.add(
             sink=sys.stderr,
@@ -26,6 +25,12 @@ class Agent(BaseAction):
         logger.info("primitive agent")
         logger.info(f"Version: {__version__}")
 
+        # TODO: tighten logic for determining if we're running in a container
+        RUNNING_IN_CONTAINER = False
+        if job_run_id is not None:
+            logger.info("Running in container...")
+            RUNNING_IN_CONTAINER = True
+
         # Wait for monitor to make database
         wait_for_db()
 
@@ -100,6 +105,10 @@ class Agent(BaseAction):
                     # This should probably eventually be another daemon?
                     uploader.scan()
 
+                if RUNNING_IN_CONTAINER:
+                    logger.info("Running in container, exiting after job run")
+                    break
+
                 sleep(5)
         except KeyboardInterrupt:
             logger.info("Stopping primitive agent...")

primitive/agent/commands.py

@@ -7,8 +7,9 @@ if typing.TYPE_CHECKING:
 
 
 @click.command("agent")
+@click.option("--job-run-id", type=str, help="Explicit Job Run to pull")
 @click.pass_context
-def cli(context):
+def cli(context, job_run_id: typing.Optional[str] = None):
     """agent"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
-    primitive.agent.execute()
+    primitive.agent.execute(job_run_id=job_run_id)

primitive/agent/runner.py

@@ -3,7 +3,6 @@ import os
 import re
 import shutil
 import typing
-from abc import abstractmethod
 from enum import Enum, IntEnum
 from pathlib import Path, PurePath
 from typing import Dict, List, TypedDict
@@ -13,6 +12,7 @@ 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
 
 try:
@@ -57,16 +57,6 @@ class LogLevel(Enum):
     WARNING = "WARNING"
 
 
-# Log Counter
-class LogCounter:
-    count = 0
-
-    @classmethod
-    def next(cls) -> int:
-        cls.count += 1
-        return cls.count
-
-
 class Runner:
     def __init__(
         self,
@@ -79,7 +69,7 @@ class Runner:
         self.job_run = job_run
         self.job_settings = job_run["jobSettings"]
         self.config = None
-        self.source_dir: Path = None
+        self.source_dir: Path | None = None
         self.initial_env = {}
         self.modified_env = {}
         self.file_logger = None
@@ -91,11 +81,11 @@ class Runner:
             self.file_logger = logger.add(
                 Path(get_logs_cache(self.job_run["id"]) / log_name),
                 rotation=max_log_size,
-                format=Runner.fmt,
+                format=fmt,
                 backtrace=True,
-                diagnose=True,
             )
 
+    @log_context(label="setup")
     def setup(self) -> None:
         # Attempt to download the job source code
         git_repo_full_name = self.job_run["gitCommit"]["repoFullName"]
@@ -155,66 +145,62 @@ class Runner:
             self.job_run["gitCommit"]["repoFullName"]
         )
 
+    @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"
         )
-        self.modified_env = {**self.initial_env}
 
+        self.modified_env = {**self.initial_env}
         task_failed = False
         cancelled = False
 
         for task in self.config["executes"]:
-            # the get status check here is to ensure that if cancel is called
-            # while one task is running, we do not run any OTHER laebeled tasks
-            # THIS is required for MULTI STEP JOBS
+            # Everything inside this loop should be contextualized with the task label
+            # this way we aren't jumping back and forth between the task label and "execute"
+            with logger.contextualize(label=task["label"]):
+                # the get status check here is to ensure that if cancel is called
+                # while one task is running, we do not run any OTHER labeled tasks
+                # THIS is required for MULTI STEP JOBS
+                status = self.primitive.jobs.get_job_status(self.job_run["id"])
+                status_value = status.data["jobRun"]["status"]
+                conclusion_value = status.data["jobRun"]["conclusion"]
+
+                if status_value == "completed" and conclusion_value == "cancelled":
+                    cancelled = True
+                    break
+
+                # Everything within this block should be contextualized as user logs
+                with logger.contextualize(type="user"):
+                    with asyncio.Runner() as async_runner:
+                        if task_failed := async_runner.run(self.run_task(task)):
+                            break
+
+        # FOR NONE MULTI STEP JOBS
+        # we still have to check that the job was cancelled here as well
+        with logger.contextualize(label="conclusion"):
             status = self.primitive.jobs.get_job_status(self.job_run["id"])
             status_value = status.data["jobRun"]["status"]
             conclusion_value = status.data["jobRun"]["conclusion"]
-
             if status_value == "completed" and conclusion_value == "cancelled":
                 cancelled = True
-                break
 
-            with logger.contextualize(label=task["label"]):
-                with asyncio.Runner() as async_runner:
-                    if task_failed := async_runner.run(self.run_task(task)):
-                        break
+            if cancelled:
+                logger.warning("Job cancelled by user")
+                return
 
-        number_of_files_produced = self.get_number_of_files_produced()
-        logger.info(
-            f"Produced {number_of_files_produced} files for {self.job['slug']} job"
-        )
+            conclusion = "success"
+            if task_failed:
+                conclusion = "failure"
+            else:
+                logger.success(f"Completed {self.job['slug']} job")
 
-        # FOR NONE MULTI STEP JOBS
-        # we still have to check that the job was cancelled here as well
-        status = self.primitive.jobs.get_job_status(self.job_run["id"])
-        status_value = status.data["jobRun"]["status"]
-        conclusion_value = status.data["jobRun"]["conclusion"]
-        if status_value == "completed" and conclusion_value == "cancelled":
-            cancelled = True
-
-        if cancelled:
-            logger.warning("Job cancelled by user")
             self.primitive.jobs.job_run_update(
                 self.job_run["id"],
-                number_of_files_produced=number_of_files_produced,
+                status="request_completed",
+                conclusion=conclusion,
             )
-            return
-
-        conclusion = "success"
-        if task_failed:
-            conclusion = "failure"
-        else:
-            logger.success(f"Completed {self.job['slug']} job")
-
-        self.primitive.jobs.job_run_update(
-            self.job_run["id"],
-            status="request_completed",
-            conclusion=conclusion,
-            number_of_files_produced=number_of_files_produced,
-        )
 
     def get_number_of_files_produced(self) -> int:
         """Returns the number of files produced by the job."""
@@ -315,6 +301,7 @@ class Runner:
                 logger.error(f"Task {task['label']} failed on '{cmd}'")
                 return True
 
+        logger.success(f"Completed {task['label']} task")
         return False
 
     async def log_cmd(self, process, stream, tags: Dict = {}) -> bool:
@@ -409,9 +396,8 @@ class Runner:
 
         return [line for line in lines if len(line) > 0]
 
+    @log_context(label="cleanup")
     def cleanup(self) -> None:
-        logger.remove(self.file_logger)
-
         if "stores" not in self.config:
             return
 
@@ -427,23 +413,13 @@ class Runner:
 
         shutil.rmtree(path=self.source_dir)
 
-    @abstractmethod
-    def fmt(record) -> str:
-        extra = record["extra"]
-        # Delimiters with empty space MUST exist for LogQL pattern matching
-        label = extra.get("label", None)
-        tag = extra.get("tag", None)
-        context = f"{label} | " if label else " | "
-        context += f"{tag} | " if tag else " | "
-
-        log = (
-            f"{LogCounter.next()} | "
-            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS!UTC}</green> | "
-            "<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | "
-            "<level>{level}</level> | "
-            f"{context}"
-            "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - "
-            "<level>{message}</level>\n"
+        number_of_files_produced = self.get_number_of_files_produced()
+        logger.info(
+            f"Produced {number_of_files_produced} files for {self.job['slug']} job"
+        )
+        self.primitive.jobs.job_run_update(
+            self.job_run["id"],
+            number_of_files_produced=number_of_files_produced,
         )
 
-        return log
+        logger.remove(self.file_logger)

primitive/utils/logging.py

@@ -0,0 +1,45 @@
+import json
+from datetime import timezone
+from functools import wraps
+
+from loguru import logger
+
+
+def log_context(**context):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with logger.contextualize(**context):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def fmt(record) -> str:
+    extra = record["extra"]
+    label = extra.get("label", None)
+    tag = extra.get("tag", None)
+    type = extra.get("type", "system")
+
+    context_object = {
+        "label": label,
+        "type": type,
+        "utc": record["time"]
+        .astimezone(timezone.utc)
+        .strftime("%Y-%m-%d %H:%M:%S.%f")[:-3],
+        "level": record["level"].name,
+        "name": record["name"],
+        "function": record["function"],
+        "line": record["line"],
+        "message": record["message"],
+    }
+
+    if tag:
+        context_object["tag"] = tag
+
+    # Loguru will fail if you return a string that doesn't select
+    # something within its record
+    record["extra"]["serialized"] = json.dumps(context_object)
+    return "{extra[serialized]}\n"

{primitive-0.2.35.dist-info → primitive-0.2.37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: primitive
-Version: 0.2.35
+Version: 0.2.37
 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

{primitive-0.2.35.dist-info → primitive-0.2.37.dist-info}/RECORD

@@ -1,11 +1,11 @@
-primitive/__about__.py,sha256=hEyqfjOzciSUTNTdMBjYfDhPQ78HUi4rV_w9imBssec,130
+primitive/__about__.py,sha256=k_4SlrPFVlD5xhSWVsBLRpaxZRnyooKrAmVPOvghuL0,130
 primitive/__init__.py,sha256=bwKdgggKNVssJFVPfKSxqFMz4IxSr54WWbmiZqTMPNI,106
 primitive/cli.py,sha256=g7EtHI9MATAB0qQu5w-WzbXtxz_8zu8z5E7sETmMkKU,2509
 primitive/client.py,sha256=h8WZVnQylVe0vbpuyC8YZHl2JyITSPC-1HbUcmrE5pc,3623
 primitive/agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-primitive/agent/actions.py,sha256=VHlNTw2M-T1MRajpIBu2weT8AsdDdYtlgbBnRBYYeco,3692
-primitive/agent/commands.py,sha256=cK7d3OcN5Z65gQWVZFQ-Y9ddw9Pes4f9OVBpeMsj5sE,255
-primitive/agent/runner.py,sha256=u1uEJiXUi2ps7wHFYo5iOW4WURo4ObnApUbzhsG-6AU,15784
+primitive/agent/actions.py,sha256=vq4_CZy7Mn37jUg-JPuJlMLtcx4sMXp6tUNvXGZOOyQ,4128
+primitive/agent/commands.py,sha256=kqa-PGqmzS-APd4BSMAkX4l8SdK5N1PRQwq9S8jBjsw,390
+primitive/agent/runner.py,sha256=ZlFC7Eq0NXOSCOvMtDu9gWKF2sa4QMkeG-wH-R-9iOc,15463
 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
@@ -92,12 +92,13 @@ primitive/utils/chunk_size.py,sha256=PAuVuirUTA9oRXyjo1c6MWxo31WVBRkWMuWw-AS58Bw
 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/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/shell.py,sha256=Z4zxmOaSyGCrS0D6I436iQci-ewHLt4UxVg1CD9Serc,2171
 primitive/utils/text.py,sha256=XiESMnlhjQ534xE2hMNf08WehE1SKaYFRNih0MmnK0k,829
-primitive-0.2.35.dist-info/METADATA,sha256=PLtjAKPLgi5K7IkgS3mqqeroGABE4fnmDTnSLYZE0P4,3569
-primitive-0.2.35.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-primitive-0.2.35.dist-info/entry_points.txt,sha256=p1K8DMCWka5FqLlqP1sPek5Uovy9jq8u51gUsP-z334,48
-primitive-0.2.35.dist-info/licenses/LICENSE.txt,sha256=B8kmQMJ2sxYygjCLBk770uacaMci4mPSoJJ8WoDBY_c,1098
-primitive-0.2.35.dist-info/RECORD,,
+primitive-0.2.37.dist-info/METADATA,sha256=r-F7XhBMe84oxARfJ9ZMbwVAcV5jQojsGrhxvWlK9E0,3569
+primitive-0.2.37.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+primitive-0.2.37.dist-info/entry_points.txt,sha256=p1K8DMCWka5FqLlqP1sPek5Uovy9jq8u51gUsP-z334,48
+primitive-0.2.37.dist-info/licenses/LICENSE.txt,sha256=B8kmQMJ2sxYygjCLBk770uacaMci4mPSoJJ8WoDBY_c,1098
+primitive-0.2.37.dist-info/RECORD,,