ob-metaflow-extensions 1.1.151__py2.py3-none-any.whl → 1.6.2__py2.py3-none-any.whl

metaflow_extensions/outerbounds/plugins/apps/core/app_deploy_decorator.py

@@ -0,0 +1,333 @@
+from typing import List, Optional
+from metaflow.exception import MetaflowException
+from metaflow.decorators import StepDecorator, FlowDecorator
+from metaflow import current
+from metaflow.metaflow_config import KUBERNETES_CONTAINER_IMAGE
+from .deployer import AppDeployer, DeployedApp, PackagedCode
+from .perimeters import PerimeterExtractor
+import os
+
+
+class AppDeployFlowDecorator(FlowDecorator):
+    """
+    Deploy and manage Outerbounds Apps from within a Metaflow flow.
+
+    This decorator enables deploying apps during flow execution, automatically
+    tagging them with Metaflow metadata (flow name, run ID, step, etc.) for
+    easy tracking and management.
+
+    When applied to a flow, it exposes `current.apps` which provides:
+    - Access to the flow's code package and image
+    - Ability to list apps deployed in the current run
+    - Automatic tagging with Metaflow context
+
+    Examples
+    --------
+
+    ```python
+    from metaflow import FlowSpec, step, current, app_deploy
+    from metaflow.apps import AppDeployer
+
+    @app_deploy
+    class MyFlow(FlowSpec):
+
+        @step
+        def start(self):
+            # Deploy an app using the flow's code package
+            deployer = AppDeployer(
+                name="my-service",
+                port=8000,
+                image=current.apps.current_image,
+                code_package=current.apps.metaflow_code_package,
+                commands=["python server.py"],
+            )
+            self.app = deployer.deploy()
+            self.next(self.end)
+
+        @step
+        def end(self):
+            # List all apps deployed in this run
+            apps = current.apps.list()
+            print(f"Deployed {len(apps)} app(s)")
+    ```
+
+
+    MF Add To Current
+    -----------------
+    apps -> metaflow_extensions.outerbounds.plugins.apps.core.app_deploy_decorator.FlowAppManager
+
+        @@ Returns
+        ----------
+        FlowAppManager
+    """
+
+    name = "app_deploy"
+    defaults = {}
+
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
+        self._attributes_with_user_values = (
+            set(attributes.keys()) if attributes is not None else set()
+        )
+
+        super().__init__(attributes, statically_defined, inserted_by)
+
+    def flow_init(
+        self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
+    ):
+
+        from metaflow import decorators
+
+        decorators._attach_decorators(flow, ["app_deploy_internal"])
+        decorators._process_late_attached_decorator(
+            ["app_deploy_internal"],
+            flow,
+            graph,
+            environment,
+            flow_datastore,
+            logger,
+        )
+
+
+class AppDeployInternalDecorator(StepDecorator):
+
+    name = "app_deploy_internal"
+    defaults = {}
+
+    packaged_code = None
+
+    package_url = None
+    package_sha = None
+
+    MAX_ENTROPY = 6
+    MAX_NAME_LENGTH = 150
+
+    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
+        self.logger = logger
+        self.environment = environment
+        self.step = step
+        self.flow_datastore = flow_datastore
+
+    def _resolve_package_url_and_sha(self):
+        return os.environ.get("METAFLOW_CODE_URL", self.package_url), os.environ.get(
+            "METAFLOW_CODE_SHA", self.package_sha
+        )
+
+    def _extract_project_info(self):
+        project = current.get("project_name")
+        branch = current.get("branch_name")
+        is_production = current.get("is_production")
+        return project, branch, is_production
+
+    def task_pre_step(
+        self,
+        step_name,
+        task_datastore,
+        metadata,
+        run_id,
+        task_id,
+        flow,
+        graph,
+        retry_count,
+        max_user_code_retries,
+        ubf_context,
+        inputs,
+    ):
+        perimeter, api_server = PerimeterExtractor.during_programmatic_access()
+        package_url, package_sha = self._resolve_package_url_and_sha()
+        if package_url is None or package_sha is None:
+            # TODO: Think through if we need this abstraction
+            raise MetaflowException(
+                "METAFLOW_CODE_URL or METAFLOW_CODE_SHA is not set. "
+                "Please set METAFLOW_CODE_URL and METAFLOW_CODE_SHA in your environment."
+            )
+        self.packaged_code = PackagedCode(package_url, package_sha)
+        image = os.environ.get("FASTBAKERY_IMAGE", None)
+
+        project, branch, is_production = self._extract_project_info()
+
+        project_info = {}
+        if project is not None:
+            project_info["metaflow/project"] = project
+            project_info["metaflow/branch"] = branch
+            project_info["metaflow/is_production"] = is_production
+
+        default_tags = {
+            "metaflow/flow_name": flow.name,
+            "metaflow/step_name": step_name,
+            "metaflow/run_id": run_id,
+            "metaflow/task_id": task_id,
+            "metaflow/retry_count": retry_count,
+            "metaflow/pathspec": current.pathspec,
+            **project_info,
+        }
+
+        # Set state on AppDeployer using the new API
+        AppDeployer._set_state("perimeter", perimeter)
+        AppDeployer._set_state("api_url", api_server)
+        AppDeployer._set_state(
+            "default_tags", [{k: str(v)} for k, v in default_tags.items()]
+        )
+
+        current._update_env(
+            {
+                "apps": FlowAppManager(
+                    flow.name,
+                    run_id,
+                    self.packaged_code,
+                    image,
+                )
+            }
+        )
+
+    def task_post_step(
+        self, step_name, flow, graph, retry_count, max_user_code_retries
+    ):
+        pass
+
+    def runtime_init(self, flow, graph, package, run_id):
+        # Set some more internal state.
+        self.flow = flow
+        self.graph = graph
+        self.package = package
+        self.run_id = run_id
+
+    def runtime_task_created(
+        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
+    ):
+        # To execute the Kubernetes job, the job container needs to have
+        # access to the code package. We store the package in the datastore
+        # which the pod is able to download as part of it's entrypoint.
+        if not is_cloned:
+            self._save_package_once(self.flow_datastore, self.package)
+
+    @classmethod
+    def _save_package_once(cls, flow_datastore, package):
+        if cls.packaged_code is None:
+            cls.package_url, cls.package_sha = flow_datastore.save_data(
+                [package.blob], len_hint=1
+            )[0]
+            cls.packaged_code = PackagedCode(cls.package_url, cls.package_sha)
+            os.environ["METAFLOW_CODE_URL"] = cls.package_url
+            os.environ["METAFLOW_CODE_SHA"] = cls.package_sha
+
+
+class FlowAppManager:
+    """
+    Manager for apps deployed within a Metaflow flow execution.
+
+    Accessible via `current.apps` when using the `@app_deploy` decorator.
+    Provides access to the flow's code package, container image, and
+    methods to list apps deployed in the current run.
+
+    Attributes
+    ----------
+    metaflow_code_package : PackagedCode
+        The code package for the current flow, ready to use with AppDeployer.
+    current_image : str
+        The container image used by the current task (from fast_bakery or similar).
+    default_image : str
+        The default Kubernetes container image from Metaflow config.
+
+    Examples
+    --------
+
+    ```python
+    # python myflow.py --environment=fast-bakery run --with kubernetes
+    from metaflow.apps import AppDeployer
+
+    @pypi(packages={"flask": ">=2.0", "requests": ">=2.28"})
+    @step
+    def deploy_step(self):
+        image = current.apps.current_image
+        if image is None:
+            image = current.apps.default_image
+        # Use the flow's code package directly
+        deployer = AppDeployer(
+            name="my-app",
+            port=8000,
+            image=image,
+            code_package=current.apps.metaflow_code_package,
+            commands=["python app.py"],
+        )
+        deployed = deployer.deploy()
+
+        # List apps from this run
+        apps = current.apps.list()
+    ```
+    """
+
+    def __init__(
+        self,
+        flow_name: str,
+        run_id: str,
+        package: PackagedCode,
+        image: Optional[str] = None,
+    ) -> None:
+        self._package = package
+        self._image = image
+        self._runid = run_id
+        self._flowname = flow_name
+
+    @property
+    def metaflow_code_package(self) -> PackagedCode:
+        """
+        The flow's code package for use with AppDeployer.
+
+        Returns
+        -------
+        PackagedCode
+            A namedtuple with `url` and `key` fields pointing to the
+            packaged metaflow's code package stored in the datastore.
+        """
+        return self._package
+
+    @property
+    def current_image(self) -> str:
+        """
+        The container image for the current task.
+
+        Returns
+        -------
+        str
+            Image URI from fast bakery image or None if not set.
+        """
+        return self._image
+
+    @property
+    def default_image(self) -> str:
+        """
+        The default Kubernetes container image from Metaflow config.
+
+        Returns
+        -------
+        str
+            The KUBERNETES_CONTAINER_IMAGE from Metaflow configuration.
+        """
+        return KUBERNETES_CONTAINER_IMAGE
+
+    def list(self) -> List["DeployedApp"]:
+        """
+        List apps deployed in the current Metaflow run.
+
+        Returns apps tagged with this flow's name and run ID.
+
+        Returns
+        -------
+        List[DeployedApp]
+            Apps deployed during this flow execution.
+
+        Examples
+        --------
+
+        ```python
+        apps = current.apps.list()
+        for app in apps:
+            print(f"{app.name}: {app.public_url}")
+        ```
+        """
+        return AppDeployer.list_deployments(
+            tags=[
+                {"metaflow/run_id": self._runid},
+                {"metaflow/flow_name": self._flowname},
+            ]
+        )