Gergő Jedlicska
61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
import os
|
|
import sys
|
|
|
|
import papermill as pm
|
|
from speckle_automate import AutomationContext, AutomationResult, AutomationStatus
|
|
|
|
if __name__ == "__main__":
|
|
args = sys.argv[1:]
|
|
|
|
# we need to pass in the token via an env var, the params cell is executed
|
|
# automatically exposing the token in the output file
|
|
token_env_var = "SPECKLE_TOKEN"
|
|
|
|
speckle_token = os.environ.get(token_env_var) or args[2]
|
|
|
|
os.environ[token_env_var] = speckle_token
|
|
|
|
automation_run_data, function_inputs = args[0], args[1]
|
|
|
|
automation_context = AutomationContext.initialize(
|
|
automation_run_data, speckle_token
|
|
)
|
|
|
|
output_file = "function.output.ipynb"
|
|
try:
|
|
run_result = pm.execute_notebook(
|
|
"function.ipynb",
|
|
output_file,
|
|
parameters={
|
|
"automation_run_data": automation_run_data,
|
|
"function_inputs": function_inputs,
|
|
"token_env_var": token_env_var,
|
|
"execute_function": True,
|
|
},
|
|
)
|
|
|
|
# find the cell by its tag
|
|
function_execution_cells = [
|
|
c for c in run_result.cells if "function_execution" in c.metadata.tags
|
|
]
|
|
|
|
# should have only found 1
|
|
if len(function_execution_cells) != 1:
|
|
msg = (
|
|
"Expected 1 cell to have the tag `function_inputs`,"
|
|
f" found {len(function_execution_cells)}"
|
|
)
|
|
raise ValueError(msg)
|
|
automation_result = AutomationResult.model_validate_json(
|
|
function_execution_cells[0].outputs[-1]["text"]
|
|
)
|
|
automation_context._automation_result = automation_result
|
|
automation_context.store_file_result(output_file)
|
|
except Exception as ex:
|
|
automation_context.mark_run_failed(str(ex))
|
|
finally:
|
|
exit_code = (
|
|
0 if automation_context.run_status == AutomationStatus.SUCCEEDED else 1
|
|
)
|
|
exit(exit_code)
|