runnable 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

runnable/nodes.py

@@ -0,0 +1,489 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+import runnable.context as context
+from runnable import defaults, exceptions
+from runnable.datastore import StepAttempt
+from runnable.defaults import TypeMapVariable
+
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
+# --8<-- [start:docs]
+
+
+class BaseNode(ABC, BaseModel):
+    """
+    Base class with common functionality provided for a Node of a graph.
+
+    A node of a graph could be a
+        * single execution node as task, success, fail.
+        * Could be graph in itself as parallel, dag and map.
+        * could be a convenience function like as-is.
+
+    The name is relative to the DAG.
+    The internal name of the node, is absolute name in dot path convention.
+        This has one to one mapping to the name in the run log
+    The internal name of a node, should always be odd when split against dot.
+
+    The internal branch name, only applies for branched nodes, is the branch it belongs to.
+    The internal branch name should always be even when split against dot.
+    """
+
+    node_type: str = Field(serialization_alias="type")
+    name: str
+    internal_name: str = Field(exclude=True)
+    internal_branch_name: str = Field(default="", exclude=True)
+    is_composite: bool = Field(default=False, exclude=True)
+
+    @property
+    def _context(self):
+        return context.run_context
+
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=False)
+
+    @field_validator("name")
+    @classmethod
+    def validate_name(cls, name: str):
+        if "." in name or "%" in name:
+            raise ValueError("Node names cannot have . or '%' in them")
+        return name
+
+    def _command_friendly_name(self, replace_with=defaults.COMMAND_FRIENDLY_CHARACTER) -> str:
+        """
+        Replace spaces with special character for spaces.
+        Spaces in the naming of the node is convenient for the user but causes issues when used programmatically.
+
+        Returns:
+            str: The command friendly name of the node
+        """
+        return self.internal_name.replace(" ", replace_with)
+
+    @classmethod
+    def _get_internal_name_from_command_name(cls, command_name: str) -> str:
+        """
+        Replace Magnus specific character (%) with whitespace.
+        The opposite of _command_friendly_name.
+
+        Args:
+            command_name (str): The command friendly node name
+
+        Returns:
+            str: The internal name of the step
+        """
+        return command_name.replace(defaults.COMMAND_FRIENDLY_CHARACTER, " ")
+
+    @classmethod
+    def _resolve_map_placeholders(cls, name: str, map_variable: TypeMapVariable = None) -> str:
+        """
+        If there is no map step used, then we just return the name as we find it.
+
+        If there is a map variable being used, replace every occurrence of the map variable placeholder with
+        the value sequentially.
+
+        For example:
+        1). dag:
+              start_at: step1
+              steps:
+                step1:
+                    type: map
+                    iterate_on: y
+                    iterate_as: y_i
+                    branch:
+                        start_at: map_step1
+                        steps:
+                            map_step1: # internal_name step1.placeholder.map_step1
+                                type: task
+                                command: a.map_func
+                                command_type: python
+                                next: map_success
+                            map_success:
+                                type: success
+                            map_failure:
+                                type: fail
+
+            and if y is ['a', 'b', 'c'].
+
+            This method would be called 3 times with map_variable = {'y_i': 'a'}, map_variable = {'y_i': 'b'} and
+            map_variable = {'y_i': 'c'} corresponding to the three branches.
+
+        For nested map branches, we would get the map_variables ordered hierarchically.
+
+        Args:
+            name (str): The name to resolve
+            map_variable (dict): The dictionary of map variables
+
+        Returns:
+            [str]: The resolved name
+        """
+        if not map_variable:
+            return name
+
+        for _, value in map_variable.items():
+            name = name.replace(defaults.MAP_PLACEHOLDER, str(value), 1)
+
+        return name
+
+    def _get_step_log_name(self, map_variable: TypeMapVariable = None) -> str:
+        """
+        For every step in the dag, there is a corresponding step log name.
+        This method returns the step log name in dot path convention.
+
+        All node types except a map state has a "static" defined step_log names and are equivalent to internal_name.
+        For nodes belonging to map state, the internal name has a placeholder that is replaced at runtime.
+
+        Args:
+            map_variable (dict): If the node is of type map, the names are based on the current iteration state of the
+            parameter.
+
+        Returns:
+            str: The dot path name of the step log name
+        """
+        return self._resolve_map_placeholders(self.internal_name, map_variable=map_variable)
+
+    def _get_branch_log_name(self, map_variable: TypeMapVariable = None) -> str:
+        """
+        For nodes that are internally branches, this method returns the branch log name.
+        The branch log name is in dot path convention.
+
+        For nodes that are not map, the internal branch name is equivalent to the branch name.
+        For map nodes, the internal branch name has a placeholder that is replaced at runtime.
+
+        Args:
+            map_variable (dict): If the node is of type map, the names are based on the current iteration state of the
+            parameter.
+
+        Returns:
+            str: The dot path name of the branch log
+        """
+        return self._resolve_map_placeholders(self.internal_branch_name, map_variable=map_variable)
+
+    def __str__(self) -> str:  # pragma: no cover
+        """
+        String representation of the node.
+
+        Returns:
+            str: The string representation of the node.
+        """
+        return f"Node of type {self.node_type} and name {self.internal_name}"
+
+    @abstractmethod
+    def _get_on_failure_node(self) -> str:
+        """
+        If the node defines a on_failure node in the config, return this or None.
+
+        The naming is relative to the dag, the caller is supposed to resolve it to the correct graph
+
+        Returns:
+            str: The on_failure node defined by the dag or ''
+        This is a base implementation which the BaseNode does not satisfy
+        """
+        ...
+
+    @abstractmethod
+    def _get_next_node(self) -> str:
+        """
+        Return the next node as defined by the config.
+
+        Returns:
+            str: The node name, relative to the dag, as defined by the config
+        """
+        ...
+
+    @abstractmethod
+    def _is_terminal_node(self) -> bool:
+        """
+        Returns whether a node has a next node
+
+        Returns:
+            bool: True or False of whether there is next node.
+        """
+        ...
+
+    @abstractmethod
+    def _get_catalog_settings(self) -> Dict[str, Any]:
+        """
+        If the node defines a catalog settings, return it or None
+
+        Returns:
+            dict: catalog settings defined as per the node or None
+        """
+        ...
+
+    @abstractmethod
+    def _get_branch_by_name(self, branch_name: str):
+        """
+        Retrieve a branch by name.
+
+        The name is expected to follow a dot path convention.
+
+        Args:
+            branch_name (str): [description]
+
+        Raises:
+            Exception: [description]
+        """
+        ...
+
+    def _get_neighbors(self) -> List[str]:
+        """
+        Gets the connecting neighbor nodes, either the "next" node or "on_failure" node.
+
+        Returns:
+            list: List of connected neighbors for a given node. Empty if terminal node.
+        """
+        neighbors = []
+        try:
+            next_node = self._get_next_node()
+            neighbors += [next_node]
+        except exceptions.TerminalNodeError:
+            pass
+
+        try:
+            fail_node = self._get_on_failure_node()
+            if fail_node:
+                neighbors += [fail_node]
+        except exceptions.TerminalNodeError:
+            pass
+
+        return neighbors
+
+    @abstractmethod
+    def _get_executor_config(self, executor_type: str) -> str:
+        """
+        Return the executor config of the node, if defined, or empty dict
+
+        Args:
+            executor_type (str): The executor type that the config refers to.
+
+        Returns:
+            dict: The executor config, if defined or an empty dict
+        """
+        ...
+
+    @abstractmethod
+    def _get_max_attempts(self) -> int:
+        """
+        The number of max attempts as defined by the config or 1.
+
+        Returns:
+            int: The number of maximum retries as defined by the config or 1.
+        """
+        ...
+
+    @abstractmethod
+    def execute(self, mock=False, map_variable: TypeMapVariable = None, **kwargs) -> StepAttempt:
+        """
+        The actual function that does the execution of the command in the config.
+
+        Should only be implemented for task, success, fail and as-is and never for
+        composite nodes.
+
+        Args:
+            executor (magnus.executor.BaseExecutor): The executor class
+            mock (bool, optional): Don't run, just pretend. Defaults to False.
+            map_variable (str, optional): The value of the map iteration variable, if part of a map node.
+                Defaults to ''.
+
+        Raises:
+            NotImplementedError: Base class, hence not implemented.
+        """
+        ...
+
+    @abstractmethod
+    def execute_as_graph(self, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        This function would be called to set up the execution of the individual
+        branches of a composite node.
+
+        Function should only be implemented for composite nodes like dag, map, parallel.
+
+        Args:
+            executor (magnus.executor.BaseExecutor): The executor.
+
+        Raises:
+            NotImplementedError: Base class, hence not implemented.
+        """
+        ...
+
+    @abstractmethod
+    def fan_out(self, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        This function would be called to set up the execution of the individual
+        branches of a composite node.
+
+        Function should only be implemented for composite nodes like dag, map, parallel.
+
+        Args:
+            executor (magnus.executor.BaseExecutor): The executor.
+            map_variable (str, optional): The value of the map iteration variable, if part of a map node.
+
+        Raises:
+            Exception: If the node is not a composite node.
+        """
+        ...
+
+    @abstractmethod
+    def fan_in(self, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        This function would be called to tear down the execution of the individual
+        branches of a composite node.
+
+        Function should only be implemented for composite nodes like dag, map, parallel.
+
+        Args:
+            executor (magnus.executor.BaseExecutor): The executor.
+            map_variable (str, optional): The value of the map iteration variable, if part of a map node.
+
+        Raises:
+            Exception: If the node is not a composite node.
+        """
+        ...
+
+    @classmethod
+    @abstractmethod
+    def parse_from_config(cls, config: Dict[str, Any]) -> "BaseNode":
+        """
+        Parse the config from the user and create the corresponding node.
+
+        Args:
+            config (Dict[str, Any]): The config of the node from the yaml or from the sdk.
+
+        Returns:
+            BaseNode: The corresponding node.
+        """
+        ...
+
+
+# --8<-- [end:docs]
+class TraversalNode(BaseNode):
+    next_node: str = Field(serialization_alias="next")
+    on_failure: str = Field(default="")
+    overrides: Dict[str, str] = Field(default_factory=dict)
+
+    def _get_on_failure_node(self) -> str:
+        """
+        If the node defines a on_failure node in the config, return this or None.
+
+        The naming is relative to the dag, the caller is supposed to resolve it to the correct graph
+
+        Returns:
+            str: The on_failure node defined by the dag or ''
+        This is a base implementation which the BaseNode does not satisfy
+        """
+        return self.on_failure
+
+    def _get_next_node(self) -> str:
+        """
+        Return the next node as defined by the config.
+
+        Returns:
+            str: The node name, relative to the dag, as defined by the config
+        """
+
+        return self.next_node
+
+    def _is_terminal_node(self) -> bool:
+        """
+        Returns whether a node has a next node
+
+        Returns:
+            bool: True or False of whether there is next node.
+        """
+        return False
+
+    def _get_executor_config(self, executor_type) -> str:
+        return self.overrides.get(executor_type) or ""
+
+
+class CatalogStructure(BaseModel):
+    model_config = ConfigDict(extra="forbid")  # Need to forbid
+
+    get: List[str] = Field(default_factory=list)
+    put: List[str] = Field(default_factory=list)
+
+
+class ExecutableNode(TraversalNode):
+    catalog: Optional[CatalogStructure] = Field(default=None)
+    max_attempts: int = Field(default=1, ge=1)
+
+    def _get_catalog_settings(self) -> Dict[str, Any]:
+        """
+        If the node defines a catalog settings, return it or None
+
+        Returns:
+            dict: catalog settings defined as per the node or None
+        """
+        if self.catalog:
+            return self.catalog.model_dump()
+        return {}
+
+    def _get_max_attempts(self) -> int:
+        return self.max_attempts
+
+    def _get_branch_by_name(self, branch_name: str):
+        raise Exception("This is an executable node and does not have branches")
+
+    def execute_as_graph(self, map_variable: TypeMapVariable = None, **kwargs):
+        raise Exception("This is an executable node and does not have a graph")
+
+    def fan_in(self, map_variable: TypeMapVariable = None, **kwargs):
+        raise Exception("This is an executable node and does not have a fan in")
+
+    def fan_out(self, map_variable: TypeMapVariable = None, **kwargs):
+        raise Exception("This is an executable node and does not have a fan out")
+
+
+class CompositeNode(TraversalNode):
+    def _get_catalog_settings(self) -> Dict[str, Any]:
+        """
+        If the node defines a catalog settings, return it or None
+
+        Returns:
+            dict: catalog settings defined as per the node or None
+        """
+        raise Exception("This is a composite node and does not have a catalog settings")
+
+    def _get_max_attempts(self) -> int:
+        raise Exception("This is a composite node and does not have a max_attempts")
+
+    def execute(self, mock=False, map_variable: TypeMapVariable = None, **kwargs) -> StepAttempt:
+        raise Exception("This is a composite node and does not have an execute function")
+
+
+class TerminalNode(BaseNode):
+    def _get_on_failure_node(self) -> str:
+        raise exceptions.TerminalNodeError()
+
+    def _get_next_node(self) -> str:
+        raise exceptions.TerminalNodeError()
+
+    def _is_terminal_node(self) -> bool:
+        return True
+
+    def _get_catalog_settings(self) -> Dict[str, Any]:
+        raise exceptions.TerminalNodeError()
+
+    def _get_branch_by_name(self, branch_name: str):
+        raise exceptions.TerminalNodeError()
+
+    def _get_executor_config(self, executor_type) -> str:
+        raise exceptions.TerminalNodeError()
+
+    def _get_max_attempts(self) -> int:
+        return 1
+
+    def execute_as_graph(self, map_variable: TypeMapVariable = None, **kwargs):
+        raise exceptions.TerminalNodeError()
+
+    def fan_in(self, map_variable: TypeMapVariable = None, **kwargs):
+        raise exceptions.TerminalNodeError()
+
+    def fan_out(self, map_variable: TypeMapVariable = None, **kwargs):
+        raise exceptions.TerminalNodeError()
+
+    @classmethod
+    def parse_from_config(cls, config: Dict[str, Any]) -> "TerminalNode":
+        return cls(**config)

runnable/parameters.py

@@ -0,0 +1,183 @@
+import inspect
+import json
+import logging
+import os
+from typing import Any, Dict, Optional, Type, Union
+
+from pydantic import BaseModel, ConfigDict
+from typing_extensions import Callable
+
+from runnable import defaults
+from runnable.defaults import TypeMapVariable
+from runnable.utils import remove_prefix
+
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
+
+def get_user_set_parameters(remove: bool = False) -> Dict[str, Any]:
+    """
+    Scans the environment variables for any user returned parameters that have a prefix MAGNUS_PRM_.
+
+    This function does not deal with any type conversion of the parameters.
+    It just deserializes the parameters and returns them as a dictionary.
+
+    Args:
+        remove (bool, optional): Flag to remove the parameter if needed. Defaults to False.
+
+    Returns:
+        dict: The dictionary of found user returned parameters
+    """
+    parameters = {}
+    for env_var, value in os.environ.items():
+        if env_var.startswith(defaults.PARAMETER_PREFIX):
+            key = remove_prefix(env_var, defaults.PARAMETER_PREFIX)
+            try:
+                parameters[key.lower()] = json.loads(value)
+            except json.decoder.JSONDecodeError:
+                logger.error(f"Parameter {key} could not be JSON decoded, adding the literal value")
+                parameters[key.lower()] = value
+
+            if remove:
+                del os.environ[env_var]
+    return parameters
+
+
+def set_user_defined_params_as_environment_variables(params: Dict[str, Any]):
+    """
+    Sets the user set parameters as environment variables.
+
+    At this point in time, the params are already in Dict or some kind of literal
+
+    Args:
+        parameters (Dict[str, Any]): The parameters to set as environment variables
+        update (bool, optional): Flag to update the environment variables. Defaults to True.
+
+    """
+    for key, value in params.items():
+        logger.info(f"Storing parameter {key} with value: {value}")
+        environ_key = defaults.PARAMETER_PREFIX + key
+
+        os.environ[environ_key] = serialize_parameter_as_str(value)
+
+
+def cast_parameters_as_type(value: Any, newT: Optional[Type] = None) -> Union[Any, BaseModel, Dict[str, Any]]:
+    """
+    Casts the environment variable to the given type.
+
+    Note: Only pydantic models special, everything else just goes through.
+
+    Args:
+        value (Any): The value to cast
+        newT (T): The type to cast to
+
+    Returns:
+        T: The casted value
+
+    Examples:
+        >>> class MyBaseModel(BaseModel):
+        ...     a: int
+        ...     b: str
+        >>>
+        >>> class MyDict(dict):
+        ...     pass
+        >>>
+        >>> cast_parameters_as_type({"a": 1, "b": "2"}, MyBaseModel)
+        MyBaseModel(a=1, b="2")
+        >>> cast_parameters_as_type({"a": 1, "b": "2"}, MyDict)
+        MyDict({'a': 1, 'b': '2'})
+        >>> cast_parameters_as_type(MyBaseModel(a=1, b="2"), MyBaseModel)
+        MyBaseModel(a=1, b="2")
+        >>> cast_parameters_as_type(MyDict({"a": 1, "b": "2"}), MyBaseModel)
+        MyBaseModel(a=1, b="2")
+        >>> cast_parameters_as_type({"a": 1, "b": "2"}, MyDict[str, int])
+        MyDict({'a': 1, 'b': '2'})
+        >>> cast_parameters_as_type({"a": 1, "b": "2"}, Dict[str, int])
+        MyDict({'a': 1, 'b': '2'})
+        >>> with pytest.warns(UserWarning):
+        ...     cast_parameters_as_type(1, MyBaseModel)
+        MyBaseModel(a=1, b=None)
+        >>> with pytest.raises(TypeError):
+        ...     cast_parameters_as_type(1, MyDict)
+    """
+    if not newT:
+        return value
+
+    if issubclass(newT, BaseModel):
+        return newT(**value)
+
+    if issubclass(newT, Dict):
+        return dict(value)
+
+    if type(value) != newT:
+        logger.warning(f"Casting {value} of {type(value)} to {newT} seems wrong!!")
+
+    return newT(value)
+
+
+def serialize_parameter_as_str(value: Any) -> str:
+    if isinstance(value, BaseModel):
+        return json.dumps(value.model_dump())
+
+    return json.dumps(value)
+
+
+def filter_arguments_for_func(
+    func: Callable[..., Any], params: Dict[str, Any], map_variable: TypeMapVariable = None
+) -> Dict[str, Any]:
+    """
+    Inspects the function to be called as part of the pipeline to find the arguments of the function.
+    Matches the function arguments to the parameters available either by command line or by up stream steps.
+
+
+    Args:
+        func (Callable): The function to inspect
+        parameters (dict): The parameters available for the run
+
+    Returns:
+        dict: The parameters matching the function signature
+    """
+    function_args = inspect.signature(func).parameters
+
+    # Update parameters with the map variables
+    params.update(map_variable or {})
+
+    unassigned_params = set(params.keys())
+    bound_args = {}
+    for name, value in function_args.items():
+        if name not in params:
+            # No parameter of this name was provided
+            if value.default == inspect.Parameter.empty:
+                # No default value is given in the function signature. error as parameter is required.
+                raise ValueError(f"Parameter {name} is required for {func.__name__} but not provided")
+            # default value is given in the function signature, nothing further to do.
+            continue
+
+        if issubclass(value.annotation, BaseModel):
+            # We try to cast it as a pydantic model.
+            named_param = params[name]
+
+            if not isinstance(named_param, dict):
+                # A case where the parameter is a one attribute model
+                named_param = {name: named_param}
+
+            bound_model = bind_args_for_pydantic_model(named_param, value.annotation)
+            bound_args[name] = bound_model
+            unassigned_params = unassigned_params.difference(bound_model.model_fields.keys())
+        else:
+            # simple python data type.
+            bound_args[name] = cast_parameters_as_type(params[name], value.annotation)  # type: ignore
+
+        unassigned_params.remove(name)
+
+        params = {key: params[key] for key in unassigned_params}  # remove keys from params if they are assigned
+
+    return bound_args
+
+
+def bind_args_for_pydantic_model(params: Dict[str, Any], model: Type[BaseModel]) -> BaseModel:
+    class EasyModel(model):  # type: ignore
+        model_config = ConfigDict(extra="ignore")
+
+    swallow_all = EasyModel(**params)
+    bound_model = model(**swallow_all.model_dump())
+    return bound_model