Pyflyte build imageSpec (flyteorg#1555)

* wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * pyflyte build Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * Support serialize and package Signed-off-by: Kevin Su <[email protected]> * more tests Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * move to plugin Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * wip Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * fixed tested Signed-off-by: Kevin Su <[email protected]> * fixed tested Signed-off-by: Kevin Su <[email protected]> * more tests Signed-off-by: Kevin Su <[email protected]> * Add support passing yaml in pyflyte run Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]>
Yicheng-Lu-llll · Apr 24, 2023 · 4064953 · 4064953
1 parent dfb8697
commit 4064953
Show file tree

Hide file tree

Showing 29 changed files with 840 additions and 11 deletions.
diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml
@@ -66,6 +66,7 @@ jobs:
           - flytekit-deck-standard
           - flytekit-dolt
           - flytekit-duckdb
+          - flytekit-envd
           - flytekit-greatexpectations
           - flytekit-hive
           - flytekit-k8s-pod

diff --git a/Dockerfile b/Dockerfile
@@ -16,6 +16,7 @@ RUN apt-get update && apt-get install build-essential -y
 RUN pip install -U flytekit==$VERSION \
 	flytekitplugins-pod==$VERSION \
 	flytekitplugins-deck-standard==$VERSION \
+	flytekitplugins-envd==$VERSION \
 	scikit-learn
 
 RUN useradd -u 1000 flytekit

diff --git a/flytekit/clis/sdk_in_container/build.py b/flytekit/clis/sdk_in_container/build.py
@@ -0,0 +1,127 @@
+import os
+import pathlib
+import typing
+
+import click
+from typing_extensions import OrderedDict
+
+from flytekit.clis.sdk_in_container.constants import CTX_MODULE, CTX_PROJECT_ROOT
+from flytekit.clis.sdk_in_container.run import RUN_LEVEL_PARAMS_KEY, get_entities_in_file, load_naive_entity
+from flytekit.configuration import ImageConfig, SerializationSettings
+from flytekit.core.base_task import PythonTask
+from flytekit.core.workflow import PythonFunctionWorkflow
+from flytekit.tools.script_mode import _find_project_root
+from flytekit.tools.translator import get_serializable
+
+
+def get_workflow_command_base_params() -> typing.List[click.Option]:
+    """
+    Return the set of base parameters added to every pyflyte build workflow subcommand.
+    """
+    return [
+        click.Option(
+            param_decls=["--fast"],
+            required=False,
+            is_flag=True,
+            default=False,
+            help="Use fast serialization. The image won't contain the source code. The value is false by default.",
+        ),
+    ]
+
+
+def build_command(ctx: click.Context, entity: typing.Union[PythonFunctionWorkflow, PythonTask]):
+    """
+    Returns a function that is used to implement WorkflowCommand and build an image for flyte workflows.
+    """
+
+    def _build(*args, **kwargs):
+        m = OrderedDict()
+        options = None
+        run_level_params = ctx.obj[RUN_LEVEL_PARAMS_KEY]
+
+        project, domain = run_level_params.get("project"), run_level_params.get("domain")
+        serialization_settings = SerializationSettings(
+            project=project,
+            domain=domain,
+            image_config=ImageConfig.auto_default_image(),
+        )
+        if not run_level_params.get("fast"):
+            serialization_settings.source_root = ctx.obj[RUN_LEVEL_PARAMS_KEY].get(CTX_PROJECT_ROOT)
+
+        _ = get_serializable(m, settings=serialization_settings, entity=entity, options=options)
+
+    return _build
+
+
+class WorkflowCommand(click.MultiCommand):
+    """
+    click multicommand at the python file layer, subcommands should be all the workflows in the file.
+    """
+
+    def __init__(self, filename: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._filename = pathlib.Path(filename).resolve()
+
+    def list_commands(self, ctx):
+        entities = get_entities_in_file(self._filename.__str__())
+        return entities.all()
+
+    def get_command(self, ctx, exe_entity):
+        """
+        This command uses the filename with which this command was created, and the string name of the entity passed
+          after the Python filename on the command line, to load the Python object, and then return the Command that
+          click should run.
+        :param ctx: The click Context object.
+        :param exe_entity: string of the flyte entity provided by the user. Should be the name of a workflow, or task
+          function.
+        :return:
+        """
+        rel_path = os.path.relpath(self._filename)
+        if rel_path.startswith(".."):
+            raise ValueError(
+                f"You must call pyflyte from the same or parent dir, {self._filename} not under {os.getcwd()}"
+            )
+
+        project_root = _find_project_root(self._filename)
+        rel_path = self._filename.relative_to(project_root)
+        module = os.path.splitext(rel_path)[0].replace(os.path.sep, ".")
+
+        ctx.obj[RUN_LEVEL_PARAMS_KEY][CTX_PROJECT_ROOT] = project_root
+        ctx.obj[RUN_LEVEL_PARAMS_KEY][CTX_MODULE] = module
+
+        entity = load_naive_entity(module, exe_entity, project_root)
+
+        cmd = click.Command(
+            name=exe_entity,
+            callback=build_command(ctx, entity),
+            help=f"Build an image for {module}.{exe_entity}.",
+        )
+        return cmd
+
+
+class BuildCommand(click.MultiCommand):
+    """
+    A click command group for building a image for flyte workflows & tasks in a file.
+    """
+
+    def __init__(self, *args, **kwargs):
+        params = get_workflow_command_base_params()
+        super().__init__(*args, params=params, **kwargs)
+
+    def list_commands(self, ctx):
+        return [str(p) for p in pathlib.Path(".").glob("*.py") if str(p) != "__init__.py"]
+
+    def get_command(self, ctx, filename):
+        if ctx.obj:
+            ctx.obj[RUN_LEVEL_PARAMS_KEY] = ctx.params
+        return WorkflowCommand(filename, name=filename, help="Build an image for [workflow|task]")
+
+
+_build_help = """
+This command can build an image for a workflow or a task from the command line, for fully self-contained scripts.
+"""
+
+build = BuildCommand(
+    name="build",
+    help=_build_help,
+)
diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py
@@ -6,6 +6,7 @@
 
 from flytekit import configuration
 from flytekit.clis.sdk_in_container.backfill import backfill
+from flytekit.clis.sdk_in_container.build import build
 from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE, CTX_PACKAGES, CTX_VERBOSE
 from flytekit.clis.sdk_in_container.init import init
 from flytekit.clis.sdk_in_container.local_cache import local_cache
@@ -132,6 +133,7 @@ def main(ctx, pkgs: typing.List[str], config: str, verbose: bool):
 main.add_command(run)
 main.add_command(register)
 main.add_command(backfill)
+main.add_command(build)
 main.epilog
 
 if __name__ == "__main__":

diff --git a/flytekit/configuration/__init__.py b/flytekit/configuration/__init__.py
@@ -143,12 +143,15 @@
 from io import BytesIO
 from typing import Dict, List, Optional
 
+import yaml
 from dataclasses_json import dataclass_json
 from docker_image import reference
 
 from flytekit.configuration import internal as _internal
 from flytekit.configuration.default_images import DefaultImages
 from flytekit.configuration.file import ConfigEntry, ConfigFile, get_config_file, read_file_if_exists, set_if_exists
+from flytekit.image_spec import ImageSpec
+from flytekit.image_spec.image_spec import ImageBuildEngine
 from flytekit.loggers import logger
 
 PROJECT_PLACEHOLDER = "{{ registration.project }}"
@@ -205,6 +208,13 @@ def look_up_image_info(name: str, tag: str, optional_tag: bool = False) -> Image
         :param Text tag: e.g. somedocker.com/myimage:someversion123
         :rtype: Text
         """
+        if pathlib.Path(tag).is_file():
+            with open(tag, "r") as f:
+                image_spec_dict = yaml.safe_load(f)
+                image_spec = ImageSpec(**image_spec_dict)
+                ImageBuildEngine.build(image_spec)
+                tag = image_spec.image_name()
+
         ref = reference.Reference.parse(tag)
         if not optional_tag and ref["tag"] is None:
             raise AssertionError(f"Incorrectly formatted image {tag}, missing tag value")
@@ -699,6 +709,7 @@ class SerializationSettings(object):
         fast_serialization_settings (Optional[FastSerializationSettings]): If the code is being serialized so that it
             can be fast registered (and thus omit building a Docker image) this object contains additional parameters
             for serialization.
+        source_root (Optional[str]): The root directory of the source code.
     """
 
     image_config: ImageConfig
@@ -710,6 +721,7 @@ class SerializationSettings(object):
     python_interpreter: str = DEFAULT_RUNTIME_PYTHON_INTERPRETER
     flytekit_virtualenv_root: Optional[str] = None
     fast_serialization_settings: Optional[FastSerializationSettings] = None
+    source_root: Optional[str] = None
 
     def __post_init__(self):
         if self.flytekit_virtualenv_root is None:

diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py
@@ -3,7 +3,7 @@
 import importlib
 import re
 from abc import ABC
-from typing import Callable, Dict, List, Optional, TypeVar
+from typing import Callable, Dict, List, Optional, TypeVar, Union
 
 from flytekit.configuration import ImageConfig, SerializationSettings
 from flytekit.core.base_task import PythonTask, TaskMetadata, TaskResolverMixin
@@ -13,6 +13,7 @@
 from flytekit.core.tracked_abc import FlyteTrackedABC
 from flytekit.core.tracker import TrackedInstance, extract_task_module
 from flytekit.core.utils import _get_container_definition, _serialize_pod_spec, timeit
+from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec
 from flytekit.loggers import logger
 from flytekit.models import task as _task_model
 from flytekit.models.security import Secret, SecurityContext
@@ -35,7 +36,7 @@ def __init__(
         name: str,
         task_config: T,
         task_type="python-task",
-        container_image: Optional[str] = None,
+        container_image: Optional[Union[str, ImageSpec]] = None,
         requests: Optional[Resources] = None,
         limits: Optional[Resources] = None,
         environment: Optional[Dict[str, str]] = None,
@@ -77,7 +78,7 @@ def __init__(
                     raise AssertionError(f"Secret {s} should be of type flytekit.Secret, received {type(s)}")
             sec_ctx = SecurityContext(secrets=secret_requests)
 
-        # pod_template_name overwrites the metedata.pod_template_name
+        # pod_template_name overwrites the metadata.pod_template_name
         kwargs["metadata"] = kwargs["metadata"] if "metadata" in kwargs else TaskMetadata()
         kwargs["metadata"].pod_template_name = pod_template_name
 
@@ -115,7 +116,7 @@ def task_resolver(self) -> TaskResolverMixin:
         return self._task_resolver
 
     @property
-    def container_image(self) -> Optional[str]:
+    def container_image(self) -> Optional[Union[str, ImageSpec]]:
         return self._container_image
 
     @property
@@ -180,6 +181,8 @@ def _get_container(self, settings: SerializationSettings) -> _task_model.Contain
         for elem in (settings.env, self.environment):
             if elem:
                 env.update(elem)
+        if isinstance(self.container_image, ImageSpec):
+            self.container_image.source_root = settings.source_root
         return _get_container_definition(
             image=get_registerable_container_image(self.container_image, settings.image_config),
             command=[],
@@ -249,12 +252,16 @@ def get_all_tasks(self) -> List[PythonAutoContainerTask]:  # type: ignore
 default_task_resolver = DefaultTaskResolver()
 
 
-def get_registerable_container_image(img: Optional[str], cfg: ImageConfig) -> str:
+def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: ImageConfig) -> str:
     """
-    :param img: Configured image
+    :param img: Configured image or image spec
     :param cfg: Registration configuration
     :return:
     """
+    if isinstance(img, ImageSpec):
+        ImageBuildEngine.build(img)
+        return img.image_name()
+
     if img is not None and img != "":
         matches = _IMAGE_REPLACE_REGEX.findall(img)
         if matches is None or len(matches) == 0:

diff --git a/flytekit/core/task.py b/flytekit/core/task.py
@@ -8,6 +8,7 @@
 from flytekit.core.python_function_task import PythonFunctionTask
 from flytekit.core.reference_entity import ReferenceEntity, TaskReference
 from flytekit.core.resources import Resources
+from flytekit.image_spec.image_spec import ImageSpec
 from flytekit.models.documentation import Documentation
 from flytekit.models.security import Secret
 
@@ -84,7 +85,7 @@ def task(
     interruptible: Optional[bool] = None,
     deprecated: str = "",
     timeout: Union[_datetime.timedelta, int] = 0,
-    container_image: Optional[str] = None,
+    container_image: Optional[Union[str, ImageSpec]] = None,
     environment: Optional[Dict[str, str]] = None,
     requests: Optional[Resources] = None,
     limits: Optional[Resources] = None,

diff --git a/flytekit/image_spec/__init__.py b/flytekit/image_spec/__init__.py
@@ -0,0 +1 @@
+from .image_spec import ImageSpec