flyte 2.0.0b32__py3-none-any.whl

flyte/_deploy.py

@@ -0,0 +1,408 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
+
+import cloudpickle
+import rich.repr
+
+from flyte.models import SerializationContext
+from flyte.syncify import syncify
+
+from ._environment import Environment
+from ._image import Image
+from ._initialize import ensure_client, get_client, get_init_config, requires_initialization
+from ._logging import logger
+from ._task import TaskTemplate
+from ._task_environment import TaskEnvironment
+
+if TYPE_CHECKING:
+    from flyteidl2.task import task_definition_pb2
+
+    from ._code_bundle import CopyFiles
+    from ._deployer import DeployedEnvironment, DeploymentContext
+    from ._internal.imagebuild.image_builder import ImageCache
+
+
+@rich.repr.auto
+@dataclass
+class DeploymentPlan:
+    envs: Dict[str, Environment]
+    version: Optional[str] = None
+
+
+@rich.repr.auto
+@dataclass
+class DeployedTask:
+    deployed_task: task_definition_pb2.TaskSpec
+    deployed_triggers: List[task_definition_pb2.TaskTrigger]
+
+    def get_name(self) -> str:
+        """
+        Returns the name of the deployed environment.
+        Returns:
+        """
+        return self.deployed_task.task_template.id.name
+
+    def summary_repr(self) -> str:
+        """
+        Returns a summary representation of the deployed task.
+        """
+        return (
+            f"DeployedTask(name={self.deployed_task.task_template.id.name}, "
+            f"version={self.deployed_task.task_template.id.version})"
+        )
+
+    def table_repr(self) -> List[Tuple[str, ...]]:
+        """
+        Returns a table representation of the deployed task.
+        """
+        return [
+            ("type", "task"),
+            ("name", self.deployed_task.task_template.id.name),
+            ("version", self.deployed_task.task_template.id.version),
+            ("triggers", ",".join([t.name for t in self.deployed_triggers])),
+        ]
+
+
+@rich.repr.auto
+@dataclass
+class DeployedTaskEnvironment:
+    env: TaskEnvironment
+    deployed_entities: List[DeployedTask]
+
+    def get_name(self) -> str:
+        """
+        Returns the name of the deployed environment.
+        """
+        return self.env.name
+
+    def summary_repr(self) -> str:
+        """
+        Returns a summary representation of the deployment.
+        """
+        entities = ", ".join(f"{e.summary_repr()}" for e in self.deployed_entities or [])
+        return f"Deployment(env=[{self.env.name}], entities=[{entities}])"
+
+    def table_repr(self) -> List[List[Tuple[str, ...]]]:
+        """
+        Returns a detailed representation of the deployed tasks.
+        """
+        tuples = []
+        if self.deployed_entities:
+            for e in self.deployed_entities:
+                tuples.append(e.table_repr())
+        return tuples
+
+    def env_repr(self) -> List[Tuple[str, ...]]:
+        """
+        Returns a detailed representation of the deployed environments.
+        """
+        env = self.env
+        return [
+            ("environment", env.name),
+            ("image", env.image.uri if isinstance(env.image, Image) else env.image or ""),
+        ]
+
+
+@rich.repr.auto
+@dataclass(frozen=True)
+class Deployment:
+    envs: Dict[str, DeployedEnvironment]
+
+    def summary_repr(self) -> str:
+        """
+        Returns a summary representation of the deployment.
+        """
+        envs = ", ".join(f"{e.summary_repr()}" for e in self.envs.values() or [])
+        return f"Deployment(envs=[{envs}])"
+
+    def table_repr(self) -> List[List[Tuple[str, ...]]]:
+        """
+        Returns a detailed representation of the deployed tasks.
+        """
+        tuples = []
+        for d in self.envs.values():
+            tuples.extend(d.table_repr())
+        return tuples
+
+    def env_repr(self) -> List[List[Tuple[str, ...]]]:
+        """
+        Returns a detailed representation of the deployed environments.
+        """
+        tuples = []
+        for d in self.envs.values():
+            tuples.append(d.env_repr())
+        return tuples
+
+
+async def _deploy_task(
+    task: TaskTemplate, serialization_context: SerializationContext, dryrun: bool = False
+) -> DeployedTask:
+    """
+    Deploy the given task.
+    """
+    ensure_client()
+    import grpc.aio
+    from flyteidl2.task import task_definition_pb2, task_service_pb2
+
+    import flyte.errors
+
+    from ._internal.runtime.convert import convert_upload_default_inputs
+    from ._internal.runtime.task_serde import translate_task_to_wire
+    from ._internal.runtime.trigger_serde import to_task_trigger
+
+    image_uri = task.image.uri if isinstance(task.image, Image) else task.image
+
+    try:
+        if dryrun:
+            return DeployedTask(translate_task_to_wire(task, serialization_context), [])
+
+        default_inputs = await convert_upload_default_inputs(task.interface)
+        spec = translate_task_to_wire(task, serialization_context, default_inputs=default_inputs)
+
+        msg = f"Deploying task {task.name}, with image {image_uri} version {serialization_context.version}"
+        if spec.task_template.HasField("container") and spec.task_template.container.args:
+            msg += f" from {spec.task_template.container.args[-3]}.{spec.task_template.container.args[-1]}"
+        logger.info(msg)
+        task_id = task_definition_pb2.TaskIdentifier(
+            org=spec.task_template.id.org,
+            project=spec.task_template.id.project,
+            domain=spec.task_template.id.domain,
+            version=spec.task_template.id.version,
+            name=spec.task_template.id.name,
+        )
+
+        deployable_triggers = []
+        for t in task.triggers:
+            inputs = spec.task_template.interface.inputs
+            default_inputs = spec.default_inputs
+            deployable_triggers.append(
+                await to_task_trigger(
+                    t=t, task_name=task.name, task_inputs=inputs, task_default_inputs=list(default_inputs)
+                )
+            )
+
+        try:
+            await get_client().task_service.DeployTask(
+                task_service_pb2.DeployTaskRequest(
+                    task_id=task_id,
+                    spec=spec,
+                    triggers=deployable_triggers,
+                )
+            )
+            logger.info(f"Deployed task {task.name} with version {task_id.version}")
+        except grpc.aio.AioRpcError as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                logger.info(f"Task {task.name} with image {image_uri} already exists, skipping deployment.")
+                return DeployedTask(spec, deployable_triggers)
+            raise
+
+        return DeployedTask(spec, deployable_triggers)
+    except Exception as e:
+        logger.error(f"Failed to deploy task {task.name} with image {image_uri}: {e}")
+        raise flyte.errors.DeploymentError(
+            f"Failed to deploy task {task.name} file{task.source_file} with image {image_uri}, Error: {e!s}"
+        ) from e
+
+
+async def _build_image_bg(env_name: str, image: Image) -> Tuple[str, str]:
+    """
+    Build the image in the background and return the environment name and the built image.
+    """
+    from ._build import build
+
+    logger.info(f"Building image {image.name} for environment {env_name}")
+    return env_name, await build.aio(image)
+
+
+async def _build_images(deployment: DeploymentPlan, image_refs: Dict[str, str] | None = None) -> ImageCache:
+    """
+    Build the images for the given deployment plan and update the environment with the built image.
+    """
+    from ._internal.imagebuild.image_builder import ImageCache
+
+    if image_refs is None:
+        image_refs = {}
+
+    images = []
+    image_identifier_map = {}
+    for env_name, env in deployment.envs.items():
+        if not isinstance(env.image, str):
+            if env.image._ref_name is not None:
+                if env.image._ref_name in image_refs:
+                    # If the image is set in the config, set it as the base_image
+                    image_uri = image_refs[env.image._ref_name]
+                    env.image = env.image.clone(base_image=image_uri)
+                else:
+                    raise ValueError(
+                        f"Image name '{env.image._ref_name}' not found in config. Available: {list(image_refs.keys())}"
+                    )
+                if not env.image._layers:
+                    # No additional layers, use the base_image directly without building
+                    image_identifier_map[env_name] = image_uri
+                    continue
+            logger.debug(f"Building Image for environment {env_name}, image: {env.image}")
+            images.append(_build_image_bg(env_name, env.image))
+
+        elif env.image == "auto" and "auto" not in image_identifier_map:
+            if "default" in image_refs:
+                # If the default image is set through CLI, use it instead
+                image_uri = image_refs["default"]
+                image_identifier_map[env_name] = image_uri
+                continue
+            auto_image = Image.from_debian_base()
+            images.append(_build_image_bg(env_name, auto_image))
+    final_images = await asyncio.gather(*images)
+
+    for env_name, image_uri in final_images:
+        logger.warning(f"Built Image for environment {env_name}, image: {image_uri}")
+        image_identifier_map[env_name] = image_uri
+
+    return ImageCache(image_lookup=image_identifier_map)
+
+
+async def _deploy_task_env(context: DeploymentContext) -> DeployedTaskEnvironment:
+    """
+    Deploy the given task environment.
+    """
+    ensure_client()
+    env = context.environment
+    if not isinstance(env, TaskEnvironment):
+        raise ValueError(f"Expected TaskEnvironment, got {type(env)}")
+
+    task_coros = []
+    for task in env.tasks.values():
+        task_coros.append(_deploy_task(task, context.serialization_context, dryrun=context.dryrun))
+    deployed_task_vals = await asyncio.gather(*task_coros)
+    deployed_tasks = []
+    for t in deployed_task_vals:
+        deployed_tasks.append(t)
+    return DeployedTaskEnvironment(env=env, deployed_entities=deployed_tasks)
+
+
+@requires_initialization
+async def apply(deployment_plan: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
+    import flyte.errors
+
+    from ._code_bundle import build_code_bundle
+    from ._deployer import DeploymentContext, get_deployer
+
+    cfg = get_init_config()
+
+    image_cache = await _build_images(deployment_plan, cfg.images)
+
+    if copy_style == "none" and not deployment_plan.version:
+        raise flyte.errors.DeploymentError("Version must be set when copy_style is none")
+    else:
+        # if this is an AppEnvironment.include, skip code bundling here and build a code bundle at the
+        # app._deploy._deploy_app function
+        code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
+        if deployment_plan.version:
+            version = deployment_plan.version
+        else:
+            h = hashlib.md5()
+            h.update(cloudpickle.dumps(deployment_plan.envs))
+            h.update(code_bundle.computed_version.encode("utf-8"))
+            h.update(cloudpickle.dumps(image_cache))
+            version = h.hexdigest()
+
+    sc = SerializationContext(
+        project=cfg.project,
+        domain=cfg.domain,
+        org=cfg.org,
+        code_bundle=code_bundle,
+        version=version,
+        image_cache=image_cache,
+        root_dir=cfg.root_dir,
+    )
+
+    deployment_coros = []
+    for env_name, env in deployment_plan.envs.items():
+        logger.info(f"Deploying environment {env_name}")
+        deployer = get_deployer(type(env))
+        context = DeploymentContext(environment=env, serialization_context=sc, dryrun=dryrun)
+        deployment_coros.append(deployer(context))
+    deployed_envs = await asyncio.gather(*deployment_coros)
+    envs = {}
+    for d in deployed_envs:
+        envs[d.get_name()] = d
+
+    return Deployment(envs)
+
+
+def _recursive_discover(planned_envs: Dict[str, Environment], env: Environment) -> Dict[str, Environment]:
+    """
+    Recursively deploy the environment and its dependencies, if not already deployed (present in env_tasks) and
+    return the updated env_tasks.
+    """
+    if env.name in planned_envs:
+        if planned_envs[env.name] is not env:
+            # Raise error if different TaskEnvironment objects have the same name
+            raise ValueError(f"Duplicate environment name '{env.name}' found")
+    # Add the environment to the existing envs
+    planned_envs[env.name] = env
+
+    # Recursively discover dependent environments
+    for dependent_env in env.depends_on:
+        _recursive_discover(planned_envs, dependent_env)
+    return planned_envs
+
+
+def plan_deploy(*envs: Environment, version: Optional[str] = None) -> List[DeploymentPlan]:
+    if envs is None:
+        return [DeploymentPlan({})]
+    deployment_plans = []
+    visited_envs: Set[str] = set()
+    for env in envs:
+        if env.name in visited_envs:
+            raise ValueError(f"Duplicate environment name '{env.name}' found")
+        planned_envs = _recursive_discover({}, env)
+        deployment_plans.append(DeploymentPlan(planned_envs, version=version))
+        visited_envs.update(planned_envs.keys())
+    return deployment_plans
+
+
+@syncify
+async def deploy(
+    *envs: Environment,
+    dryrun: bool = False,
+    version: str | None = None,
+    interactive_mode: bool | None = None,
+    copy_style: CopyFiles = "loaded_modules",
+) -> List[Deployment]:
+    """
+    Deploy the given environment or list of environments.
+    :param envs: Environment or list of environments to deploy.
+    :param dryrun: dryrun mode, if True, the deployment will not be applied to the control plane.
+    :param version: version of the deployment, if None, the version will be computed from the code bundle.
+    TODO: Support for interactive_mode
+    :param interactive_mode: Optional, can be forced to True or False.
+       If not provided, it will be set based on the current environment. For example Jupyter notebooks are considered
+         interactive mode, while scripts are not. This is used to determine how the code bundle is created.
+    :param copy_style: Copy style to use when running the task
+
+    :return: Deployment object containing the deployed environments and tasks.
+    """
+    if interactive_mode:
+        raise NotImplementedError("Interactive mode not yet implemented for deployment")
+    deployment_plans = plan_deploy(*envs, version=version)
+    deployments = []
+    for deployment_plan in deployment_plans:
+        deployments.append(apply(deployment_plan, copy_style=copy_style, dryrun=dryrun))
+    return await asyncio.gather(*deployments)
+
+
+@syncify
+async def build_images(envs: Environment) -> ImageCache:
+    """
+    Build the images for the given environments.
+    :param envs: Environment to build images for.
+    :return: ImageCache containing the built images.
+    """
+    cfg = get_init_config()
+    images = cfg.images if cfg else {}
+    deployment = plan_deploy(envs)
+    return await _build_images(deployment[0], images)

flyte/_deployer.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Protocol, Tuple, Type
+
+import rich.repr
+
+from flyte.models import SerializationContext
+
+from ._deploy import _deploy_task_env
+from ._environment import Environment
+from ._task_environment import TaskEnvironment
+
+
+@rich.repr.auto
+@dataclass
+class DeploymentContext:
+    """
+    Context for deployment operations.
+    """
+
+    environment: Environment
+    serialization_context: SerializationContext
+    dryrun: bool = False
+
+
+class DeployedEnvironment(Protocol):
+    """
+    Protocol for deployed environment representations.
+    """
+
+    def get_name(self) -> str:
+        """
+        Returns the name of the deployed environment.
+        Returns:
+        """
+        ...
+
+    def env_repr(self) -> List[Tuple[str, ...]]:
+        """
+        Returns a detailed representation of the deployed environment.
+        Returns:
+        """
+        ...
+
+    def table_repr(self) -> List[List[Tuple[str, ...]]]:
+        """
+        Returns a detailed representation of the deployed entities in the environment, useful for tabular display.
+        Returns:
+
+        """
+        ...
+
+    def summary_repr(self) -> str:
+        """
+        Returns a summary representation of the deployed environment.
+        Returns:
+        """
+        ...
+
+
+class Deployer(Protocol):
+    """
+    Protocol for deployment callables.
+    """
+
+    async def __call__(self, context: DeploymentContext) -> DeployedEnvironment:
+        """
+        Deploy the environment described in the context.
+
+        Args:
+            context: Deployment context containing environment, serialization context, and dryrun flag
+
+        Returns:
+            Deployment result
+        """
+        ...
+
+
+_ENVTYPE_REGISTRY: Dict[Type[Environment], Deployer] = {
+    TaskEnvironment: _deploy_task_env,
+}
+
+
+def register_deployer(env_type: Type[Environment], deployer: Deployer) -> None:
+    """
+    Register a deployer for a specific environment type.
+
+    Args:
+        env_type: Type of environment this deployer handles
+        deployer: Deployment callable that conforms to the Deployer protocol
+    """
+    _ENVTYPE_REGISTRY[env_type] = deployer
+
+
+def get_deployer(env_type: Type[Environment | TaskEnvironment]) -> Deployer:
+    """
+    Get the registered deployer for an environment type.
+
+    Args:
+        env_type: Type of environment to get deployer for
+
+    Returns:
+        Deployer for the environment type, defaults to task environment deployer
+    """
+    for tpe, v in _ENVTYPE_REGISTRY.items():
+        if issubclass(env_type, tpe):
+            return v
+    raise ValueError(f"No deployer registered for environment type {env_type}")

flyte/_doc.py

@@ -0,0 +1,29 @@
+import inspect
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+
+@dataclass
+class Documentation:
+    """
+    This class is used to store the documentation of a task.
+
+    It can be set explicitly or extracted from the docstring of the task.
+    """
+
+    description: str
+
+    def __help__str__(self):
+        return self.description
+
+
+def extract_docstring(func: Optional[Callable]) -> Documentation:
+    """
+    Extracts the description from a docstring.
+    """
+    if func is None:
+        return Documentation(description="")
+    docstring = inspect.getdoc(func)
+    if not docstring:
+        return Documentation(description="")
+    return Documentation(description=docstring)

flyte/_docstring.py

@@ -0,0 +1,32 @@
+from typing import TYPE_CHECKING, Callable, Dict, Optional
+
+if TYPE_CHECKING:
+    pass
+
+
+class Docstring(object):
+    def __init__(self, docstring: Optional[str] = None, callable_: Optional[Callable] = None):
+        import docstring_parser
+
+        self._parsed_docstring: docstring_parser.Docstring
+
+        if docstring is not None:
+            self._parsed_docstring = docstring_parser.parse(docstring)
+        elif callable_.__doc__ is not None:
+            self._parsed_docstring = docstring_parser.parse(callable_.__doc__)
+
+    @property
+    def input_descriptions(self) -> Dict[str, Optional[str]]:
+        return {p.arg_name: p.description for p in self._parsed_docstring.params}
+
+    @property
+    def output_descriptions(self) -> Dict[str, Optional[str]]:
+        return {p.return_name: p.description for p in self._parsed_docstring.many_returns if p.return_name is not None}
+
+    @property
+    def short_description(self) -> Optional[str]:
+        return self._parsed_docstring.short_description
+
+    @property
+    def long_description(self) -> Optional[str]:
+        return self._parsed_docstring.long_description

flyte/_environment.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import rich.repr
+
+from ._image import Image
+from ._pod import PodTemplate
+from ._resources import Resources
+from ._secret import Secret, SecretRequest
+
+# Global registry to track all Environment instances in load order
+_ENVIRONMENT_REGISTRY: List[Environment] = []
+
+
+def list_loaded_environments() -> List[Environment]:
+    """
+    Return a list of all Environment objects in the order they were loaded.
+    This is useful for deploying environments in the order they were defined.
+    """
+    return _ENVIRONMENT_REGISTRY
+
+
+def is_snake_or_kebab_with_numbers(s: str) -> bool:
+    return re.fullmatch(r"^[a-z0-9]+([_-][a-z0-9]+)*$", s) is not None
+
+
+@rich.repr.auto
+@dataclass(init=True, repr=True)
+class Environment:
+    """
+    :param name: Name of the environment
+    :param image: Docker image to use for the environment. If set to "auto", will use the default image.
+    :param resources: Resources to allocate for the environment.
+    :param env_vars: Environment variables to set for the environment.
+    :param secrets: Secrets to inject into the environment.
+    :param pod_template: Pod template to use for the environment.
+    :param description: Description of the environment.
+    :param interruptible: Whether the environment is interruptible and can be scheduled on spot/preemptible instances
+    :param depends_on: Environment dependencies to hint, so when you deploy the environment, the dependencies are
+        also deployed. This is useful when you have a set of environments that depend on each other.
+    """
+
+    name: str
+    depends_on: List[Environment] = field(default_factory=list)
+    pod_template: Optional[Union[str, PodTemplate]] = None
+    description: Optional[str] = None
+    secrets: Optional[SecretRequest] = None
+    env_vars: Optional[Dict[str, str]] = None
+    resources: Optional[Resources] = None
+    interruptible: bool = False
+    image: Union[str, Image, Literal["auto"]] = "auto"
+
+    def _validate_name(self):
+        if not is_snake_or_kebab_with_numbers(self.name):
+            raise ValueError(f"Environment name '{self.name}' must be in snake_case or kebab-case format.")
+
+    def __post_init__(self):
+        if not isinstance(self.image, (Image, str)):
+            raise TypeError(f"Expected image to be of type str or Image, got {type(self.image)}")
+        if self.secrets and not isinstance(self.secrets, (str, Secret, List)):
+            raise TypeError(f"Expected secrets to be of type SecretRequest, got {type(self.secrets)}")
+        for dep in self.depends_on:
+            if not isinstance(dep, Environment):
+                raise TypeError(f"Expected depends_on to be of type List[Environment], got {type(dep)}")
+        if self.resources is not None and not isinstance(self.resources, Resources):
+            raise TypeError(f"Expected resources to be of type Resources, got {type(self.resources)}")
+        if self.env_vars is not None and not isinstance(self.env_vars, dict):
+            raise TypeError(f"Expected env_vars to be of type Dict[str, str], got {type(self.env_vars)}")
+        if self.pod_template is not None and not isinstance(self.pod_template, (str, PodTemplate)):
+            raise TypeError(f"Expected pod_template to be of type str or PodTemplate, got {type(self.pod_template)}")
+        self._validate_name()
+        # Automatically register this environment instance in load order
+        _ENVIRONMENT_REGISTRY.append(self)
+
+    def add_dependency(self, *env: Environment):
+        """
+        Add a dependency to the environment.
+        """
+        for e in env:
+            if not isinstance(e, Environment):
+                raise TypeError(f"Expected Environment, got {type(e)}")
+            if e.name == self.name:
+                raise ValueError("Cannot add self as a dependency")
+            if e in self.depends_on:
+                continue
+        self.depends_on.extend(env)
+
+    def clone_with(
+        self,
+        name: str,
+        image: Optional[Union[str, Image, Literal["auto"]]] = None,
+        resources: Optional[Resources] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        secrets: Optional[SecretRequest] = None,
+        depends_on: Optional[List[Environment]] = None,
+        description: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Environment:
+        raise NotImplementedError
+
+    def _get_kwargs(self) -> Dict[str, Any]:
+        """
+        Get the keyword arguments for the environment.
+        """
+        kwargs: Dict[str, Any] = {
+            "depends_on": self.depends_on,
+            "image": self.image,
+        }
+        if self.resources is not None:
+            kwargs["resources"] = self.resources
+        if self.secrets is not None:
+            kwargs["secrets"] = self.secrets
+        if self.env_vars is not None:
+            kwargs["env_vars"] = self.env_vars
+        if self.pod_template is not None:
+            kwargs["pod_template"] = self.pod_template
+        if self.description is not None:
+            kwargs["description"] = self.description
+        return kwargs