huanld/speckle_automate_python_example
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: migrate to speckle automate sdk

Browse Source
This commit is contained in:
Gergő Jedlicska
2023-09-20 10:47:19 +02:00
parent 7379a23452
commit db8e01e0b0
6 changed files with 61 additions and 554 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.vscode/launch.json
Vendored
+4 -3
View File
@@ -13,10 +13,11 @@
"justMyCode": true,
"envFile": "${workspaceFolder}/.env",
"args": [
"run",
"{\"projectId\": \"843d07eb10\", \"modelId\": \"base design\", \"versionId\": \"2a32ccfee1\", \"speckleServerUrl\": \"https://latest.speckle.systems\"}",
// make sure to use camelCase for variable names
"{\"speckleTypeToCount\": \"Objects.Geometry.Brep\"}"
],
},
"{\"forbiddenSpeckleType\": \"Objects.Geometry.Brep\"}"
]
}
]
}
automation_context.py
-483
View File
@@ -1,483 +0,0 @@
"""WIP module for an automate python sdk."""
import json
import os
import sys
import time
import traceback
from collections import defaultdict
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Callable, TypeVar, overload
import httpx
from gql import gql
from pydantic import BaseModel, ConfigDict
from specklepy.api import operations
from specklepy.api.client import SpeckleClient
from specklepy.objects.base import Base
from specklepy.transports.memory import MemoryTransport
from specklepy.transports.server import ServerTransport
from stringcase import camelcase
class AutomateBase(BaseModel):
"""Use this class as a base model for automate related DTO."""
model_config = ConfigDict(alias_generator=camelcase, populate_by_name=True)
class AutomationRunData(BaseModel):
"""Values of the project / model that triggered the run of this function."""
project_id: str
model_id: str
branch_name: str
version_id: str
speckle_server_url: str
automation_id: str
automation_revision_id: str
automation_run_id: str
function_id: str
function_revision: str
model_config = ConfigDict(
alias_generator=camelcase, populate_by_name=True, protected_namespaces=()
)
class AutomationStatus(str, Enum):
"""Set the status of the automation."""
INITIALIZING = "INITIALIZING"
RUNNING = "RUNNING"
FAILED = "FAILED"
SUCCEEDED = "SUCCEEDED"
class ObjectResultLevel(str, Enum):
"""Possible status message levels for object reports."""
INFO = "INFO"
WARNING = "WARNING"
ERROR = "ERROR"
class ObjectResult(AutomateBase):
"""An object level result."""
level: ObjectResultLevel
status_message: str
class AutomationResult(AutomateBase):
"""Schema accepted by the Speckle server as a result for an automation run."""
elapsed: float = 0
result_view: str | None = None
result_versions: list[str] = field(default_factory=list)
blobs: list[str] = field(default_factory=list)
run_status: AutomationStatus = AutomationStatus.RUNNING
status_message: str | None = None
object_results: dict[str, list[ObjectResult]] = field(
default_factory=lambda: defaultdict(list) # typing: ignore
)
T = TypeVar("T", bound=AutomateBase)
@dataclass
class AutomationContext:
"""A WIP umbrella class for automate sdk functionality.
Potentially turn this into a context manager, to handle function enter exit status
changes.
"""
automation_run_data: AutomationRunData
speckle_client: SpeckleClient
_server_transport: ServerTransport
_speckle_token: str
#: keep a memory transponrt at hand, to speed up things if needed
_memory_transport: MemoryTransport = field(default_factory=MemoryTransport)
#: added for performance measuring
_init_time: float = field(default_factory=time.perf_counter)
_automation_result: AutomationResult = field(default_factory=AutomationResult)
@classmethod
def initialize(
cls, automation_run_data: str | AutomationRunData, speckle_token: str
) -> "AutomationContext":
"""Bootstrap the AutomateSDK from raw data.
Todo:
----
* bootstrap a structlog logger instance
* expose a logger, that ppl can use instead of print
* log an initialization message
"""
# parse the json value if its not an initialized project data instance
automation_run_data = (
automation_run_data
if isinstance(automation_run_data, AutomationRunData)
else AutomationRunData.model_validate_json(automation_run_data)
)
speckle_client = SpeckleClient(
automation_run_data.speckle_server_url,
automation_run_data.speckle_server_url.startswith("https"),
)
speckle_client.authenticate_with_token(speckle_token)
if not speckle_client.account:
msg = (
f"Could not autenticate to {automation_run_data.speckle_server_url}",
"with the provided token",
)
raise ValueError(msg)
server_transport = ServerTransport(
automation_run_data.project_id, speckle_client
)
return cls(automation_run_data, speckle_client, server_transport, speckle_token)
@property
def run_status(self) -> AutomationStatus:
"""Get the status of the automation run."""
return self._automation_result.run_status
def elapsed(self) -> float:
"""Return the elapsed time in seconds since the initialization time."""
return time.perf_counter() - self._init_time
def receive_version(self) -> Base:
"""Receive the Speckle project version that triggered this automation run."""
commit = self.speckle_client.commit.get(
self.automation_run_data.project_id, self.automation_run_data.version_id
)
if not commit.referencedObject:
raise ValueError("The commit has no referencedObject, cannot receive it.")
base = operations.receive(
commit.referencedObject, self._server_transport, self._memory_transport
)
print(
f"It took {self.elapsed():2f} seconds to receive",
f" the speckle version {self.automation_run_data.version_id}",
)
return base
def create_new_version_in_project(
self, root_object: Base, model_id: str, version_message: str = ""
) -> None:
"""Save a base model to a new version on the project.
Args:
root_object (Base): The Speckle base object for the new version.
model_id (str): For now please use a `branchName`!
version_message (str): The message for the new version.
"""
if model_id == self.automation_run_data.model_id:
raise ValueError(
f"The target model id: {model_id} cannot match the model id"
f" that triggered this automation: {self.automation_run_data.model_id}"
)
root_object_id = operations.send(
root_object,
[self._server_transport, self._memory_transport],
use_default_cache=False,
)
version_id = self.speckle_client.commit.create(
stream_id=self.automation_run_data.project_id,
object_id=root_object_id,
branch_name=model_id,
message=version_message,
source_application="SpeckleAutomate",
)
self._automation_result.result_versions.append(version_id)
def report_run_status(self) -> None:
"""Report the current run status to the Speckle server triggered the automation.
Once the automation function exits, send the status to the speckle server.
Return the result from the server, it should be a link to the stored automation
result.
"""
query = gql(
"""
mutation ReportFunctionRunStatus(
$automationId: String!,
$automationRevisionId: String!,
$automationRunId: String!,
$versionId: String!,
$functionId: String!,
$runStatus: AutomationRunStatus!
$elapsed: Float!
$resultVersionIds: [String!]!
$statusMessage: String
$objectResults: JSONObject
){
automationMutations {
functionRunStatusReport(input: {
automationId: $automationId
automationRevisionId: $automationRevisionId
automationRunId: $automationRunId
versionId: $versionId
functionRuns: [
{
functionId: $functionId
status: $runStatus,
elapsed: $elapsed,
resultVersionIds: $resultVersionIds,
statusMessage: $statusMessage
results: $objectResults
}]
})
}
}
"""
)
if self.run_status in [AutomationStatus.SUCCEEDED, AutomationStatus.FAILED]:
object_results = {
"version": "1.0.0",
"values": {
"speckleObjects": self._automation_result.model_dump(by_alias=True)[
"objectResults"
],
"blobs": self._automation_result.blobs,
},
}
else:
object_results = None
params = {
"automationId": self.automation_run_data.automation_id,
"automationRevisionId": self.automation_run_data.automation_revision_id,
"automationRunId": self.automation_run_data.automation_run_id,
"versionId": self.automation_run_data.version_id,
"functionId": self.automation_run_data.function_id,
"runStatus": self.run_status.value,
"elapsed": self.elapsed(),
"resultVersionIds": self._automation_result.result_versions,
"objectResults": object_results,
}
self.speckle_client.httpclient.execute(query, params)
def store_file_result(self, file_path: Path | str) -> None:
"""Save a file attached to the project of this automation."""
path_obj = (
Path(file_path).resolve() if isinstance(file_path, str) else file_path
)
class UploadResult(AutomateBase):
blob_id: str
file_name: str
upload_status: int
class BlobUploadResponse(AutomateBase):
upload_results: list[UploadResult]
if not path_obj.exists():
raise ValueError("The given file path doesn't exist")
files = {path_obj.name: open(str(path_obj), "rb")}
url = (
f"{self.automation_run_data.speckle_server_url}/api/stream/"
f"{self.automation_run_data.project_id}/blob"
)
data = (
httpx.post(
url,
files=files,
headers={"authorization": f"Bearer {self._speckle_token}"},
)
.raise_for_status()
.json()
)
upload_response = BlobUploadResponse.model_validate(data)
if len(upload_response.upload_results) != 1:
raise ValueError("Expecting one upload result.")
for upload_result in upload_response.upload_results:
self._automation_result.blobs.append(upload_result.blob_id)
def mark_run_failed(self, status_message: str) -> None:
"""Mark the current run a failure."""
self._mark_run(AutomationStatus.FAILED, status_message)
def mark_run_success(self, status_message: str | None) -> None:
"""Mark the current run a success with an optional message."""
self._mark_run(AutomationStatus.SUCCEEDED, status_message)
def _mark_run(self, status: AutomationStatus, status_message: str | None) -> None:
duration = self.elapsed()
self._automation_result.status_message = status_message
self._automation_result.run_status = status
self._automation_result.elapsed = duration
msg = f"Automation run {status.value} after {duration:2f} seconds."
print("\n".join([msg, status_message]) if status_message else msg)
def add_object_error(self, object_id: str, error_cause: str) -> None:
"""Add an error to a given objec id."""
self._add_object_result(object_id, ObjectResultLevel.ERROR, error_cause)
def add_object_warning(self, object_id: str, warning: str) -> None:
"""Add a warning to a given object id."""
self._add_object_result(object_id, ObjectResultLevel.WARNING, warning)
def add_object_info(self, object_id: str, info: str) -> None:
"""Add an info message to a given object."""
self._add_object_result(object_id, ObjectResultLevel.INFO, info)
def _add_object_result(
self, object_id: str, level: ObjectResultLevel, status_message: str
) -> None:
print(
f"Object {object_id} was marked with {level.value.upper()}",
f" cause: {status_message}",
)
self._automation_result.object_results[object_id].append(
ObjectResult(level=level, status_message=status_message)
)
AutomateFunction = Callable[[AutomationContext, T], None]
AutomateFunctionWithoutInputs = Callable[[AutomationContext], None]
@overload
def execute_automate_function(
automate_function: AutomateFunction[T],
input_schema: type[T],
) -> None:
...
@overload
def execute_automate_function(automate_function: AutomateFunctionWithoutInputs) -> None:
...
def execute_automate_function(
automate_function: AutomateFunction[T] | AutomateFunctionWithoutInputs,
input_schema: type[T] | None = None,
):
"""Runs the provided automate function with the input schema."""
# first arg is the python file name, we do not need that
args = sys.argv[1:]
if len(args) < 2:
raise ValueError("too few arguments specified need minimum 2")
if len(args) > 4:
raise ValueError("too many arguments specified, max supported is 4")
# we rely on a command name convention to decide what to do.
# this is here, so that the function authors do not see any of this
command = args[0]
if command == "generate_schema":
path = Path(args[1])
schema = json.dumps(
input_schema.model_json_schema(by_alias=True) if input_schema else {}
)
path.write_text(schema)
elif command == "run":
automation_run_data = args[1]
function_inputs = args[2]
speckle_token = os.environ.get("SPECKLE_TOKEN", None)
if not speckle_token and len(args) != 4:
raise ValueError("Cannot get speckle token from arguments or environment")
speckle_token = speckle_token if speckle_token else args[3]
inputs = (
input_schema.model_validate_json(function_inputs)
if input_schema
else input_schema
)
if inputs:
automate_sdk = run_function(
automate_function, # type: ignore
automation_run_data,
speckle_token,
inputs,
)
else:
automate_sdk = run_function(
automate_function, # type: ignore
automation_run_data,
speckle_token,
)
exit_code = 0 if automate_sdk.run_status == AutomationStatus.SUCCEEDED else 1
exit(exit_code)
else:
raise NotImplementedError(f"Command: '{command}' is not supported.")
@overload
def run_function(
automate_function: AutomateFunction[T],
automation_run_data: AutomationRunData | str,
speckle_token: str,
inputs: T,
) -> AutomationContext:
...
@overload
def run_function(
automate_function: AutomateFunctionWithoutInputs,
automation_run_data: AutomationRunData | str,
speckle_token: str,
) -> AutomationContext:
...
def run_function(
automate_function: AutomateFunction[T] | AutomateFunctionWithoutInputs,
automation_run_data: AutomationRunData | str,
speckle_token: str,
inputs: T | None = None,
) -> AutomationContext:
"""Run the provided function with the automate sdk context."""
automate_sdk = AutomationContext.initialize(automation_run_data, speckle_token)
automate_sdk.report_run_status()
try:
# avoiding complex type gymnastics here on the internals.
# the external type overloads make this correct
if inputs:
automate_function(automate_sdk, inputs) # type: ignore
else:
automate_function(automate_sdk) # type: ignore
# the function author forgot to mark the function success
if automate_sdk.run_status not in [
AutomationStatus.FAILED,
AutomationStatus.SUCCEEDED,
]:
automate_sdk.mark_run_success(
"WARNING: Automate assumed a success status,"
" but it was not marked as so by the function."
)
except Exception:
trace = traceback.format_exc()
print(trace)
automate_sdk.mark_run_failed(
"Function error. Check the automation run logs for details."
)
finally:
automate_sdk.report_run_status()
return automate_sdk
main.py
+37 -2
View File
@@ -2,11 +2,13 @@
use the automation_context module to wrap your function in an Autamate context helper
"""
from automation_context import (
from speckle_automate import (
AutomateBase,
AutomationContext,
execute_automate_function,
)
from flatten import flatten_base
@@ -14,6 +16,8 @@ class FunctionInputs(AutomateBase):
"""These are function author defined values.
Automate will make sure to supply them matching the types specified here.
Please use the pydantic model schema to define your inputs:
ttps://docs.pydantic.dev/latest/usage/models/
"""
forbidden_speckle_type: str
@@ -23,7 +27,16 @@ def automate_function(
automate_context: AutomationContext,
function_inputs: FunctionInputs,
) -> None:
"""Hey, trying the automate sdk experience here."""
"""This is an example Speckle Automate function.
Args:
automate_context: A context helper object, that carries relevant information
about the runtime context of this function.
It gives access to the Speckle project data, that triggered this run.
It also has conveniece methods attach result data to the Speckle model.
function_inputs: An instance object matching the defined schema.
"""
# the context provides a conveniet way, to receive the triggering version
version_root_object = automate_context.receive_version()
count = 0
@@ -39,6 +52,7 @@ def automate_function(
count += 1
if count > 0:
# this is how a run is marked with a failure cause
automate_context.mark_run_failed(
"Automation failed: "
f"Found {count} object that have a forbidden speckle type: "
@@ -48,6 +62,27 @@ def automate_function(
else:
automate_context.mark_run_success("No forbidden types found.")
# if the function generates file results, this is how it can be
# attached to the Speckle project / model
# automate_context.store_file_result("./report.pdf")
def automate_function_without_inputs(automate_context: AutomationContext) -> None:
"""A function example without inputs.
If your function does not need any input variables,
besides what the automation context provides,
the inputs argument can be omitted.
"""
pass
# make sure to call the function with the executor
if __name__ == "__main__":
# NOTE: always pass in the automate function by its reference, do not invoke it!
# pass in the function reference with the inputs schema to the executor
execute_automate_function(automate_function, FunctionInputs)
# if the function has no arguments, the executor can handle it like so
# execute_automate_function(automate_function_without_inputs)
poetry.lock
Generated
+12 -42
View File
@@ -23,7 +23,6 @@ files = [
]
[package.dependencies]
exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""}
idna = ">=2.8"
sniffio = ">=1.1"
@@ -109,8 +108,6 @@ mypy-extensions = ">=0.4.3"
packaging = ">=22.0"
pathspec = ">=0.9.0"
platformdirs = ">=2"
tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
typing-extensions = {version = ">=4.0.1", markers = "python_version < \"3.11\""}
[package.extras]
colorama = ["colorama (>=0.4.3)"]
@@ -255,20 +252,6 @@ wrapt = ">=1.10,<2"
[package.extras]
dev = ["PyTest", "PyTest-Cov", "bump2version (<1)", "sphinx (<2)", "tox"]
[[package]]
name = "exceptiongroup"
version = "1.1.3"
description = "Backport of PEP 654 (exception groups)"
optional = false
python-versions = ">=3.7"
files = [
{file = "exceptiongroup-1.1.3-py3-none-any.whl", hash = "sha256:343280667a4585d195ca1cf9cef84a4e178c4b6cf2274caef9859782b567d5e3"},
{file = "exceptiongroup-1.1.3.tar.gz", hash = "sha256:097acd85d473d75af5bb98e41b61ff7fe35efe6675e4f9370ec6ec5126d160e9"},
]
[package.extras]
test = ["pytest (>=6)"]
[[package]]
name = "gql"
version = "3.4.1"
@@ -508,7 +491,6 @@ files = [
[package.dependencies]
mypy-extensions = ">=1.0.0"
tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
typing-extensions = ">=4.1.0"
[package.extras]
@@ -729,11 +711,9 @@ files = [
[package.dependencies]
colorama = {version = "*", markers = "sys_platform == \"win32\""}
exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""}
iniconfig = "*"
packaging = "*"
pluggy = ">=0.12,<2.0"
tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""}
[package.extras]
testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
@@ -812,13 +792,13 @@ files = [
[[package]]
name = "specklepy"
version = "2.16.2"
version = "2.17.0"
description = "The Python SDK for Speckle 2.0"
optional = false
python-versions = ">=3.7.2,<4.0"
python-versions = ">=3.8.0,<4.0"
files = [
{file = "specklepy-2.16.2-py3-none-any.whl", hash = "sha256:3337a9512d73cdf0528c9a40f0318aa68ee859326e4197cd908a9ab159184365"},
{file = "specklepy-2.16.2.tar.gz", hash = "sha256:a48dbc17e289cc85e6adbc6415430292f2d681bdda0dad5b299ea635dd5430d4"},
{file = "specklepy-2.17.0-py3-none-any.whl", hash = "sha256:90c18812666dd3beeabe6b784f57de26227f7651d949c7cabe041e533b895581"},
{file = "specklepy-2.17.0.tar.gz", hash = "sha256:6c6044351beb6e68e81d296189966efdc470096febfa588e965949567ffaac1a"},
]
[package.dependencies]
@@ -826,6 +806,7 @@ appdirs = ">=1.4.4,<2.0.0"
attrs = ">=23.1.0,<24.0.0"
Deprecated = ">=1.2.13,<2.0.0"
gql = {version = ">=3.3.0,<4.0.0", extras = ["requests", "websockets"]}
httpx = ">=0.25.0,<0.26.0"
pydantic = ">=2.0,<3.0"
stringcase = ">=1.2.0,<2.0.0"
ujson = ">=5.3.0,<6.0.0"
@@ -840,26 +821,15 @@ files = [
{file = "stringcase-1.2.0.tar.gz", hash = "sha256:48a06980661908efe8d9d34eab2b6c13aefa2163b3ced26972902e3bdfd87008"},
]
[[package]]
name = "tomli"
version = "2.0.1"
description = "A lil' TOML parser"
optional = false
python-versions = ">=3.7"
files = [
{file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"},
{file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"},
]
[[package]]
name = "typing-extensions"
version = "4.7.1"
description = "Backported and Experimental Type Hints for Python 3.7+"
version = "4.8.0"
description = "Backported and Experimental Type Hints for Python 3.8+"
optional = false
python-versions = ">=3.7"
python-versions = ">=3.8"
files = [
{file = "typing_extensions-4.7.1-py3-none-any.whl", hash = "sha256:440d5dd3af93b060174bf433bccd69b0babc3b15b1a8dca43789fd7f61514b36"},
{file = "typing_extensions-4.7.1.tar.gz", hash = "sha256:b75ddc264f0ba5615db7ba217daeb99701ad295353c45f9e95963337ceeeffb2"},
{file = "typing_extensions-4.8.0-py3-none-any.whl", hash = "sha256:8f92fc8806f9a6b641eaa5318da32b44d401efaac0f6678c9bc448ba3605faa0"},
{file = "typing_extensions-4.8.0.tar.gz", hash = "sha256:df8e4339e9cb77357558cbdbceca33c303714cf861d1eef15e1070055ae8b7ef"},
]
[[package]]
@@ -1199,5 +1169,5 @@ multidict = ">=4.0"
[metadata]
lock-version = "2.0"
python-versions = "^3.10"
content-hash = "a4fd24e538d011e9a91f29c80a4221504c7e35e07cecdac214e9710683974c65"
python-versions = "^3.11"
content-hash = "780374d9cb4a657b90c78126b3caba29f2b4dbf202f9953705bfd512f34a88c6"
pyproject.toml
+2 -5
View File
@@ -7,11 +7,8 @@ readme = "README.md"
packages = [{include = "src/speckle_automate_py"}]
[tool.poetry.dependencies]
python = "^3.10"
specklepy = "^2.16.2"
pydantic = "^2.1.1"
stringcase = "^1.2.0"
httpx = "^0.25.0"
python = "^3.11"
specklepy = "^2.17.0"
[tool.poetry.group.dev.dependencies]
black = "^23.3.0"
tests/test_function.py
+6 -19
View File
@@ -6,17 +6,17 @@ from pathlib import Path
import pytest
from gql import gql
from specklepy.api import operations
from specklepy.api.client import SpeckleClient
from specklepy.objects.base import Base
from specklepy.transports.server import ServerTransport
from automation_context import (
from speckle_automate import (
AutomationContext,
AutomationRunData,
AutomationStatus,
run_function,
)
from specklepy.api import operations
from specklepy.api.client import SpeckleClient
from specklepy.objects.base import Base
from specklepy.transports.server import ServerTransport
from main import FunctionInputs, automate_function
@@ -158,16 +158,3 @@ def test_function_run(automation_run_data: AutomationRunData, speckle_token: str
)
assert automate_sdk.run_status == AutomationStatus.FAILED
def test_file_uploads(automation_run_data: AutomationRunData, speckle_token: str):
"""Test file store capabilities of the automate sdk."""
automate_context = AutomationContext.initialize(automation_run_data, speckle_token)
path = Path(f"./{crypto_random_string(10)}").resolve()
path.write_text("foobar")
automate_context.store_file_result(path)
os.remove(path)
assert len(automate_context._automation_result.blobs) == 1