import datetime
import logging
import os
import re
import shlex
import sys
from contextlib import contextmanager
from pathlib import Path
from typing import Any
from typing import Dict
from typing import Generator
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union
from ewokscore.events.contexts import RawExecInfoType
from ewokscore.events.contexts import job_context
from ewokscore.graph import TaskGraph
from ewokscore.graph.inputs import graph_inputs_as_table
from tabulate import tabulate
from . import _engines
from . import _requirements
from . import graph_cache
from .errors import AbortException
try:
from ewoksjob.bindings import submit_graph as _submit_graph
from ewoksjob.client.futures import FutureInterface
except ImportError:
_submit_graph = None
FutureInterface = Any
try:
from pyicat_plus.client import defaults as icat_defaults
from pyicat_plus.client.main import IcatClient
except ImportError:
IcatClient = None
icat_defaults = None
__all__ = ["execute_graph", "load_graph", "save_graph", "convert_graph", "submit_graph"]
logger = logging.getLogger(__name__)
[docs]
def execute_graph(
graph,
engine: Optional[str] = None,
inputs: Optional[List[dict]] = None,
load_options: Optional[dict] = None,
varinfo: Optional[dict] = None,
execinfo: RawExecInfoType = None,
task_options: Optional[dict] = None,
outputs: Optional[List[dict]] = None,
merge_outputs: Optional[bool] = True,
environment: Optional[dict] = None,
convert_destination: Optional[Any] = None,
save_options: Optional[dict] = None,
upload_parameters: Optional[dict] = None,
**execute_options,
) -> Optional[dict]:
with job_context(execinfo, engine=engine) as execinfo:
with _upload_context(upload_parameters):
if environment:
environment = {k: str(v) for k, v in environment.items()}
os.environ.update(environment)
# Load the graph
if load_options is None:
load_options = dict()
graph = load_graph(graph, inputs=inputs, **load_options)
# Save the graph (with inputs)
if convert_destination is not None:
convert_graph(graph, convert_destination, save_options=save_options)
# Execute the graph
engine_api = _engines.get_execution_engine(engine)
result = engine_api.execute_graph(
graph,
varinfo=varinfo,
execinfo=execinfo,
task_options=task_options,
outputs=outputs,
merge_outputs=merge_outputs,
**execute_options,
)
return result
@contextmanager
def _upload_context(upload_parameters: Optional[dict]) -> Generator[None, None, None]:
if upload_parameters is None:
yield
return
if IcatClient is None:
raise RuntimeError("requires the 'pyicat-plus' package")
metadata = upload_parameters.setdefault("metadata", {})
if "startDate" not in metadata:
metadata["startDate"] = datetime.datetime.now().astimezone()
yield
# Only upload when no exception is raised.
if "endDate" not in metadata:
metadata["endDate"] = datetime.datetime.now().astimezone()
metadata_urls = upload_parameters.pop(
"metadata_urls", icat_defaults.METADATA_BROKERS
)
client = IcatClient(metadata_urls=metadata_urls)
logger.info(
"Upload processed dataset '%s' metadata to ICAT: %s",
upload_parameters.get("dataset"),
upload_parameters.get("path"),
)
client.store_processed_data(**upload_parameters)
[docs]
def submit_graph(
graph,
_celery_options: Optional[dict] = None,
resolve_graph_remotely: Optional[bool] = None,
load_options: Optional[dict] = None,
**options,
) -> FutureInterface:
"""Submit a workflow to be executed remotely. The workflow is
resolved on the client-side by default (e.g. load from a file)
but can optionally be resolved remotely.
"""
if _submit_graph is None:
raise RuntimeError("requires the 'ewoksjob' package")
return _submit_graph(
graph,
_convert_graph=_load_graph,
_celery_options=_celery_options,
resolve_graph_remotely=resolve_graph_remotely,
load_options=load_options,
**options,
)
def _load_graph(graph, load_options: Optional[dict] = None) -> dict:
return convert_graph(
graph, None, load_options=load_options, save_requirements=False
)
@graph_cache.cache
def load_graph(
graph: Any,
inputs: Optional[List[dict]] = None,
representation: Optional[str] = None,
root_dir: Optional[Union[str, Path]] = None,
root_module: Optional[str] = None,
**load_options,
) -> TaskGraph:
"""When load option `graph_cache_max_size > 0` is provided, the graph will cached in memory.
When the graph comes from external storage (for example a file) any changes
to the external graph will require flushing the cache with `graph_cache_max_size = 0`.
"""
engine_api, representation = _engines.get_serialization_engine(
graph, representation=representation
)
return engine_api.deserialize_graph(
graph,
inputs=inputs,
representation=representation,
root_dir=root_dir,
root_module=root_module,
**load_options,
)
[docs]
def save_graph(
graph: TaskGraph,
destination,
representation: Optional[str] = None,
**save_options,
) -> Union[str, dict]:
engine_api, representation = _engines.get_serialization_engine(
destination, representation=representation
)
return engine_api.serialize_graph(
graph, destination, representation=representation, **save_options
)
[docs]
def convert_graph(
source,
destination,
inputs: Optional[List[dict]] = None,
load_options: Optional[dict] = None,
save_options: Optional[dict] = None,
save_requirements: bool = True,
package_manager_name: Optional[str] = None,
package_manager_command: Union[None, str, Tuple[str, ...]] = None,
) -> Union[str, dict]:
if load_options is None:
load_options = dict()
if save_options is None:
save_options = dict()
graph = load_graph(source, inputs=inputs, **load_options)
if save_requirements:
if not package_manager_command:
package_manager_command = tuple()
elif isinstance(package_manager_command, str):
package_manager_command = _split_command(package_manager_command)
try:
_requirements.add_requirements(
graph,
manager_name=package_manager_name,
manager_command=package_manager_command,
)
except Exception:
logger.exception("Continue after failure to add workflow requirements")
return save_graph(graph, destination, **save_options)
[docs]
def show_graph(
source,
inputs: Optional[List[dict]] = None,
load_options: Optional[dict] = None,
original_source: Optional[str] = None,
column_widths: Optional[Dict[str, Optional[int]]] = None,
) -> Union[str, dict]:
if load_options is None:
load_options = dict()
graph = load_graph(source, inputs=inputs, **load_options)
_print_graph(graph, column_widths=column_widths, original_source=original_source)
def _print_graph(
graph: TaskGraph,
column_widths: Optional[Dict[str, Optional[int]]] = None,
original_source: Optional[str] = None,
) -> None:
column_names, rows, metadata, footnotes = graph_inputs_as_table(
graph, column_widths=column_widths
)
print()
if original_source:
print(f"Workflow: {original_source}")
else:
print("Workflow:")
for key, value in metadata.items():
label = key.replace("_", " ").capitalize()
print(f"{label}: {value}")
if rows:
print(tabulate(rows, headers=column_names, tablefmt="fancy_grid"))
else:
print("No workflow input parameters detected!")
for footnote in footnotes:
print(footnote)
[docs]
def install_graph(
source,
skip_prompt: bool = False,
package_manager_name: Optional[str] = None,
package_manager_command: Union[None, str, Tuple[str, ...]] = None,
load_options: Optional[dict] = None,
) -> None:
if load_options is None:
load_options = dict()
graph = load_graph(source, **load_options)
requirements = _requirements.get_requirements(graph)
if not package_manager_command:
package_manager_command = tuple()
elif isinstance(package_manager_command, str):
package_manager_command = _split_command(package_manager_command)
if skip_prompt:
_requirements.install_requirements(
requirements,
manager_name=package_manager_name,
manager_command=package_manager_command,
)
return
answer = input(
f"{requirements.__info__()}\n\nThis will install the packages above. Do you want to proceed (y/N)?"
)
if answer.lower() == "y" or answer.lower() == "yes":
_requirements.install_requirements(
requirements,
manager_name=package_manager_name,
manager_command=package_manager_command,
)
else:
raise AbortException()
def _split_command(command: str) -> Tuple[str, ...]:
if sys.platform == "win32":
return tuple(_shlex_split(command))
return tuple(shlex.split(command))
def _shlex_split(command: str) -> Tuple[str, ...]:
r"""
Split a string on spaces, but treat escaped spaces (\ ) as part of the token.
"""
parts = re.split(r"(?<!\\) ", command)
parts = [p.replace(r"\ ", " ") for p in parts]
return tuple(p for p in parts if p)
