primitive 0.1.59__tar.gz → 0.1.62__tar.gz

{primitive-0.1.59 → primitive-0.1.62}/.vscode/settings.json

@@ -3,6 +3,7 @@
     "Checkin",
     "loguru",
     "machdep",
+    "pkey",
     "testsuites",
     "untar",
     "Untaring",

{primitive-0.1.59 → primitive-0.1.62}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: primitive
-Version: 0.1.59
+Version: 0.1.62
 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
@@ -21,6 +21,7 @@ Requires-Python: >=3.11
 Requires-Dist: click
 Requires-Dist: gql[all]
 Requires-Dist: loguru
+Requires-Dist: paramiko[all]
 Requires-Dist: primitive-pal==0.1.4
 Requires-Dist: pyyaml
 Description-Content-Type: text/markdown

{primitive-0.1.59 → primitive-0.1.62}/pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
   "loguru",
   "pyyaml",
   "primitive-pal == 0.1.4",
+  "paramiko[all]"
 ]
 
 [tool.uv]

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/__about__.py

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

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/agent/actions.py

@@ -129,6 +129,7 @@ class Agent(BaseAction):
                             source_dir=source_dir,
                             job_id=job_run["id"],
                             job_slug=job_run["job"]["slug"],
+                            max_log_size=500 * 1024,
                         )
                     except Exception as e:
                         # Log Error

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/agent/process.py

@@ -34,7 +34,7 @@ class Process:
         self.sel.register(self.process.stdout, selectors.EVENT_READ)
         self.sel.register(self.process.stderr, selectors.EVENT_READ)
 
-    def probe_logs(self):
+    def log(self):
         for key, _ in self.sel.select():
             data = key.fileobj.readline()
             if not data:
@@ -48,7 +48,7 @@ class Process:
 
     def wait(self):
         while True:
-            self.probe_logs()
+            self.log()
             if not self.is_running():
                 break
 

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/agent/provision.py

@@ -2,6 +2,7 @@ import sys
 from subprocess import Popen, PIPE
 from pathlib import Path
 from typing import Dict
+from ..utils.shell import env_string_to_dict
 
 
 class ProvisionPython:
@@ -34,14 +35,7 @@ class ProvisionPython:
         output, _ = proc.communicate()
 
         # Split the output into lines and parse it into a dictionary
-        env_vars = {}
-
-        for line in output.splitlines():
-            var_line = line.split("=", 1)
-
-            if len(var_line) == 2:
-                key, value = var_line
-                env_vars[key] = value
+        env_vars = env_string_to_dict(output)
 
         cmd = f"python -m pip install -r {self.requirements_path}"
         proc = Popen(

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/agent/runner.py

@@ -7,11 +7,10 @@ from typing import Dict, Iterable, List, Optional, TypedDict
 
 import yaml
 from loguru import logger
-
-from ..utils.cache import get_artifacts_cache
-from ..utils.files import find_files_for_extension
 from .process import Process
 from .provision import ProvisionPython
+from ..utils.cache import get_artifacts_cache, get_logs_cache
+from ..utils.files import find_files_for_extension
 
 try:
     from yaml import CLoader as Loader
@@ -56,12 +55,11 @@ class AgentRunner:
         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
 
         logger.enable("primitive")
-        self.logger_handle = logger.add(
-            Path(self.artifacts_dir / "runner_{time}.log"),
-            rotation=self.max_log_size,  # Rotate when the log file reaches 10MB
-        )
+        self.swap_logs(label="init")
 
         logger.info(f"Scanning directory for job {self.job_slug}")
 
@@ -113,6 +111,9 @@ class AgentRunner:
         conclusion = None
         total_errors = 0
         for step in self.steps():
+            # Swap logger
+            self.swap_logs(label=step["name"])
+
             logger.info(f"Beginning step {step['name']}")
 
             # Update workdir
@@ -138,7 +139,7 @@ class AgentRunner:
                 while proc.is_running():
                     # Check job status
                     status = self.primitive.jobs.get_job_status(self.job_id)
-                    status_value = status["jobRun"]["status"]
+                    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
@@ -184,6 +185,16 @@ class AgentRunner:
         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)
+
+        self.logger_handle = logger.add(
+            Path(self.logs_dir / f"{label}_{{time}}.primitive.log"),
+            rotation=self.max_log_size,
+        )
+
     def provision(self) -> Optional[Dict]:
         match self.job["provision"]:
             case "python":
@@ -199,8 +210,6 @@ class AgentRunner:
                 return prov.create_env()
 
     def collect_artifacts(self, step: JobStep) -> None:
-        # str(PurePath(file_path).relative_to(Path(source))
-
         # Search each artifact type
         for artifact in step["artifacts"]:
             files = find_files_for_extension(self.source_dir, artifact["extension"])

primitive-0.1.62/src/primitive/agent/uploader.py

@@ -0,0 +1,104 @@
+import typing
+from typing import Dict
+import shutil
+import os
+from pathlib import Path, PurePath
+
+from loguru import logger
+
+from ..utils.cache import get_artifacts_cache, get_logs_cache
+
+if typing.TYPE_CHECKING:
+    import primitive.client
+
+
+class Uploader:
+    def __init__(
+        self,
+        primitive: "primitive.client.Primitive",
+    ):
+        self.primitive = primitive
+
+    def upload_dir(self, cache: Path) -> Dict:
+        file_ids = []
+        job_run_id = cache.name
+
+        files = None
+        has_walk = getattr(cache, "walk", None)
+        if has_walk:
+            files = sorted(
+                [
+                    current_path / file
+                    for current_path, _, current_path_files in cache.walk()
+                    for file in current_path_files
+                ],
+                key=lambda p: p.stat().st_size,
+            )
+        else:
+            files = sorted(
+                [
+                    Path(Path(current_path) / file)
+                    for current_path, _, current_path_files in os.walk(cache)
+                    for file in current_path_files
+                ],
+                key=lambda p: p.stat().st_size,
+            )
+
+        for file in files:
+            response = self.primitive.files.upload_file_via_api(
+                path=file,
+                key_prefix=str(PurePath(file).relative_to(cache.parent).parent),
+                job_run_id=job_run_id,
+            )
+            upload_id = response.json()["data"]["fileUpload"]["id"]
+
+            if upload_id:
+                file_ids.append(upload_id)
+                continue
+
+            logger.error(f"Unable to upload file {file}")
+
+        # Clean up job cache
+        shutil.rmtree(path=cache)
+
+        return {job_run_id: file_ids}
+
+    def scan(self) -> None:
+        # Scan artifacts directory
+        artifacts_dir = get_artifacts_cache()
+        logs_dir = get_logs_cache()
+
+        artifacts = sorted(
+            [
+                artifacts_cache
+                for artifacts_cache in artifacts_dir.iterdir()
+                if artifacts_cache.is_dir()
+            ],
+            key=lambda p: p.stat().st_ctime,
+        )
+
+        logs = sorted(
+            [logs_cache for logs_cache in logs_dir.iterdir() if logs_cache.is_dir()],
+            key=lambda p: p.stat().st_ctime,
+        )
+
+        log_files = {
+            job_id: files
+            for log_path in logs
+            for job_id, files in self.upload_dir(log_path).items()
+        }
+
+        artifact_files = {
+            job_id: files
+            for artifact_path in artifacts
+            for job_id, files in self.upload_dir(artifact_path).items()
+        }
+
+        files_by_id = {
+            job_id: log_files.get(job_id, []) + artifact_files.get(job_id, [])
+            for job_id in log_files.keys() | artifact_files.keys()
+        }
+
+        # Update job run
+        for job_id, files in files_by_id.items():
+            self.primitive.jobs.job_run_update(id=job_id, file_ids=files)

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/exec/actions.py

@@ -1,9 +1,14 @@
+import tempfile
 import typing
 
+from primitive.exec.interactive import interactive_shell
+
 if typing.TYPE_CHECKING:
     pass
 
 
+from paramiko import SSHClient
+
 from primitive.utils.actions import BaseAction
 
 
@@ -41,7 +46,48 @@ class Exec(BaseAction):
             reservation_id=reservation["id"], desired_status="in_progress"
         )
 
-        print(f"Executing command: {command} on {hardware['name']}")
+        ssh_credentials = self.primitive.hardware.get_hardware_ssh_credentials(
+            hardware_id=hardware["id"]
+        )
+
+        ssh_hostname = ssh_credentials["ssh_hostname"]
+        ssh_username = ssh_credentials["ssh_username"]
+        ssh_password = ssh_credentials["ssh_password"]
+        ssh_private_key = ssh_credentials["ssh_private_key"]
+
+        ssh_client = SSHClient()
+        ssh_client.load_system_host_keys()
+        keyfile = None
+        if ssh_private_key:
+            keyfile = tempfile.NamedTemporaryFile()
+            keyfile.write(ssh_private_key.encode())
+            keyfile.flush()
+            ssh_client.connect(
+                hostname=ssh_hostname,
+                username=ssh_username,
+                key_filename=keyfile.name,
+            )
+        else:
+            ssh_client.connect(
+                hostname=ssh_hostname,
+                username=ssh_username,
+                pkey=ssh_private_key,
+                password=ssh_password,
+            )
+
+        if command:
+            formatted_command = " ".join(command)
+            stdin, stdout, stderr = ssh_client.exec_command(formatted_command)
+            print(stdout.read())
+            ssh_client.close()
+        else:
+            channel = ssh_client.get_transport().open_session()
+            channel.get_pty()
+            channel.invoke_shell()
+            interactive_shell(channel)
+            ssh_client.close()
+            if keyfile:
+                keyfile.close()
 
         if created_reservation_on_behalf_of_user:
             print("Cleaning up reservation.")

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/exec/commands.py

@@ -13,7 +13,7 @@ if typing.TYPE_CHECKING:
     type=str,
     required=True,
 )
-@click.argument("command", nargs=-1, required=True)
+@click.argument("command", nargs=-1, required=False)
 def cli(context, hardware_identifier: str, command: str) -> None:
     """Exec"""
     primitive: Primitive = context.obj.get("PRIMITIVE")

primitive-0.1.62/src/primitive/exec/interactive.py

@@ -0,0 +1,99 @@
+# https://github.com/paramiko/paramiko/blob/main/demos/interactive.py
+
+import socket
+import sys
+
+
+def _to_unicode(s):
+    """
+    decode a string as ascii or utf8 if possible (as required by the sftp
+    protocol).  if neither works, just return a byte string because the server
+    probably doesn't know the filename's encoding.
+    """
+    try:
+        return s.encode("ascii")
+    except (UnicodeError, AttributeError):
+        try:
+            return s.decode("utf-8")
+        except UnicodeError:
+            return s
+
+
+# windows does not have termios...
+try:
+    import termios
+    import tty
+
+    has_termios = True
+except ImportError:
+    has_termios = False
+
+
+def interactive_shell(chan):
+    if has_termios:
+        posix_shell(chan)
+    else:
+        windows_shell(chan)
+
+
+def posix_shell(chan):
+    import select
+
+    oldtty = termios.tcgetattr(sys.stdin)
+    try:
+        tty.setraw(sys.stdin.fileno())
+        tty.setcbreak(sys.stdin.fileno())
+        chan.settimeout(0.0)
+
+        while True:
+            r, w, e = select.select([chan, sys.stdin], [], [])
+            if chan in r:
+                try:
+                    x = _to_unicode(chan.recv(1024))
+                    if len(x) == 0:
+                        sys.stdout.write("\r\n*** EOF\r\n")
+                        break
+                    sys.stdout.write(x)
+                    sys.stdout.flush()
+                except socket.timeout:
+                    pass
+            if sys.stdin in r:
+                x = sys.stdin.read(1)
+                if len(x) == 0:
+                    break
+                chan.send(x)
+
+    finally:
+        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
+
+
+# thanks to Mike Looijmans for this code
+def windows_shell(chan):
+    import threading
+
+    sys.stdout.write(
+        "Line-buffered terminal emulation. Press F6 or ^Z to send EOF.\r\n\r\n"
+    )
+
+    def writeall(sock):
+        while True:
+            data = sock.recv(256)
+            if not data:
+                sys.stdout.write("\r\n*** EOF ***\r\n\r\n")
+                sys.stdout.flush()
+                break
+            sys.stdout.write(data)
+            sys.stdout.flush()
+
+    writer = threading.Thread(target=writeall, args=(chan,))
+    writer.start()
+
+    try:
+        while True:
+            d = sys.stdin.read(1)
+            if not d:
+                break
+            chan.send(d)
+    except EOFError:
+        # user hit ^Z or F6
+        pass

primitive-0.1.62/src/primitive/files/actions.py

@@ -0,0 +1,189 @@
+import hashlib
+import sys
+import threading
+from pathlib import Path
+from typing import Dict, Optional
+
+import requests
+from gql import gql
+from loguru import logger
+
+from primitive.graphql.sdk import create_requests_session
+from primitive.utils.actions import BaseAction
+
+from ..utils.auth import guard
+from .graphql.mutations import (
+    file_update_mutation,
+    pending_file_create_mutation,
+)
+
+
+# this class can be used in multithreaded S3 client uploader
+# this requires getting an S3 access token to this machine however
+# we are using presigned urls instead at this time Oct 29th, 2024
+class ProgressPercentage(object):
+    def __init__(self, filepath: Path) -> None:
+        self._filename = filepath.name
+        self._size = float(filepath.stat().st_size)
+        self._seen_so_far = 0
+        self._lock = threading.Lock()
+
+    def __call__(self, bytes_amount):
+        # To simplify, assume this is hooked up to a single filename
+        with self._lock:
+            self._seen_so_far += bytes_amount
+            percentage = (self._seen_so_far / self._size) * 100
+            sys.stdout.write(
+                "\r%s  %s / %s  (%.2f%%)"
+                % (self._filename, self._seen_so_far, self._size, percentage)
+            )
+            sys.stdout.flush()
+
+
+class Files(BaseAction):
+    def _pending_file_create(
+        self,
+        file_name: str,
+        file_size: int,
+        file_checksum: str,
+        file_path: str,
+        key_prefix: str,
+        is_public: bool = False,
+    ):
+        mutation = gql(pending_file_create_mutation)
+        input = {
+            "filePath": file_path,
+            "fileName": file_name,
+            "fileSize": file_size,
+            "fileChecksum": file_checksum,
+            "keyPrefix": key_prefix,
+            "isPublic": is_public,
+        }
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result.data.get("pendingFileCreate")
+
+    def _update_file_status(
+        self,
+        file_id: str,
+        is_uploading: Optional[bool] = None,
+        is_complete: Optional[bool] = None,
+    ):
+        mutation = gql(file_update_mutation)
+        input: Dict[str, str | bool] = {
+            "id": file_id,
+        }
+        if is_uploading is not None:
+            input["isUploading"] = is_uploading
+        if is_complete is not None:
+            input["isComplete"] = is_complete
+
+        variables = {"input": input}
+        result = self.primitive.session.execute(
+            mutation, variable_values=variables, get_execution_result=True
+        )
+        return result
+
+    @guard
+    def upload_file_direct(
+        self,
+        path: Path,
+        is_public: False,
+        key_prefix: str = "",
+        file_id: Optional[str] = None,
+    ):
+        logger.enable("primitive")
+        if path.exists() is False:
+            raise Exception(f"File {path} does not exist.")
+
+        file_size = path.stat().st_size
+        if file_size == 0:
+            raise Exception(f"{path} is empty.")
+
+        file_checksum = hashlib.md5(path.read_bytes()).hexdigest()
+
+        if not file_id:
+            pending_file_create = self._pending_file_create(
+                file_name=path.name,
+                file_size=path.stat().st_size,
+                file_checksum=file_checksum,
+                file_path=str(path),
+                key_prefix=key_prefix,
+                is_public=is_public,
+            )
+        file_id = pending_file_create.get("id")
+        presigned_url = pending_file_create.get("presignedUrlForUpload")
+
+        if not file_id:
+            raise Exception("No file_id found or provided.")
+        if not presigned_url:
+            raise Exception("No presigned_url returned.")
+
+        self._update_file_status(file_id, is_uploading=True)
+        with open(path, "rb") as object_file:
+            object_text = object_file.read()
+            response = requests.put(presigned_url, data=object_text)
+            if response.ok:
+                logger.info(f"File {path} uploaded successfully.")
+                update_file_status_result = self._update_file_status(
+                    file_id, is_uploading=False, is_complete=True
+                )
+            else:
+                message = f"Failed to upload file {path}. {response.status_code}: {response.text}"
+                logger.error(message)
+                raise Exception(message)
+        file_pk = update_file_status_result.data.get("fileUpdate").get("pk")
+        file_access_url = f"{self.primitive.host_config.get("transport")}://{self.primitive.host}/files/{file_pk}/presigned-url/"
+        logger.info(f"Available at: {file_access_url}")
+        return update_file_status_result
+
+    @guard
+    def upload_file_via_api(
+        self,
+        path: Path,
+        is_public: bool = False,
+        key_prefix: str = "",
+        job_run_id: str = "",
+    ):
+        """
+        This method uploads a file via the Primitive API.
+        This does NOT upload the file straight to S3
+        """
+        file_path = str(path.resolve())
+        if path.exists() is False:
+            raise FileNotFoundError(f"File not found at {file_path}")
+
+        if is_public:
+            operations = (
+                """{ "query": "mutation fileUpload($input: FileUploadInput!) { fileUpload(input: $input) { ... on File { id } } }", "variables": { "input": { "fileObject": null, "isPublic": true, "filePath": \""""
+                + file_path
+                + """\", "keyPrefix": \""""
+                + key_prefix
+                + """\", "jobRunId": \""""
+                + job_run_id
+                + """\" } } }"""
+            )  # noqa
+
+        else:
+            operations = (
+                """{ "query": "mutation fileUpload($input: FileUploadInput!) { fileUpload(input: $input) { ... on File { id } } }", "variables": { "input": { "fileObject": null, "isPublic": false, "filePath": \""""
+                + file_path
+                + """\", "keyPrefix": \""""
+                + key_prefix
+                + """\", "jobRunId": \""""
+                + job_run_id
+                + """\" } } }"""
+            )  # noqa
+        body = {
+            "operations": ("", operations),
+            "map": ("", '{"fileObject": ["variables.input.fileObject"]}'),
+            "fileObject": (path.name, open(path, "rb")),
+        }
+
+        session = create_requests_session(self.primitive.host_config)
+        transport = self.primitive.host_config.get("transport")
+        url = f"{transport}://{self.primitive.host}/"
+        response = session.post(url, files=body)
+        return response

{primitive-0.1.59 → primitive-0.1.62}/src/primitive/files/commands.py

@@ -1,7 +1,10 @@
+import json
+import typing
+from pathlib import Path
+
 import click
+
 from ..utils.printer import print_result
-from pathlib import Path
-import typing
 
 if typing.TYPE_CHECKING:
     from ..client import Primitive
@@ -19,10 +22,19 @@ def cli(context):
 @click.argument("path", type=click.Path(exists=True))
 @click.option("--public", "-p", help="Is this a Public file", is_flag=True)
 @click.option("--key-prefix", "-k", help="Key Prefix", default="")
-def file_upload_command(context, path, public, key_prefix):
+@click.option("--direct", "-k", help="direct", is_flag=True)
+def file_upload_command(context, path, public, key_prefix, direct):
     """File Upload"""
     primitive: Primitive = context.obj.get("PRIMITIVE")
     path = Path(path)
-    result = primitive.files.file_upload(path, is_public=public, key_prefix=key_prefix)
-    message = result.json()
+    if direct:
+        result = primitive.files.upload_file_direct(
+            path, is_public=public, key_prefix=key_prefix
+        )
+    else:
+        result = primitive.files.upload_file_via_api(
+            path, is_public=public, key_prefix=key_prefix
+        )
+
+    message = json.dumps(result.data)
     print_result(message=message, context=context)