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

runnable/extensions/executor/local_container/implementation.py

@@ -0,0 +1,367 @@
+import logging
+from pathlib import Path
+from typing import Dict, cast
+
+from pydantic import Field
+from rich import print
+
+from runnable import defaults, integration, utils
+from runnable.datastore import StepLog
+from runnable.defaults import TypeMapVariable
+from runnable.extensions.executor import GenericExecutor
+from runnable.extensions.nodes import TaskNode
+from runnable.integration import BaseIntegration
+from runnable.nodes import BaseNode
+
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
+
+class LocalContainerExecutor(GenericExecutor):
+    """
+    In the mode of local-container, we execute all the commands in a container.
+
+    Ensure that the local compute has enough resources to finish all your jobs.
+
+    The image of the run, could either be provided as default in the configuration of the execution engine
+    i.e.:
+    execution:
+      type: 'local-container'
+      config:
+        docker_image: the image you want the code to run in.
+
+    or default image could be over-ridden for a single node by providing a docker_image in the step config.
+    i.e:
+    dag:
+      steps:
+        step:
+          executor_config:
+            local-container:
+                docker_image: The image that you want that single step to run in.
+    This image would only be used for that step only.
+
+    This mode does not build the docker image with the latest code for you, it is still left for the user to build
+    and ensure that the docker image provided is the correct one.
+
+    Example config:
+    execution:
+      type: local-container
+      config:
+        docker_image: The default docker image to use if the node does not provide one.
+    """
+
+    service_name: str = "local-container"
+    docker_image: str
+    auto_remove_container: bool = True
+    run_in_local: bool = False
+    environment: Dict[str, str] = Field(default_factory=dict)
+
+    _container_log_location = "/tmp/run_logs/"
+    _container_catalog_location = "/tmp/catalog/"
+    _container_secrets_location = "/tmp/dotenv"
+    _volumes: Dict[str, Dict[str, str]] = {}
+
+    def add_code_identities(self, node: BaseNode, step_log: StepLog, **kwargs):
+        """
+        Call the Base class to add the git code identity and add docker identity
+
+        Args:
+            node (BaseNode): The node we are adding the code identity
+            step_log (Object): The step log corresponding to the node
+        """
+
+        super().add_code_identities(node, step_log)
+
+        if node.node_type in ["success", "fail"]:
+            # Need not add code identities if we are in a success or fail node
+            return
+
+        executor_config = self._resolve_executor_config(node)
+
+        docker_image = executor_config.get("docker_image", None)
+        if docker_image:
+            code_id = self._context.run_log_store.create_code_identity()
+
+            code_id.code_identifier = utils.get_local_docker_image_id(docker_image)
+            code_id.code_identifier_type = "docker"
+            code_id.code_identifier_dependable = True
+            code_id.code_identifier_url = "local docker host"
+            step_log.code_identities.append(code_id)
+
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        We are already in the container, we just execute the node.
+        The node is already prepared for execution.
+        """
+        return self._execute_node(node, map_variable, **kwargs)
+
+    def execute_job(self, node: TaskNode):
+        """
+        Set up the step log and call the execute node
+
+        Args:
+            node (BaseNode): _description_
+        """
+
+        step_log = self._context.run_log_store.create_step_log(node.name, node._get_step_log_name(map_variable=None))
+
+        self.add_code_identities(node=node, step_log=step_log)
+
+        step_log.step_type = node.node_type
+        step_log.status = defaults.PROCESSING
+        self._context.run_log_store.add_step_log(step_log, self._context.run_id)
+
+        command = utils.get_job_execution_command(node)
+        self._spin_container(node=node, command=command)
+
+        # Check the step log status and warn if necessary. Docker errors are generally suppressed.
+        step_log = self._context.run_log_store.get_step_log(
+            node._get_step_log_name(map_variable=None), self._context.run_id
+        )
+        if step_log.status != defaults.SUCCESS:
+            msg = (
+                "Node execution inside the container failed. Please check the logs.\n"
+                "Note: If you do not see any docker issue from your side and the code works properly on local execution"
+                "please raise a bug report."
+            )
+            logger.warning(msg)
+
+    def trigger_job(self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        We come into this step via execute from graph, use trigger job to spin up the container.
+
+
+        If the config has "run_in_local: True", we compute it on local system instead of container.
+        In local container execution, we just spin the container to execute magnus execute_single_node.
+
+        Args:
+            node (BaseNode): The node we are currently executing
+            map_variable (str, optional): If the node is part of the map branch. Defaults to ''.
+        """
+        executor_config = self._resolve_executor_config(node)
+        auto_remove_container = executor_config.get("auto_remove_container", True)
+
+        logger.debug("Here is the resolved executor config")
+        logger.debug(executor_config)
+
+        if executor_config.get("run_in_local", False):
+            # Do not change config but only validate the configuration.
+            # Trigger the job on local system instead of a container
+            integration.validate(self, self._context.run_log_store)
+            integration.validate(self, self._context.catalog_handler)
+            integration.validate(self, self._context.secrets_handler)
+
+            self.execute_node(node=node, map_variable=map_variable, **kwargs)
+            return
+
+        command = utils.get_node_execution_command(node, map_variable=map_variable)
+
+        self._spin_container(
+            node=node,
+            command=command,
+            map_variable=map_variable,
+            auto_remove_container=auto_remove_container,
+            **kwargs,
+        )
+
+        step_log = self._context.run_log_store.get_step_log(node._get_step_log_name(map_variable), self._context.run_id)
+        if step_log.status != defaults.SUCCESS:
+            msg = (
+                "Node execution inside the container failed. Please check the logs.\n"
+                "Note: If you do not see any docker issue from your side and the code works properly on local execution"
+                "please raise a bug report."
+            )
+            logger.warning(msg)
+            step_log.status = defaults.FAIL
+            self._context.run_log_store.add_step_log(step_log, self._context.run_id)
+
+    def _spin_container(
+        self,
+        node: BaseNode,
+        command: str,
+        map_variable: TypeMapVariable = None,
+        auto_remove_container: bool = True,
+        **kwargs,
+    ):
+        """
+        During the flow run, we have to spin up a container with the docker image mentioned
+        and the right log locations
+        """
+        # Conditional import
+        import docker  # pylint: disable=C0415
+
+        try:
+            client = docker.from_env()
+            api_client = docker.APIClient()
+        except Exception as ex:
+            logger.exception("Could not get access to docker")
+            raise Exception("Could not get the docker socket file, do you have docker installed?") from ex
+
+        try:
+            logger.info(f"Running the command {command}")
+            #  Overrides global config with local
+            executor_config = self._resolve_executor_config(node)
+
+            docker_image = executor_config.get("docker_image", None)
+            environment = executor_config.get("environment", {})
+            environment.update(self._context.variables)
+            if not docker_image:
+                raise Exception(
+                    f"Please provide a docker_image using executor_config of the step {node.name} or at global config"
+                )
+
+            # TODO: Should consider using getpass.getuser() when running the docker container? Volume permissions
+            container = client.containers.create(
+                image=docker_image,
+                command=command,
+                auto_remove=False,
+                volumes=self._volumes,
+                network_mode="host",
+                environment=environment,
+            )
+
+            # print(container.__dict__)
+
+            container.start()
+            stream = api_client.logs(container=container.id, timestamps=True, stream=True, follow=True)
+            while True:
+                try:
+                    output = next(stream).decode("utf-8")
+                    output = output.strip("\r\n")
+                    logger.info(output)
+                    print(output)
+                except StopIteration:
+                    logger.info("Docker Run completed")
+                    break
+
+            exit_status = api_client.inspect_container(container.id)["State"]["ExitCode"]
+
+            if auto_remove_container:
+                container.remove(force=True)
+
+            if exit_status != 0:
+                msg = f"Docker command failed with exit code {exit_status}"
+                raise Exception(msg)
+
+        except Exception as _e:
+            logger.exception("Problems with spinning/running the container")
+            raise _e
+
+
+class LocalContainerComputeFileSystemRunLogstore(BaseIntegration):
+    """
+    Integration between local container and file system run log store
+    """
+
+    executor_type = "local-container"
+    service_type = "run_log_store"  # One of secret, catalog, datastore
+    service_provider = "file-system"  # The actual implementation of the service
+
+    def validate(self, **kwargs):
+        if self.executor._is_parallel_execution():  # pragma: no branch
+            msg = (
+                "Run log generated by file-system run log store are not thread safe. "
+                "Inconsistent results are possible because of race conditions to write to the same file.\n"
+                "Consider using partitioned run log store like database for consistent results."
+            )
+            logger.warning(msg)
+
+    def configure_for_traversal(self, **kwargs):
+        from runnable.extensions.run_log_store.file_system.implementation import FileSystemRunLogstore
+
+        self.executor = cast(LocalContainerExecutor, self.executor)
+        self.service = cast(FileSystemRunLogstore, self.service)
+
+        write_to = self.service.log_folder_name
+        self.executor._volumes[str(Path(write_to).resolve())] = {
+            "bind": f"{self.executor._container_log_location}",
+            "mode": "rw",
+        }
+
+    def configure_for_execution(self, **kwargs):
+        from runnable.extensions.run_log_store.file_system.implementation import FileSystemRunLogstore
+
+        self.executor = cast(LocalContainerExecutor, self.executor)
+        self.service = cast(FileSystemRunLogstore, self.service)
+
+        self.service.log_folder = self.executor._container_log_location
+
+
+class LocalContainerComputeFileSystemCatalog(BaseIntegration):
+    """
+    Integration pattern between Local container and File System catalog
+    """
+
+    executor_type = "local-container"
+    service_type = "catalog"  # One of secret, catalog, datastore
+    service_provider = "file-system"  # The actual implementation of the service
+
+    def configure_for_traversal(self, **kwargs):
+        from runnable.extensions.catalog.file_system.implementation import FileSystemCatalog
+
+        self.executor = cast(LocalContainerExecutor, self.executor)
+        self.service = cast(FileSystemCatalog, self.service)
+
+        catalog_location = self.service.catalog_location
+        self.executor._volumes[str(Path(catalog_location).resolve())] = {
+            "bind": f"{self.executor._container_catalog_location}",
+            "mode": "rw",
+        }
+
+    def configure_for_execution(self, **kwargs):
+        from runnable.extensions.catalog.file_system.implementation import FileSystemCatalog
+
+        self.executor = cast(LocalContainerExecutor, self.executor)
+        self.service = cast(FileSystemCatalog, self.service)
+
+        self.service.catalog_location = self.executor._container_catalog_location
+
+
+class LocalContainerComputeDotEnvSecrets(BaseIntegration):
+    """
+    Integration between local container and dot env secrets
+    """
+
+    executor_type = "local-container"
+    service_type = "secrets"  # One of secret, catalog, datastore
+    service_provider = "dotenv"  # The actual implementation of the service
+
+    def validate(self, **kwargs):
+        logger.warning("Using dot env for non local deployments is not ideal, consider options")
+
+    def configure_for_traversal(self, **kwargs):
+        from runnable.extensions.secrets.dotenv.implementation import DotEnvSecrets
+
+        self.executor = cast(LocalContainerExecutor, self.executor)
+        self.service = cast(DotEnvSecrets, self.service)
+
+        secrets_location = self.service.secrets_location
+        self.executor._volumes[str(Path(secrets_location).resolve())] = {
+            "bind": f"{self.executor._container_secrets_location}",
+            "mode": "ro",
+        }
+
+    def configure_for_execution(self, **kwargs):
+        from runnable.extensions.secrets.dotenv.implementation import DotEnvSecrets
+
+        self.executor = cast(LocalContainerExecutor, self.executor)
+        self.service = cast(DotEnvSecrets, self.service)
+
+        self.service.location = self.executor._container_secrets_location
+
+
+class LocalContainerComputeEnvSecretsManager(BaseIntegration):
+    """
+    Integration between local container and env secrets manager
+    """
+
+    executor_type = "local-container"
+    service_type = "secrets"  # One of secret, catalog, datastore
+    service_provider = "env-secrets-manager"  # The actual implementation of the service
+
+    def validate(self, **kwargs):
+        msg = (
+            "Local container executions cannot be used with environment secrets manager. "
+            "Please use a supported secrets manager"
+        )
+        logger.exception(msg)
+        raise Exception(msg)

runnable/extensions/executor/mocked/implementation.py

@@ -0,0 +1,220 @@
+import copy
+import logging
+from typing import Any, Dict, Type, cast
+
+from pydantic import ConfigDict, Field
+
+from runnable import context, defaults
+from runnable.defaults import TypeMapVariable
+from runnable.extensions.executor import GenericExecutor
+from runnable.extensions.nodes import TaskNode
+from runnable.integration import BaseIntegration
+from runnable.nodes import BaseNode
+from runnable.tasks import BaseTaskType
+
+logger = logging.getLogger(defaults.LOGGER_NAME)
+
+
+def create_executable(params: Dict[str, Any], model: Type[BaseTaskType], node_name: str) -> BaseTaskType:
+    class EasyModel(model):  # type: ignore
+        model_config = ConfigDict(extra="ignore")
+
+    swallow_all = EasyModel(**params, node_name=node_name)
+    return swallow_all
+
+
+class MockedExecutor(GenericExecutor):
+    service_name: str = "mocked"
+
+    enable_parallel: bool = defaults.ENABLE_PARALLEL
+
+    patches: Dict[str, Any] = Field(default_factory=dict)
+
+    @property
+    def _context(self):
+        return context.run_context
+
+    def _set_up_for_re_run(self, parameters: Dict[str, Any]) -> None:
+        raise Exception("MockedExecutor does not support re-run")
+
+    def execute_from_graph(self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        This is the entry point to from the graph execution.
+
+        While the self.execute_graph is responsible for traversing the graph, this function is responsible for
+        actual execution of the node.
+
+        If the node type is:
+            * task : We can delegate to _execute_node after checking the eligibility for re-run in cases of a re-run
+            * success: We can delegate to _execute_node
+            * fail: We can delegate to _execute_node
+
+        For nodes that are internally graphs:
+            * parallel: Delegate the responsibility of execution to the node.execute_as_graph()
+            * dag: Delegate the responsibility of execution to the node.execute_as_graph()
+            * map: Delegate the responsibility of execution to the node.execute_as_graph()
+
+        Transpilers will NEVER use this method and will NEVER call ths method.
+        This method should only be used by interactive executors.
+
+        Args:
+            node (Node): The node to execute
+            map_variable (dict, optional): If the node if of a map state, this corresponds to the value of iterable.
+                    Defaults to None.
+        """
+        step_log = self._context.run_log_store.create_step_log(node.name, node._get_step_log_name(map_variable))
+
+        self.add_code_identities(node=node, step_log=step_log)
+
+        step_log.step_type = node.node_type
+        step_log.status = defaults.PROCESSING
+
+        # Add the step log to the database as per the situation.
+        # If its a terminal node, complete it now
+        if node.node_type in ["success", "fail"]:
+            self._context.run_log_store.add_step_log(step_log, self._context.run_id)
+            self._execute_node(node, map_variable=map_variable, **kwargs)
+            return
+
+        # We call an internal function to iterate the sub graphs and execute them
+        if node.is_composite:
+            self._context.run_log_store.add_step_log(step_log, self._context.run_id)
+            node.execute_as_graph(map_variable=map_variable, **kwargs)
+            return
+
+        node_to_send: TaskNode = cast(TaskNode, node).model_copy(deep=True)
+        if node.name not in self.patches:
+            # node is not patched, so mock it
+            step_log.mock = True
+        else:
+            # node is mocked, change the executable to python with the
+            # command as the patch value
+            executable_type = node_to_send.executable.__class__
+            executable = create_executable(
+                self.patches[node.name],
+                executable_type,
+                node_name=node.name,
+            )
+            node_to_send.executable = executable
+            pass
+
+        # Executor specific way to trigger a job
+        self._context.run_log_store.add_step_log(step_log, self._context.run_id)
+        self.trigger_job(node=node_to_send, map_variable=map_variable, **kwargs)
+
+    def trigger_job(self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        Call this method only if we are responsible for traversing the graph via
+        execute_from_graph().
+
+        We are not prepared to execute node as of now.
+
+        Args:
+            node (BaseNode): The node to execute
+            map_variable (str, optional): If the node if of a map state, this corresponds to the value of iterable.
+                    Defaults to ''.
+
+        NOTE: We do not raise an exception as this method is not required by many extensions
+        """
+        self.prepare_for_node_execution()
+        self.execute_node(node=node, map_variable=map_variable, **kwargs)
+
+    def _is_step_eligible_for_rerun(self, node: BaseNode, map_variable: TypeMapVariable = None):
+        """
+        In case of a re-run, this method checks to see if the previous run step status to determine if a re-run is
+        necessary.
+            * True: If its not a re-run.
+            * True: If its a re-run and we failed in the last run or the corresponding logs do not exist.
+            * False: If its a re-run and we succeeded in the last run.
+
+        Most cases, this logic need not be touched
+
+        Args:
+            node (Node): The node to check against re-run
+            map_variable (dict, optional): If the node if of a map state, this corresponds to the value of iterable..
+                        Defaults to None.
+
+        Returns:
+            bool: Eligibility for re-run. True means re-run, False means skip to the next step.
+        """
+        return True
+
+    def _resolve_executor_config(self, node: BaseNode):
+        """
+        The overrides section can contain specific over-rides to an global executor config.
+        To avoid too much clutter in the dag definition, we allow the configuration file to have overrides block.
+        The nodes can over-ride the global config by referring to key in the overrides.
+
+        This function also applies variables to the effective node config.
+
+        For example:
+        # configuration.yaml
+        execution:
+          type: cloud-implementation
+          config:
+            k1: v1
+            k3: v3
+            overrides:
+             custom_config:
+                k1: v11
+                k2: v2 # Could be a mapping internally.
+
+        # in pipeline definition.yaml
+        dag:
+          steps:
+            step1:
+              overrides:
+                cloud-implementation: custom_config
+
+        This method should resolve the node_config to {'k1': 'v11', 'k2': 'v2', 'k3': 'v3'}
+
+        Args:
+            node (BaseNode): The current node being processed.
+
+        """
+        effective_node_config = copy.deepcopy(self.model_dump())
+
+        return effective_node_config
+
+    def execute_job(self, node: TaskNode):
+        pass
+
+    def execute_node(self, node: BaseNode, map_variable: TypeMapVariable = None, **kwargs):
+        """
+        For local execution, we just execute the node.
+
+        Args:
+            node (BaseNode): _description_
+            map_variable (dict[str, str], optional): _description_. Defaults to None.
+        """
+        self._execute_node(node=node, map_variable=map_variable, **kwargs)
+
+
+class LocalContainerComputeFileSystemRunLogstore(BaseIntegration):
+    """
+    Integration between local container and file system run log store
+    """
+
+    executor_type = "local-container"
+    service_type = "run_log_store"  # One of secret, catalog, datastore
+    service_provider = "file-system"  # The actual implementation of the service
+
+    def validate(self, **kwargs):
+        if self.executor._is_parallel_execution():  # pragma: no branch
+            msg = "Mocked executor does not support parallel execution. "
+            logger.warning(msg)
+
+
+class LocalContainerComputeChunkedFSRunLogstore(BaseIntegration):
+    """
+    Integration between local container and file system run log store
+    """
+
+    executor_type = "local-container"
+    service_type = "run_log_store"  # One of secret, catalog, datastore
+    service_provider = "chunked-fs"  # The actual implementation of the service
+
+    def validate(self, **kwargs):
+        if self.executor._is_parallel_execution():  # pragma: no branch
+            msg = "Mocked executor does not support parallel execution. "
+            logger.warning(msg)

runnable/extensions/experiment_tracker/mlflow/implementation.py

@@ -0,0 +1,94 @@
+import functools
+import logging
+from typing import Any, Union
+
+from pydantic import ConfigDict, PrivateAttr
+
+from runnable import defaults
+from runnable.experiment_tracker import BaseExperimentTracker
+
+logger = logging.getLogger(defaults.NAME)
+
+
+class MLFlowExperimentTracker(BaseExperimentTracker):
+    """
+    A MLFlow experiment tracker.
+
+    TODO: Need to set up credentials from secrets
+    """
+
+    service_name: str = "mlflow"
+
+    server_url: str
+    autolog: bool = False
+
+    _default_experiment_name: str = PrivateAttr(default="Default")
+    _active_run_id: str = PrivateAttr(default="")
+    _client: Any = PrivateAttr(default=None)
+
+    model_config = ConfigDict(extra="forbid")
+
+    def model_post_init(self, __context: Any) -> None:
+        try:
+            import mlflow
+        except ImportError:
+            raise Exception("You need to install mlflow to use MLFlowExperimentTracker.")
+
+        self._client = mlflow
+
+        self._client.set_tracking_uri(self.server_url)
+
+        if self.autolog:
+            self._client.autolog(log_models=False)
+
+    @functools.cached_property
+    def experiment_id(self):
+        experiment_name = self._default_experiment_name
+
+        # If a tag is provided, we should create that as our experiment
+        if self._context.tag:
+            experiment_name = self._context.tag
+
+        experiment = self._client.get_experiment_by_name(experiment_name)
+        if not experiment:
+            # Create the experiment and get it.
+            experiment = self._client.create_experiment(experiment_name)
+            experiment = self._client.get_experiment(experiment)
+
+        return experiment.experiment_id
+
+    @functools.cached_property
+    def run_name(self):
+        return self._context.run_id
+
+    @property
+    def client_context(self):
+        if self._active_run_id:
+            return self._client.start_run(
+                run_id=self._active_run_id, experiment_id=self.experiment_id, run_name=self.run_name
+            )
+
+        active_run = self._client.start_run(run_name=self.run_name, experiment_id=self.experiment_id)
+        self._active_run_id = active_run.info.run_id
+        return active_run
+
+    def log_metric(self, key: str, value: Union[int, float], step: int = 0):
+        """
+        Sets the metric in the experiment tracking.
+
+        Args:
+            key (str): The key against you want to store the value
+            value (Any): The value of the metric
+        """
+        if not isinstance(value, float) or isinstance(value, int):
+            msg = f"Only float/int values are accepted as metrics. Setting the metric {key} as parameter {key}_{step}"
+            logger.warning(msg)
+            self.log_parameter(key=key, value=value, step=step)
+            return
+
+        with self.client_context as _:
+            self._client.log_metric(key, float(value), step=step or None)
+
+    def log_parameter(self, key: str, value: Any, step: int = 0):
+        with self.client_context as _:
+            self._client.log_param(key + f"_{str(step)}", value)