flyte 0.0.1b0__py3-none-any.whl

flyte/_deploy.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+
+import rich.repr
+
+from ._api_commons import syncer
+from ._datastructures import SerializationContext
+from ._environment import Environment
+from ._image import Image
+from ._initialize import get_client, get_common_config, requires_client, requires_initialization
+from ._logging import logger
+from ._task import TaskTemplate
+from ._task_environment import TaskEnvironment
+
+if TYPE_CHECKING:
+    from flyte._protos.workflow import task_definition_pb2
+
+    from ._code_bundle import CopyFiles
+    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 Deployment:
+    envs: Dict[str, Environment]
+    deployed_tasks: List[task_definition_pb2.TaskSpec] | None = None
+
+    def summary_repr(self) -> str:
+        """
+        Returns a summary representation of the deployment.
+        """
+        env_names = ", ".join(self.envs.keys())
+        task_names_versions = ", ".join(
+            f"{task.task_template.id.name} (v{task.task_template.id.version})" for task in self.deployed_tasks or []
+        )
+        return f"Deployment(envs=[{env_names}], tasks=[{task_names_versions}])"
+
+
+@requires_client
+async def _deploy_task(
+    task: TaskTemplate, serialization_context: SerializationContext, dryrun: bool = False
+) -> task_definition_pb2.TaskSpec:
+    """
+    Deploy the given task.
+    """
+    from ._internal.runtime.task_serde import translate_task_to_wire
+    from ._protos.workflow import task_definition_pb2, task_service_pb2
+
+    image_uri = task.image.uri if isinstance(task.image, Image) else task.image
+
+    spec = translate_task_to_wire(task, serialization_context)
+    if dryrun:
+        return spec
+    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,
+    )
+
+    await get_client().task_service.DeployTask(task_service_pb2.DeployTaskRequest(task_id=task_id, spec=spec))
+    logger.info(f"Deployed task {task.name} with version {task_id.version}")
+    return spec
+
+
+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) -> ImageCache:
+    """
+    Build the images for the given deployment plan and update the environment with the built image.
+    """
+    from ._internal.imagebuild.image_builder import ImageCache
+
+    images = []
+    image_identifier_map = {}
+    for env_name, env in deployment.envs.items():
+        if not isinstance(env.image, str):
+            logger.info(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:
+            auto_image = Image.auto()
+            image_identifier_map["auto"] = auto_image.uri
+    final_images = await asyncio.gather(*images)
+
+    for env_name, image_uri in final_images:
+        logger.info(f"Built Image for environment {env_name}, image: {image_uri}")
+        env = deployment.envs[env_name]
+        if isinstance(env.image, Image):
+            image_identifier_map[env.image.identifier] = env.image.uri
+
+    return ImageCache(image_lookup=image_identifier_map)
+
+
+@requires_initialization
+async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool = False) -> Deployment:
+    from ._code_bundle import build_code_bundle
+
+    cfg = get_common_config()
+    image_cache = await build_images(deployment)
+    if copy_style == "none":
+        code_bundle = None
+        assert deployment.version is not None, "Version must be set when copy_style is none"
+    else:
+        code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
+        deployment.version = code_bundle.computed_version
+
+    sc = SerializationContext(
+        project=cfg.project,
+        domain=cfg.domain,
+        org=cfg.org,
+        code_bundle=code_bundle,
+        version=deployment.version,
+        image_cache=image_cache,
+        root_dir=cfg.root_dir,
+    )
+
+    tasks = []
+    for env_name, env in deployment.envs.items():
+        logger.info(f"Deploying environment {env_name}")
+        if isinstance(env, TaskEnvironment):
+            for task in env.tasks.values():
+                tasks.append(_deploy_task(task, dryrun=dryrun, serialization_context=sc))
+    return Deployment(envs=deployment.envs, deployed_tasks=await asyncio.gather(*tasks))
+
+
+def _recursive_discover(
+    planned_envs: Dict[str, Environment], envs: Environment | List[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 isinstance(envs, Environment):
+        envs = [envs]
+    for env in envs:
+        # Skip if the environment is already planned
+        if env.name in planned_envs:
+            continue
+        # Recursively discover dependent environments
+        for dependent_env in env.env_dep_hints:
+            _recursive_discover(planned_envs, dependent_env)
+        # Add the environment to the existing envs
+        planned_envs[env.name] = env
+    return planned_envs
+
+
+def plan_deploy(*envs: Environment, version: Optional[str] = None) -> DeploymentPlan:
+    if envs is None:
+        return DeploymentPlan({})
+    planned_envs = _recursive_discover({}, *envs)
+    return DeploymentPlan(planned_envs, version=version)
+
+
+@syncer.wrap
+async def deploy(
+    *envs: Environment,
+    dryrun: bool = False,
+    version: str | None = None,
+    interactive_mode: bool | None = None,
+    copy_style: CopyFiles = "loaded_modules",
+) -> 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 = plan_deploy(*envs, version=version)
+    return await apply(deployment, copy_style=copy_style, dryrun=dryrun)

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,43 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union
+
+import rich.repr
+
+from flyte._secret import SecretRequest
+
+from ._image import Image
+from ._resources import Resources
+
+if TYPE_CHECKING:
+    from kubernetes.client import V1PodTemplate
+
+
+@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: Environment variables to set for the environment.
+    :param secrets: Secrets to inject into the environment.
+    :param env_dep_hints: 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
+    env_dep_hints: List[Environment] = field(default_factory=list)
+    pod_template: Optional[Union[str, "V1PodTemplate"]] = None
+    description: Optional[str] = None
+    secrets: Optional[SecretRequest] = None
+    env: Optional[Dict[str, str]] = None
+    resources: Optional[Resources] = None
+    image: Union[str, Image, Literal["auto"]] = "auto"
+
+    def add_dependency(self, *env: Environment):
+        """
+        Add a dependency to the environment.
+        """
+        self.env_dep_hints.extend(env)

flyte/_group.py

@@ -0,0 +1,31 @@
+from contextlib import contextmanager
+
+from ._context import internal_ctx
+from ._datastructures import GroupData
+
+
+@contextmanager
+def group(name: str):
+    """
+    Create a new group with the given name. The method is intended to be used as a context manager.
+
+    Example:
+    ```python
+    @task
+    async def my_task():
+        ...
+        with group("my_group"):
+            t1(x,y)  # tasks in this block will be grouped under "my_group"
+        ...
+    ```
+
+    :param name: The name of the group
+    """
+    ctx = internal_ctx()
+    if ctx.data.task_context is None:
+        yield
+        return
+    tctx = ctx.data.task_context
+    new_tctx = tctx.replace(group_data=GroupData(name))
+    with ctx.replace_task_context(new_tctx):
+        yield

flyte/_hash.py

@@ -0,0 +1,23 @@
+from typing import Callable, Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class HashOnReferenceMixin(object):
+    def __hash__(self):
+        return hash(id(self))
+
+
+class HashMethod(Generic[T]):
+    """
+    Flyte-specific object used to wrap the hash function for a specific type
+    """
+
+    def __init__(self, function: Callable[[T], str]):
+        self._function = function
+
+    def calculate(self, obj: T) -> str:
+        """
+        Calculate hash for `obj`.
+        """
+        return self._function(obj)