mreged master

Signed-off-by: Kevin Su <[email protected]>

.github/workflows/build_image.yml

@@ -1,4 +1,4 @@
-name: Publish Python Package
+name: Publish Official flytekit Images
 
 on:
   workflow_dispatch:
@@ -8,18 +8,18 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
     steps:
       - uses: actions/checkout@v4
         with:
           fetch-depth: "0"
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v1
+        uses: docker/setup-qemu-action@v3
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v1
+        uses: docker/setup-buildx-action@v3
       - name: Login to GitHub Container Registry
-        uses: docker/login-action@v1
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"
@@ -54,12 +54,12 @@ jobs:
         with:
           fetch-depth: "0"
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v1
+        uses: docker/setup-qemu-action@v3
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v1
+        uses: docker/setup-buildx-action@v3
       - name: Login to GitHub Container Registry
-        uses: docker/login-action@v1
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"

.github/workflows/pythonbuild.yml

@@ -54,7 +54,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache pip
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           # This path is specific to Ubuntu
           path: ~/.cache/pip
@@ -97,7 +97,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache pip
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           # This path is specific to Ubuntu
           path: ~/.cache/pip
@@ -149,7 +149,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache pip
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           # This path is specific to Ubuntu
           path: ~/.cache/pip
@@ -186,7 +186,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache pip
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           # This path is specific to Ubuntu
           path: ~/.cache/pip
@@ -230,7 +230,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache pip
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           # This path is specific to Ubuntu
           path: ~/.cache/pip
@@ -274,6 +274,10 @@ jobs:
           AWS_SECRET_ACCESS_KEY: miniostorage
         run: |
           make ${{ matrix.makefile-cmd }}
+      - name: Setup tmate session
+        if: ${{ failure() }}
+        uses: mxschmitt/action-tmate@v3
+        timeout-minutes: 60
       - name: Codecov
         uses: codecov/[email protected]
         with:
@@ -394,7 +398,7 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
       - name: Cache pip
-        uses: actions/cache@v3
+        uses: actions/cache@v4
         with:
           # This path is specific to Ubuntu
           path: ~/.cache/pip
@@ -439,7 +443,7 @@ jobs:
         uses: actions/setup-python@v4
         with:
           python-version: 3.12
-      - uses: actions/cache@v3
+      - uses: actions/cache@v4
         with:
           path: ~/.cache/pip
           key: ${{ runner.os }}-pip-${{ hashFiles('**/dev-requirements.in') }}

.github/workflows/pythonpublish.yml

@@ -1,4 +1,4 @@
-name: Publish Python Package
+name: Publish Python Packages and Official Images
 
 on:
   release:
@@ -88,13 +88,13 @@ jobs:
         with:
           fetch-depth: "0"
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v1
+        uses: docker/setup-qemu-action@v3
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v1
+        uses: docker/setup-buildx-action@v3
       - name: Login to GitHub Container Registry
         if: ${{ github.event_name == 'release' }}
-        uses: docker/login-action@v1
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"
@@ -178,13 +178,13 @@ jobs:
         with:
           fetch-depth: "0"
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v1
+        uses: docker/setup-qemu-action@v3
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v1
+        uses: docker/setup-buildx-action@v3
       - name: Login to GitHub Container Registry
         if: ${{ github.event_name == 'release' }}
-        uses: docker/login-action@v1
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"
@@ -247,13 +247,13 @@ jobs:
         with:
           fetch-depth: "0"
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v1
+        uses: docker/setup-qemu-action@v3
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v1
+        uses: docker/setup-buildx-action@v3
       - name: Login to GitHub Container Registry
         if: ${{ github.event_name == 'release' }}
-        uses: docker/login-action@v1
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"
@@ -289,13 +289,13 @@ jobs:
         with:
           fetch-depth: "0"
       - name: Set up QEMU
-        uses: docker/setup-qemu-action@v1
+        uses: docker/setup-qemu-action@v3
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v1
+        uses: docker/setup-buildx-action@v3
       - name: Login to GitHub Container Registry
         if: ${{ github.event_name == 'release' }}
-        uses: docker/login-action@v1
+        uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"

.gitignore

@@ -39,3 +39,4 @@ coverage.xml
 
 # Version file is auto-generated by setuptools_scm
 flytekit/_version.py
+testing

flytekit/__init__.py

@@ -178,6 +178,9 @@
    :toctree: generated/
 
    HashMethod
+   Cache
+   CachePolicy
+   VersionParameters
 
 Artifacts
 =========
@@ -223,6 +226,7 @@
 from flytekit.core.artifact import Artifact
 from flytekit.core.base_sql_task import SQLTask
 from flytekit.core.base_task import SecurityContext, TaskMetadata, kwtypes
+from flytekit.core.cache import Cache, CachePolicy, VersionParameters
 from flytekit.core.checkpointer import Checkpoint
 from flytekit.core.condition import conditional
 from flytekit.core.container_task import ContainerTask

flytekit/clis/sdk_in_container/init.py

@@ -57,6 +57,4 @@ def init(template, project_name):
                 processed_contents = project_template_regex.sub(project_name_bytes, zip_contents)
                 dest_file.write(processed_contents)
 
-    click.echo(
-        f"Visit the {project_name} directory and follow the next steps in the Getting started guide (https://docs.flyte.org/en/latest/user_guide/getting_started_with_workflow_development/index.html) to proceed."
-    )
+    click.echo(f"Project initialized in directory {project_name}.")

flytekit/clis/sdk_in_container/run.py

@@ -1052,9 +1052,25 @@ def _create_command(
         r = run_level_params.remote_instance()
         flyte_ctx = r.context
 
+        final_inputs_with_defaults = loaded_entity.python_interface.inputs_with_defaults
+        if isinstance(loaded_entity, LaunchPlan):
+            # For LaunchPlans it is essential to handle fixed inputs and default inputs in a special way
+            # Fixed inputs are inputs that are always passed to the launch plan and cannot be overridden
+            # Default inputs are inputs that are optional and have a default value
+            # The final inputs to the launch plan are a combination of the fixed inputs and the default inputs
+            all_inputs = loaded_entity.python_interface.inputs_with_defaults
+            default_inputs = loaded_entity.saved_inputs
+            pmap = loaded_entity.parameters
+            final_inputs_with_defaults = {}
+            for name, _ in pmap.parameters.items():
+                _type, v = all_inputs[name]
+                if name in default_inputs:
+                    v = default_inputs[name]
+                final_inputs_with_defaults[name] = _type, v
+
         # Add options for each of the workflow inputs
         params = []
-        for input_name, input_type_val in loaded_entity.python_interface.inputs_with_defaults.items():
+        for input_name, input_type_val in final_inputs_with_defaults.items():
             literal_var = loaded_entity.interface.inputs.get(input_name)
             python_type, default_val = input_type_val
             required = type(None) not in get_args(python_type) and default_val is None

flytekit/configuration/plugin.py

@@ -19,11 +19,12 @@
 """
 
 import os
-from typing import Optional, Protocol, runtime_checkable
+from typing import List, Optional, Protocol, runtime_checkable
 
 from click import Group
 from importlib_metadata import entry_points
 
+from flytekit import CachePolicy
 from flytekit.configuration import Config, get_config_file
 from flytekit.loggers import logger
 from flytekit.remote import FlyteRemote
@@ -53,6 +54,10 @@ def get_default_image() -> Optional[str]:
     def get_auth_success_html(endpoint: str) -> Optional[str]:
         """Get default success html for auth. Return None to use flytekit's default success html."""
 
+    @staticmethod
+    def get_default_cache_policies() -> List[CachePolicy]:
+        """Get default cache policies for tasks."""
+
 
 class FlytekitPlugin:
     @staticmethod
@@ -103,6 +108,11 @@ def get_auth_success_html(endpoint: str) -> Optional[str]:
         """Get default success html. Return None to use flytekit's default success html."""
         return None
 
+    @staticmethod
+    def get_default_cache_policies() -> List[CachePolicy]:
+        """Get default cache policies for tasks."""
+        return []
+
 
 def _get_plugin_from_entrypoint():
     """Get plugin from entrypoint."""

flytekit/core/cache.py

@@ -0,0 +1,97 @@
+import hashlib
+from dataclasses import dataclass
+from typing import Callable, Generic, List, Optional, Protocol, Tuple, Union, runtime_checkable
+
+from typing_extensions import ParamSpec, TypeVar
+
+from flytekit.core.pod_template import PodTemplate
+from flytekit.image_spec.image_spec import ImageSpec
+
+P = ParamSpec("P")
+FuncOut = TypeVar("FuncOut")
+
+
+@dataclass
+class VersionParameters(Generic[P, FuncOut]):
+    """
+    Parameters used for version hash generation.
+
+    param func: The function to generate a version for. This is an optional parameter and can be any callable
+                 that matches the specified parameter and return types.
+    :type func: Optional[Callable[P, FuncOut]]
+    :param container_image: The container image to generate a version for. This can be a string representing the
+                            image name or an ImageSpec object.
+    :type container_image: Optional[Union[str, ImageSpec]]
+    """
+
+    func: Callable[P, FuncOut]
+    container_image: Optional[Union[str, ImageSpec]] = None
+    pod_template: Optional[PodTemplate] = None
+    pod_template_name: Optional[str] = None
+
+
+@runtime_checkable
+class CachePolicy(Protocol):
+    def get_version(self, salt: str, params: VersionParameters) -> str: ...
+
+
+@dataclass
+class Cache:
+    """
+    Cache configuration for a task.
+
+    :param version: The version of the task. If not provided, the version will be generated based on the cache policies.
+    :type version: Optional[str]
+    :param serialize: Boolean that indicates if identical (ie. same inputs) instances of this task should be executed in
+          serial when caching is enabled. This means that given multiple concurrent executions over identical inputs,
+          only a single instance executes and the rest wait to reuse the cached results.
+    :type serialize: bool
+    :param ignored_inputs: A tuple of input names to ignore when generating the version hash.
+    :type ignored_inputs: Union[Tuple[str, ...], str]
+    :param salt: A salt used in the hash generation.
+    :type salt: str
+    :param policies: A list of cache policies to generate the version hash.
+    :type policies: Optional[Union[List[CachePolicy], CachePolicy]]
+    """
+
+    version: Optional[str] = None
+    serialize: bool = False
+    ignored_inputs: Union[Tuple[str, ...], str] = ()
+    salt: str = ""
+    policies: Optional[Union[List[CachePolicy], CachePolicy]] = None
+
+    def __post_init__(self):
+        if isinstance(self.ignored_inputs, str):
+            self._ignored_inputs = (self.ignored_inputs,)
+        else:
+            self._ignored_inputs = self.ignored_inputs
+
+        # Normalize policies so that self._policies is always a list
+        if self.policies is None:
+            from flytekit.configuration.plugin import get_plugin
+
+            self._policies = get_plugin().get_default_cache_policies()
+        elif isinstance(self.policies, CachePolicy):
+            self._policies = [self.policies]
+
+        if self.version is None and not self._policies:
+            raise ValueError("If version is not defined then at least one cache policy needs to be set")
+
+    def get_ignored_inputs(self) -> Tuple[str, ...]:
+        return self._ignored_inputs
+
+    def get_version(self, params: VersionParameters) -> str:
+        if self.version is not None:
+            return self.version
+
+        task_hash = ""
+        for policy in self._policies:
+            try:
+                task_hash += policy.get_version(self.salt, params)
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to generate version for cache policy {policy}. Please consider setting the version in the Cache definition, e.g. Cache(version='v1.2.3')"
+                ) from e
+
+        hash_obj = hashlib.sha256(task_hash.encode())
+        return hash_obj.hexdigest()

flytekit/core/constants.py

@@ -38,3 +38,7 @@
 CACHE_KEY_METADATA = "cache-key-metadata"
 
 SERIALIZATION_FORMAT = "serialization-format"
+
+# Shared memory mount name and path
+SHARED_MEMORY_MOUNT_NAME = "flyte-shared-memory"
+SHARED_MEMORY_MOUNT_PATH = "/dev/shm"