digitalhub 0.8.0b0__py3-none-any.whl → 0.8.0b2__py3-none-any.whl

digitalhub/entities/run/entity.py

@@ -0,0 +1,461 @@
+from __future__ import annotations
+
+import time
+import typing
+from typing import Any
+
+import requests
+
+from digitalhub.entities._base.crud import (
+    list_entity_api_base,
+    list_entity_api_ctx,
+    logs_api,
+    read_entity_api_ctx,
+    resume_api,
+    stop_api,
+)
+from digitalhub.entities._base.entity.unversioned import UnversionedEntity
+from digitalhub.entities._base.state import State
+from digitalhub.entities._builders.spec import build_spec
+from digitalhub.entities._builders.status import build_status
+from digitalhub.entities.entity_types import EntityTypes
+from digitalhub.registry.registry import registry
+from digitalhub.runtimes.builder import build_runtime
+from digitalhub.utils.exceptions import EntityError
+from digitalhub.utils.logger import LOGGER
+
+if typing.TYPE_CHECKING:
+    from digitalhub.entities._base.entity.material import MaterialEntity
+    from digitalhub.entities._base.metadata import Metadata
+    from digitalhub.entities.run.spec import RunSpec
+    from digitalhub.entities.run.status import RunStatus
+    from digitalhub.runtimes.base import Runtime
+
+
+class Run(UnversionedEntity):
+    """
+    A class representing a run.
+    """
+
+    ENTITY_TYPE = EntityTypes.RUN.value
+
+    def __init__(
+        self,
+        project: str,
+        uuid: str,
+        kind: str,
+        metadata: Metadata,
+        spec: RunSpec,
+        status: RunStatus,
+        user: str | None = None,
+    ) -> None:
+        super().__init__(project, uuid, kind, metadata, spec, status, user)
+
+        self.spec: RunSpec
+        self.status: RunStatus
+
+    ##############################
+    #  Run Methods
+    ##############################
+
+    def build(self) -> None:
+        """
+        Build run.
+
+        Returns
+        -------
+        None
+        """
+        runtime = self._get_runtime()
+        executable = self._get_executable(runtime)
+        task = self._get_task(runtime)
+        new_spec = runtime.build(executable, task, self.to_dict())
+        self.spec = build_spec(
+            self.kind,
+            **new_spec,
+        )
+        self._set_state(State.BUILT.value)
+        self.save()
+
+    def run(self) -> Run:
+        """
+        Run run.
+
+        Returns
+        -------
+        Run
+            Run object.
+        """
+        self.refresh()
+        if self.spec.local_execution:
+            if not self._is_ready_to_run():
+                raise EntityError("Run is not in a state to run.")
+            self._set_state(State.RUNNING.value)
+            self.save(update=True)
+
+        # Try to get inputs if they exist
+        try:
+            self.spec.inputs = self.inputs(as_dict=True)
+        except EntityError:
+            pass
+
+        try:
+            status = self._get_runtime().run(self.to_dict())
+        except Exception as e:
+            self.refresh()
+            if self.spec.local_execution:
+                self._set_state(State.ERROR.value)
+            self._set_message(str(e))
+            self.save(update=True)
+            raise e
+
+        self.refresh()
+        if not self.spec.local_execution:
+            status.pop("state", None)
+        new_status = {**self.status.to_dict(), **status}
+        self._set_status(new_status)
+        self.save(update=True)
+        return self
+
+    def wait(self, log_info: bool = True) -> Run:
+        """
+        Wait for run to finish.
+
+        Parameters
+        ----------
+        log_info : bool
+            If True, log information.
+
+        Returns
+        -------
+        Run
+            Run object.
+        """
+        start = time.time()
+        while True:
+            if log_info:
+                LOGGER.info(f"Waiting for run {self.id} to finish...")
+            self.refresh()
+            time.sleep(5)
+            if self.status.state in [
+                State.STOPPED.value,
+                State.ERROR.value,
+                State.COMPLETED.value,
+            ]:
+                if log_info:
+                    current = time.time() - start
+                    LOGGER.info(f"Run {self.id} finished in {current:.2f} seconds.")
+                return self
+
+    def inputs(self, as_dict: bool = False) -> list[dict]:
+        """
+        Get inputs passed in spec as objects or as dictionaries.
+
+        Parameters
+        ----------
+        as_dict : bool
+            If True, return inputs as dictionaries.
+
+        Returns
+        -------
+        list[dict]
+            List of input objects.
+        """
+        try:
+            return self.spec.get_inputs(as_dict=as_dict)
+        except AttributeError:
+            msg = f"Run of type {self.kind} has no inputs."
+            raise EntityError(msg)
+
+    def results(self) -> dict:
+        """
+        Get results from runtime execution.
+
+        Returns
+        -------
+        dict
+            Results.
+        """
+        try:
+            return self.status.get_results()
+        except AttributeError:
+            msg = f"Run of type {self.kind} has no results."
+            raise EntityError(msg)
+
+    def result(self, key: str) -> Any:
+        """
+        Get result from runtime execution by key.
+
+        Parameters
+        ----------
+        key : str
+            Key of the result.
+
+        Returns
+        -------
+        Any
+            Result.
+        """
+        return self.results().get(key)
+
+    def outputs(self, as_key: bool = False, as_dict: bool = False) -> dict:
+        """
+        Get run objects results.
+
+        Parameters
+        ----------
+        as_key : bool
+            If True, return results as keys.
+        as_dict : bool
+            If True, return results as dictionaries.
+
+        Returns
+        -------
+        dict
+            List of output objects.
+        """
+        try:
+            return self.status.get_outputs(as_key=as_key, as_dict=as_dict)
+        except AttributeError:
+            msg = f"Run of type {self.kind} has no outputs."
+            raise EntityError(msg)
+
+    def output(self, key: str, as_key: bool = False, as_dict: bool = False) -> MaterialEntity | dict | str | None:
+        """
+        Get run object result by key.
+
+        Parameters
+        ----------
+        key : str
+            Key of the result.
+        as_key : bool
+            If True, return result as key.
+        as_dict : bool
+            If True, return result as dictionary.
+
+        Returns
+        -------
+        Entity | dict | str | None
+            Result.
+        """
+        return self.outputs(as_key=as_key, as_dict=as_dict).get(key)
+
+    def values(self) -> dict:
+        """
+        Get values from runtime execution.
+
+        Returns
+        -------
+        dict
+            Values from backend.
+        """
+        try:
+            value_list = getattr(self.spec, "values", [])
+            value_list = value_list if value_list is not None else []
+            return self.status.get_values(value_list)
+        except AttributeError:
+            msg = f"Run of type {self.kind} has no values."
+            raise EntityError(msg)
+
+    def logs(self) -> dict:
+        """
+        Get object from backend.
+        Returns empty dictionary if context is local.
+
+        Returns
+        -------
+        dict
+            Logs from backend.
+        """
+        if self._context().local:
+            return {}
+        return logs_api(self.project, self.ENTITY_TYPE, self.id)
+
+    def stop(self) -> None:
+        """
+        Stop run.
+
+        Returns
+        -------
+        None
+        """
+        if not self._context().local and not self.spec.local_execution:
+            return stop_api(self.project, self.ENTITY_TYPE, self.id)
+        try:
+            self.status.stop()
+        except AttributeError:
+            raise EntityError("Stop is not supported in local execution.")
+            return
+
+    def resume(self) -> None:
+        """
+        Resume run.
+
+        Returns
+        -------
+        None
+        """
+        if not self._context().local and not self.spec.local_execution:
+            return resume_api(self.project, self.ENTITY_TYPE, self.id)
+
+        try:
+            self.status.resume()
+        except AttributeError:
+            raise EntityError("Resume is not supported in local execution.")
+            return
+        # re-run
+        # TODO verify the logic and order
+        self.run()
+
+    def invoke(self, **kwargs) -> requests.Response:
+        """
+        Invoke run.
+
+        Parameters
+        ----------
+        kwargs
+            Keyword arguments to pass to the request.
+
+        Returns
+        -------
+        requests.Response
+            Response from service.
+        """
+        try:
+            if not self._context().local and not self.spec.local_execution:
+                local = False
+            else:
+                local = True
+            if kwargs is None:
+                kwargs = {}
+            return self.status.invoke(local, **kwargs)
+        except AttributeError:
+            msg = f"Run of type {self.kind} has no invoke operation."
+            raise EntityError(msg)
+
+    ##############################
+    #  Helpers
+    ##############################
+
+    def _is_ready_to_run(self) -> bool:
+        """
+        Check if run is in a state ready for running (BUILT or STOPPED).
+
+        Returns
+        -------
+        bool
+            True if run is in runnable state, False otherwise.
+        """
+        return (self.status.state == State.BUILT.value) or (self.status.state == State.STOPPED.value)
+
+    def _set_status(self, status: dict) -> None:
+        """
+        Set run status.
+
+        Parameters
+        ----------
+        status : dict
+            Status to set.
+
+        Returns
+        -------
+        None
+        """
+        self.status: RunStatus = build_status(self.kind, **status)
+
+    def _set_state(self, state: str) -> None:
+        """
+        Update run state.
+
+        Parameters
+        ----------
+        state : str
+            State to set.
+
+        Returns
+        -------
+        None
+        """
+        self.status.state = state
+
+    def _set_message(self, message: str) -> None:
+        """
+        Update run message.
+
+        Parameters
+        ----------
+        message : str
+            Message to set.
+
+        Returns
+        -------
+        None
+        """
+        self.status.message = message
+
+    def _get_runtime(self) -> Runtime:
+        """
+        Build runtime to build run or execute it.
+
+        Returns
+        -------
+        Runtime
+            Runtime object.
+        """
+        return build_runtime(self.kind, self.project)
+
+    def _get_executable(self, runtime: Runtime) -> dict:
+        """
+        Get object from backend. Reimplemented to avoid circular imports.
+
+        Parameters
+        ----------
+        runtime : Runtime
+            Runtime object.
+
+        Returns
+        -------
+        dict
+            Executable (function or workflow) from backend.
+        """
+        exec_kind = runtime.get_executable_kind()
+        entity_type = registry.get_entity_type(exec_kind)
+        splitted = self.spec.task.split("/")
+        exec_name = splitted[-1].split(":")[0]
+        exec_id = splitted[-1].split(":")[1]
+        return read_entity_api_ctx(
+            exec_name,
+            entity_type=entity_type,
+            project=self.project,
+            entity_id=exec_id,
+        )
+
+    def _get_task(self, runtime: Runtime) -> dict:
+        """
+        Get object from backend. Reimplemented to avoid circular imports.
+
+        Parameters
+        ----------
+        runtime : Runtime
+            Runtime object.
+
+        Returns
+        -------
+        dict
+            Task from backend.
+        """
+        executable_kind = runtime.get_executable_kind()
+        exec_string = f"{executable_kind}://{self.spec.task.split('://')[1]}"
+
+        # Local backend
+        if self._context().local:
+            tasks = list_entity_api_base(self._context().client, EntityTypes.TASK.value)
+            for i in tasks:
+                if i.get("spec").get("function") == exec_string:
+                    return i
+            raise EntityError("Task not found.")
+
+        # Remote backend
+        task_kind = self.spec.task.split("://")[0]
+        params = {"function": exec_string, "kind": task_kind}
+        return list_entity_api_ctx(self.project, EntityTypes.TASK.value, params=params)[0]

digitalhub/entities/run/spec.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+import typing
+
+from digitalhub.entities._base.spec.base import Spec, SpecParams
+from digitalhub.entities.artifact.crud import get_artifact
+from digitalhub.entities.dataitem.crud import get_dataitem
+from digitalhub.entities.entity_types import EntityTypes
+from digitalhub.entities.model.crud import get_model
+from digitalhub.entities.task.models import K8s
+from digitalhub.entities.utils import parse_entity_key
+
+if typing.TYPE_CHECKING:
+    from digitalhub.entities._base.entity.base import Entity
+
+
+ENTITY_FUNC = {
+    EntityTypes.ARTIFACT.value: get_artifact,
+    EntityTypes.DATAITEM.value: get_dataitem,
+    EntityTypes.MODEL.value: get_model,
+}
+
+
+class RunSpec(Spec):
+    """Run specification."""
+
+    def __init__(
+        self,
+        task: str,
+        local_execution: bool = False,
+        function: str | None = None,
+        node_selector: dict | None = None,
+        volumes: list | None = None,
+        resources: dict | None = None,
+        affinity: dict | None = None,
+        tolerations: list | None = None,
+        envs: list | None = None,
+        secrets: list | None = None,
+        profile: str | None = None,
+        **kwargs,
+    ) -> None:
+        self.task = task
+        self.local_execution = local_execution
+        self.function = function
+        self.node_selector = node_selector
+        self.volumes = volumes
+        self.resources = resources
+        self.affinity = affinity
+        self.tolerations = tolerations
+        self.envs = envs
+        self.secrets = secrets
+        self.profile = profile
+
+    def get_inputs(self, as_dict: bool = False) -> dict:
+        """
+        Get inputs.
+
+        Returns
+        -------
+        dict
+            The inputs.
+        """
+        inputs = {}
+        if not hasattr(self, "inputs") or self.inputs is None:
+            return inputs
+
+        for parameter, item in self.inputs.items():
+            parameter_type = self._parse_parameter(parameter)
+
+            # Get entity from key
+            if parameter_type == "key":
+                key = self._collect_key(item)
+                entity = self._collect_entity(key)
+                if as_dict:
+                    entity = entity.to_dict()
+                inputs[parameter] = entity
+
+            # Create entity from parameter
+            elif parameter_type == "create":
+                raise NotImplementedError
+
+        return inputs
+
+    @staticmethod
+    def _parse_parameter(parameter: str) -> str:
+        """
+        Parse parameter.
+
+        Parameters
+        ----------
+        parameter : str
+            Parameter.
+
+        Returns
+        -------
+        str
+            The parsed parameter.
+        """
+        if len(parameter.split(":")) == 1:
+            return "key"
+        return "create"
+
+    @staticmethod
+    def _collect_key(item: str | dict) -> str:
+        """
+        Collect key from item.
+
+        Parameters
+        ----------
+        item : str | dict
+            Key or dict representation of the entity.
+
+        Returns
+        -------
+        str
+            The key.
+        """
+        if isinstance(item, str):
+            return item
+        return item.get("key")
+
+    @staticmethod
+    def _collect_entity(key: str) -> Entity:
+        """
+        Collect entity from key.
+
+        Parameters
+        ----------
+        key : str
+            Key of the entity.
+
+        Returns
+        -------
+        Entity
+            The entity.
+        """
+        _, entity_type, _, _, _ = parse_entity_key(key)
+        return ENTITY_FUNC[entity_type](key)
+
+
+class RunParams(SpecParams, K8s):
+    """
+    Run parameters.
+    """
+
+    function: str = None
+    """The function associated with the run."""
+
+    task: str = None
+    """The task string associated with the run."""
+
+    local_execution: bool = False
+    """Flag to indicate if the run will be executed locally."""

digitalhub/entities/run/status.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+from digitalhub.entities._base.status.base import Status
+from digitalhub.entities.artifact.crud import get_artifact
+from digitalhub.entities.dataitem.crud import get_dataitem
+from digitalhub.entities.entity_types import EntityTypes
+from digitalhub.entities.model.crud import get_model
+from digitalhub.entities.utils import parse_entity_key
+
+ENTITY_FUNC = {
+    EntityTypes.ARTIFACT.value: get_artifact,
+    EntityTypes.DATAITEM.value: get_dataitem,
+    EntityTypes.MODEL.value: get_model,
+}
+
+
+class RunStatus(Status):
+    """
+    Status class for run entities.
+    """
+
+    def __init__(
+        self,
+        state: str,
+        message: str | None = None,
+        outputs: list | None = None,
+        results: dict | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(state, message)
+        self.outputs = outputs
+        self.results = results
+
+        self._any_setter(**kwargs)
+
+    def get_results(self) -> dict:
+        """
+        Get results.
+
+        Returns
+        -------
+        dict
+            The results.
+        """
+        if not hasattr(self, "results") or self.results is None:
+            return {}
+        return self.results
+
+    def get_outputs(self, as_key: bool = False, as_dict: bool = False) -> dict:
+        """
+        Get outputs.
+
+        Parameters
+        ----------
+        as_key : bool
+            If True, return outputs as keys.
+        as_dict : bool
+            If True, return outputs as dictionaries.
+
+        Returns
+        -------
+        dict
+            The outputs.
+        """
+        outputs = {}
+        if not hasattr(self, "outputs") or self.outputs is None:
+            return outputs
+
+        for parameter, key in self.outputs.items():
+            entity_type = self._get_entity_type(key)
+            entity = ENTITY_FUNC[entity_type](key)
+            if as_key:
+                entity = entity.key
+            if as_dict:
+                entity = entity.to_dict()
+            outputs[parameter] = entity
+
+        return outputs
+
+    @staticmethod
+    def _get_entity_type(key: str) -> str:
+        """
+        Get entity type.
+
+        Parameters
+        ----------
+        key : str
+            The key of the entity.
+
+        Returns
+        -------
+        str
+            The entity type.
+        """
+        _, entity_type, _, _, _ = parse_entity_key(key)
+        return entity_type
+
+    def get_values(self, values_list: list) -> dict:
+        """
+        Get values.
+
+        Parameters
+        ----------
+        values_list : list
+            The values list to search in.
+
+        Returns
+        -------
+        dict
+            The values.
+        """
+        if not hasattr(self, "results") or self.results is None:
+            return {}
+        return {k: v for k, v in self.get_results().items() if k in values_list}