pinexq-procon 2.1.0.dev3__py3-none-any.whl

pinexq/procon/dataslots/metadata.py

@@ -0,0 +1,137 @@
+import json
+import logging
+from pathlib import Path
+from typing import Callable, Protocol, Sequence, TypeVar
+
+from ..dataslots.datatypes import Metadata, SlotDescription
+
+
+LOG = logging.getLogger(__name__)
+
+
+def metadata_to_json(meta: Metadata) -> str:
+    return meta.model_dump_json()
+
+
+def json_to_metadata(data: str) -> Metadata:
+    return Metadata(**json.loads(data))
+
+
+class MetadataProxy:
+    """
+    Getter/setter methods for controlled user access to a Metadata object.
+
+    If no Metadata object is provided for initialization, this class will assign a default
+    on first access of any metadata property.
+    """
+    __slots__ = ('_metadata', '_readonly')
+
+    _metadata: Metadata
+    _readonly: bool
+
+    def __init__(self, metadata: Metadata, _readonly: bool = False):
+        self._metadata = metadata
+        self._readonly = _readonly
+
+    def __bool__(self):
+        """Return True if Metadata was set."""
+        return len(self._metadata.comment) > 0 or len(self._metadata.tags) > 0
+
+    @property
+    def comment(self) -> str:
+        """Comments assigned to the Workdata of this slot."""
+        return self._metadata.comment
+
+    @comment.setter
+    def comment(self, s: str):
+        if self._readonly:
+            raise AttributeError("Metadata on this slot is read-only! Can not set 'comment'.")
+        self._metadata.comment = s
+
+    @property
+    def tags(self) -> tuple[str, ...] | list[str]:
+        """Tags assigned to the Workdata of this slot.
+
+        Returns a list of tags if the metadata is writable (e.g. OUTPUT-slot)
+        or an immutable tuple if metadata is read-only (e.g. INPUT slots)
+        """
+        if self._readonly:
+            return tuple(self._metadata.tags)
+        return self._metadata.tags
+
+    @tags.setter
+    def tags(self, tag_list: Sequence[str]):
+        if self._readonly:
+            raise AttributeError("Metadata on this slot is read-only! Can not set 'tags'.")
+        self._metadata.tags = list(tag_list)
+
+    @property
+    def filename(self) -> str:
+        """Original filename that was uploaded as Workdata. This is independent of the
+        filename accessed by the backend of the Slot, which can vary depending on the implementation."""
+        return self._metadata.filename
+
+    @filename.setter
+    def filename(self, s: str):
+        if self._readonly:
+            raise AttributeError("Metadata on this slot is read-only! Can not set 'filename'.")
+        self._metadata.filename = s
+
+    def __repr__(self):
+        return self._metadata.__repr__()
+
+    def is_readonly(self) -> bool:
+        """True if the metadata is readonly."""
+        return self._readonly
+
+
+class MetadataHandler(Protocol):
+    """Defines the expected interface how to read/write metadata."""
+
+    def get(self, slot: SlotDescription) -> Metadata:
+        ...
+
+    def set(self, slot: SlotDescription, metadata: Metadata):
+        ...
+
+
+class LocalFileMetadataStore(MetadataHandler):
+    """Store metadata as a '.meta' sidecar file next to file defined in the slot's description uri."""
+
+    @staticmethod
+    def _metadata_path_from_description(slot: SlotDescription) -> Path:
+        slot_path = Path(slot.uri)
+        return slot_path.with_name(f"{slot_path.name}.meta")
+
+    def get(self, slot: SlotDescription) -> Metadata | None:
+        meta_path = self._metadata_path_from_description(slot)
+        if meta_path.exists():
+            return Metadata.model_validate_json(meta_path.read_text())
+        else:
+            return None
+
+    def set(self, slot: SlotDescription, metadata: Metadata):
+        meta_path = self._metadata_path_from_description(slot)
+        meta_json = metadata.model_dump_json()
+        meta_path.write_text(meta_json)
+
+
+TSetCb = TypeVar('TSetCb', bound=Callable[[SlotDescription, Metadata], None])
+TGetCb = TypeVar('TGetCb', bound=Callable[[SlotDescription,], Metadata])
+
+
+class CallbackMetadataHandler(MetadataHandler):
+    """Generic call/callback-interface to get and set metadata"""
+
+    _setter_callback: TSetCb
+    _getter_callback: TGetCb
+
+    def __init__(self, setter: TSetCb, getter: TGetCb):
+        self._setter_callback = setter
+        self._getter_callback = getter
+
+    def get(self, slot: SlotDescription) -> Metadata:
+        return self._getter_callback(slot)
+
+    def set(self, slot: SlotDescription, metadata: Metadata):
+        self._setter_callback(slot, metadata)

pinexq/procon/jobmanagement/__init__.py

@@ -0,0 +1,9 @@
+from .api_helpers import (
+    client_from_env_vars,
+    client_from_job_execution_context,
+    job_from_step_context,
+    get_client,
+    get_entrypoint_hco,
+    get_job,
+    get_grants,
+)

pinexq/procon/jobmanagement/api_helpers.py

@@ -0,0 +1,287 @@
+import logging
+import os
+import warnings
+from enum import StrEnum
+from typing import Any
+
+import httpx
+import jwt
+from httpx_caching import CachingClient
+from pinexq_client.job_management import Job, enter_jma
+from pinexq_client.job_management.hcos import EntryPointHco
+
+from ..core.exceptions import ProConException
+from ..runtime.job import RemoteExecutionContext
+from ..step import Step
+from ..step.step import ExecutionContextType
+
+
+logger = logging.getLogger(__name__)
+
+
+class JmaApiHeaders(StrEnum):
+    api_key = "x-api-key"
+    access_token = "x-access-token"
+
+
+PREFIX = "JMC_"
+API_HOST_URL = f"{PREFIX}API_HOST_URL"
+API_KEY = f"{PREFIX}API_KEY"
+ACCESS_TOKEN = f"{PREFIX}ACCESS_TOKEN"
+
+# use a client with local cache
+USE_CLIENT_WITH_CACHE = True
+
+# Environment variables that will directly translate to HTTP headers
+#  'header-name': 'ENVIRONMENT_VARIABLE'
+HEADERS_FROM_ENV = {
+    JmaApiHeaders.api_key: API_KEY,
+    JmaApiHeaders.access_token: ACCESS_TOKEN,
+}
+
+def _create_client_instance(api_endpoint: str, headers: dict) -> httpx.Client:
+    """
+    Will create a httpx client, optional with caching, to be used with the API objects.
+
+    Args:
+        api_endpoint: The endpoint for the pinexq API.
+        headers: headers to be passed to httpx.Client.
+
+    """
+    client =  httpx.Client(base_url=api_endpoint, headers=headers)
+    if USE_CLIENT_WITH_CACHE:
+        # for now, we use the persistent cache, which is also shared between instances
+        # use  if you need each client to have a own cache storage=InMemoryStorage()
+        # broken, will cache SSE stream
+        #return SyncCacheClient(
+        #    base_url=pinexq_api_endpoint,
+        #    headers=headers,
+        #    timeout=timeout)
+        return CachingClient(client)
+    else:
+        return client
+
+def client_from_env_vars() -> httpx.Client:
+    """Initializes a httpx.Client from environment variables.
+
+    The following variables are used:
+        JMC_API_KEY: (header, optional) The api-key you get from the login portal.
+        JMC_ACCESS_TOKEN: (header, optional) The JWT access token for the JM.
+    """
+    try:
+        api_url = os.environ[API_HOST_URL]
+    except KeyError:
+        raise ProConException(
+            f'Environment variable "{API_HOST_URL}" not found!'
+            f" Can not configure job management api access!"
+        )
+    headers = {hdr: os.environ[env] for hdr, env in HEADERS_FROM_ENV.items() if env in os.environ}
+    return _create_client_instance(api_endpoint=api_url, headers=headers)
+
+
+def client_from_job_execution_context(exec_context: ExecutionContextType) -> httpx.Client:
+    # This function will be deprecated and replaced by the function below!
+    """Initializes a httpx.Client from connection info embedded in the job.offer.
+
+    Tries to get the JMApi-host and headers from the environment variables first.
+    If not available, the host is extracted from the job-url in the job.offer. Headers
+    are set by the information provided job.offer.
+    """
+
+    warnings.warn(
+        "Calling `client_from_job_execution_context()` directly will be deprecated "
+        "in future releases. Please use the equivalent `get_client()` function instead.",
+        DeprecationWarning
+    )
+
+    if not exec_context.current_job_offer:
+        raise ProConException(
+            "Current execution context provides no information to create client (no job offer)."
+        )
+
+    job_context = exec_context.current_job_offer.job_execution_context
+
+    try:
+        # fixme: Redundant, as there is an explicit get_from_env_vars function
+        api_url = os.environ[API_HOST_URL]
+    except KeyError:
+        # Return only the host-url, cutoff path and query
+        api_url = httpx.URL(job_context.job_url).join("/")
+
+    # Create header from environment variables
+    # fixme: Redundant, as there is an explicit get_from_env_vars function
+    env_var_headers = {
+        hdr: os.environ[env] for hdr, env in HEADERS_FROM_ENV.items() if env in os.environ
+    }
+    # Create headers from the settings in the job.offer
+    if job_context.access_token is not None:
+        job_offer_headers = {
+            JmaApiHeaders.access_token.value: job_context.access_token,
+        }
+    else:
+        job_offer_headers = {}
+
+    # Update headers from env-vars with info from the job.offer
+    headers = env_var_headers | job_offer_headers
+    return _create_client_instance(api_endpoint=api_url, headers=headers)
+
+
+def _client_from_job_execution_context(context: ExecutionContextType) -> httpx.Client:
+    # This function will replace the deprecated one above!
+    """Initializes a httpx.Client from connection info embedded in the job.offer.
+
+    The host is extracted from the job-url in the job.offer. Headers are set by the
+    information provided job.offer.
+    """
+    if not context.current_job_offer:
+        raise ProConException(
+            "Current execution context provides no information to create client (no job offer)."
+        )
+
+    job_context = context.current_job_offer.job_execution_context
+
+    # Return only the host-url, cutoff path and query
+    api_url = httpx.URL(job_context.job_url).join("/")
+
+    # Create headers from the settings in the job.offer
+    job_offer_headers = {}
+    if job_context.access_token:
+        job_offer_headers[JmaApiHeaders.access_token.value] = job_context.access_token
+
+    # Update headers from env-vars with info from the job.offer
+    return _create_client_instance(api_endpoint=str(api_url), headers=job_offer_headers)
+
+
+def get_client(context: ExecutionContextType | None = None) -> httpx.Client:
+    """Trys to get the client fom job offer context if exec_context is provided.
+    If not possible it trys to get it from environment variables (intended for testing).
+
+    Args:
+        context: The execution context of the current Step container.
+
+    Returns:
+        An initialized HTTPX client with base_url and headers set.
+
+    Raises:
+        ...
+    """
+
+    if context:
+        try:
+            return _client_from_job_execution_context(context)
+        except ProConException as ex:
+            logger.warning(
+                f"Unable to initialize client from execution context: {str(ex)}. "
+                f"Falling back to configure client from environment variables."
+            )
+    try:
+        return client_from_env_vars()
+    except ProConException as ex:
+        raise ProConException(
+            "Unable to determine API host and credentials from neither execution "
+            "context (if running as 'remote') or environment variables!"
+        ) from ex
+
+
+def _job_from_context(context: ExecutionContextType, client: httpx.Client) -> Job:
+    """Create a PinexQ client `Job` object from a given step-function's
+    execution context and a httpx client.
+
+    Args:
+        context: The execution context of the current Step container.
+
+    Returns:
+        An initialized HTTPX client with base_url and headers set.
+    """
+    job_url = httpx.URL(context.current_job_offer.job_execution_context.job_url)
+    return Job.from_url(client=client, job_url=job_url)
+
+
+def job_from_step_context(step: Step, client: httpx.Client | None = None) -> Job:
+    """Initialize an API Job object with the job_id of the current Job.
+
+    Meant to be called inside a Step-container during execution of a Step-function.
+
+    Args:
+        step: The current Step container (referenced by `self`)
+        client: A httpx.Client initialized with the API-url.
+    """
+    # Todo: Mark this function deprecated and refer to `get_job()` instead?
+    context: RemoteExecutionContext | None = step.step_context
+    if not client:
+        client = get_client(context)
+
+    return _job_from_context(context, client)
+
+
+def get_job(context: ExecutionContextType) -> Job:
+    """Create a PinexQ client `Job` object from a step-function's execution context.
+
+    Args:
+        context: The execution context of the current Step container.
+    Returns:
+        A pinexq_client.job_management.tool.Job object initialized as the Job executing the given context.
+    """
+    client = get_client(context)
+    return _job_from_context(context, client)
+
+
+def get_entrypoint_hco(context: ExecutionContextType) -> EntryPointHco:
+    """Create a PinexQ client `Entrypoint` object from a step-function's execution context.
+
+    Args:
+        context: The execution context of the current Step container.
+    Returns:
+        An entrypoint hco object to the JMA where the current Job is running.
+    """
+    client = get_client(context)
+    return enter_jma(client)
+
+
+def get_grants(context: ExecutionContextType) -> list[str]:
+    """Try to access grants for the current user
+
+    Args:
+        context: Execution context of the current step function
+    Returns:
+        A list of grants that could be determined or an empty list.
+    """
+    access_token = _get_access_token(context)
+    if (access_token is None) or (access_token.get("grants") is None):
+        return []
+
+    try:
+        grants = access_token["grants"]
+        return grants
+    except Exception as ex:
+        raise ProConException("Can not extract Grants from access token.") from ex
+
+
+def _get_access_token(context: ExecutionContextType) -> Any | None:
+    """Trys to get access token. If one is contained in the job offer it is taken.
+    else it trys to get an AccessToken from the environment variables.
+
+    Args:
+        context: Execution context of the current step function
+    Returns:
+        A access token raw string or None.
+    """
+
+    access_token: str | None = None
+    try:
+        # This will fail, if we're not in a RemoteExecutionContext, there's no job offer or no token
+        access_token = context.current_job_offer.job_execution_context.access_token
+    except AttributeError:
+        pass
+
+    # Fall back to environment variables, if the previous step failed
+    access_token = access_token or os.environ.get(HEADERS_FROM_ENV[JmaApiHeaders.access_token])
+
+    if access_token is None:
+        return None
+
+    try:
+        token_content = jwt.decode(access_token, options={"verify_signature": False})
+        return token_content
+    except jwt.InvalidTokenError as ex:
+        raise ProConException("Access token could not be decoded!") from ex

pinexq/procon/remote/messages.py

@@ -0,0 +1,250 @@
+"""
+Message type definitions for the RabbitMQ communication.
+
+https://dev.azure.com/data-cybernetics/Rymax-One/_wiki/wikis/General/105/Processing-communication-design
+"""
+from datetime import datetime
+from enum import Enum
+from typing import Annotated, Any, Literal, Optional, Union
+from uuid import UUID, uuid4
+
+from pydantic import AnyUrl, BaseModel, Field, constr
+
+from ..dataslots import Metadata
+
+
+class MessageModel(BaseModel):
+    pass
+
+
+class HeartbeatContent(MessageModel):
+    type_: Literal['heartbeat'] = Field('heartbeat', alias='$type')
+
+
+# class JobQuota(MessageModel):
+#     type_: Literal['quota'] = Field('quota', const=True, alias='$type')
+#     runtime: Optional[timedelta]
+#     cpu_time: Optional[float]
+#     qpu_time: Optional[float]
+#     max_storage: Optional[int]
+
+class Header(BaseModel):
+    Key: str
+    Value: str
+
+
+class SlotInfo(BaseModel):
+    uri: AnyUrl
+    prebuildheaders: list[Header] | None
+    mediatype: str
+    metadata: Metadata | None = None
+
+
+class DataslotInput(BaseModel):
+    name: str
+    sources: list[SlotInfo]
+    # metadata: str
+
+
+class DataslotOutput(BaseModel):
+    name: str
+    destinations: list[SlotInfo]
+    # metadata: str
+
+
+class JobExecutionContext(BaseModel):
+    job_url: str
+    access_token: str | None
+
+
+class JobOfferContent(MessageModel):
+    type_: Literal['job.offer'] = Field('job.offer', alias='$type')
+    job_id: UUID
+    algorithm: str
+    algorithm_version: Optional[str] = None
+    parameters: Optional[Any] = None
+    # quota: Optional[JobQuota]
+    priority: Optional[int] = None
+    input_dataslots: list[DataslotInput]
+    output_dataslots: list[DataslotOutput]
+    job_execution_context: JobExecutionContext
+
+
+class JobStatus(Enum):
+    starting = 'starting'
+    running = 'running'
+    finished = 'finished'
+
+
+class JobStatusContent(MessageModel):
+    type_: Literal['job.status'] = Field('job.status', alias='$type')
+    job_id: UUID
+    status: JobStatus
+    content: Optional[Any] = None
+
+
+class JobProgressContent(MessageModel):
+    type_: Literal['job.progress'] = Field('job.progress', alias='$type')
+    job_id: UUID
+    status: JobStatus
+    content: Optional[Any] = None
+
+
+class JobCommand(Enum):
+    status = 'status'
+    abort = 'abort'
+
+
+class JobCommandContent(MessageModel):
+    type_: Literal['job.command'] = Field('job.command', alias='$type')
+    job_id: UUID
+    command: JobCommand
+
+
+class JobResultContent(MessageModel):
+    type_: Literal['job.result'] = Field('job.result', alias='$type')
+    job_id: UUID
+    result: Any
+    # quota: Optional[JobQuota]
+
+
+class MetadataUpdate(BaseModel):
+    dataslot_name: str
+    slot_index: Annotated[int, Field(ge=0)]
+    metadata: Metadata
+
+
+class JobResultMetadataContent(MessageModel):
+    type_: Literal['job.dataslot.metadata'] = Field('job.dataslot.metadata', alias='$type')
+    updates: list[MetadataUpdate]
+
+
+class WorkerStatus(Enum):
+    starting = 'starting'
+    idle = 'idle'
+    running = 'running'
+    stopped = 'stopped'
+    exiting = 'exiting'
+
+
+class WorkerStatusContent(MessageModel):
+    type_: Literal['worker.status'] = Field('worker.status', alias='$type')
+    worker_id: str
+    status: WorkerStatus
+    content: Optional[Any] = None
+
+
+class WorkerCommand(Enum):
+    start = 'start'
+    stop = 'stop'
+    status = 'status'
+    abort = 'abort'
+    restart = 'restart'
+    shutdown = 'shutdown'
+
+
+class WorkerCommandContent(MessageModel):
+    type_: Literal['worker.command'] = Field('worker.command', alias='$type')
+    worker_id: str
+    command: WorkerCommand
+
+
+class ErrorContent(MessageModel):
+    type_: Literal['error'] = Field('error', alias='$type')
+    type: str = Field('/procon/error')
+    title: str
+    detail: str
+    instance: str
+    metadata: Optional[dict] = None
+
+    @classmethod
+    def from_exception(cls, *args, exception: Exception, **kwargs) -> 'ErrorContent':
+        """Generates the error message from an exception"""
+        return cls(
+            *args,
+            detail=f'{type(exception).__name__}: {exception}',
+            **kwargs
+        )
+
+
+# ID string of a service (worker, job, api) in the header;
+#  format: <service>:<instance id>:<sub id>
+# SenderID = pydantic.constr(regex=r'^[\w-]+:[\w-]+(:[\w-]+)?$')
+SenderID = str  # don't validate the id format for now
+
+
+class MessageHeader(MessageModel):
+    type_: Literal['header'] = Field('header', alias='$type')
+    sender_id: SenderID
+    msg_id: UUID = Field(default_factory=uuid4)
+    date: datetime = Field(default_factory=datetime.utcnow)
+    version: Literal['1.0'] = '1.0'
+
+
+class MessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: Union[
+        JobOfferContent, JobStatusContent, JobProgressContent,
+        JobCommandContent, JobResultContent,
+        JobResultMetadataContent,
+        WorkerStatusContent, WorkerCommandContent,
+        HeartbeatContent, ErrorContent
+    ] = Field(..., discriminator='type_')
+    metadata: Optional[Any] = None
+
+
+class JobOfferMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: Union[JobOfferContent, ErrorContent] = Field(..., discriminator='type_')
+    metadata: Optional[Any] = None
+
+
+class WorkerCommandMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: Union[WorkerCommandContent, ErrorContent] = Field(..., discriminator='type_')
+    metadata: Optional[Any] = None
+
+
+class WorkerStatusMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: Union[WorkerStatusContent, ErrorContent] = Field(..., discriminator='type_')
+    metadata: Optional[Any] = None
+
+
+class JobCommandMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: Union[JobCommandContent, ErrorContent] = Field(..., discriminator='type_')
+    metadata: Optional[Any] = None
+
+
+class JobStatusMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: Union[JobStatusContent, ErrorContent] = Field(..., discriminator='type_')
+    metadata: Optional[Any] = None
+
+
+class JobResultMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: JobResultContent
+    metadata: Optional[Any] = None
+
+
+class JobResultMetadataMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: JobResultMetadataContent
+    metadata: Optional[Any] = None
+
+
+class JobProgressMessageBody(MessageModel):
+    type_: Literal['message'] = Field('message', alias='$type')
+    header: MessageHeader
+    content: JobProgressContent
+    metadata: Optional[Any] = None