merge master

Signed-off-by: Yee Hing Tong <[email protected]>

.github/workflows/pythonbuild.yml

@@ -316,6 +316,7 @@ jobs:
           - flytekit-aws-batch
           - flytekit-aws-sagemaker
           - flytekit-bigquery
+          - flytekit-comet-ml
           - flytekit-dask
           - flytekit-data-fsspec
           - flytekit-dbt

.gitignore

@@ -35,6 +35,7 @@ docs/source/_tags/
 .hypothesis
 .npm
 /**/target
+coverage.xml
 
 # Version file is auto-generated by setuptools_scm
 flytekit/_version.py

Dockerfile

@@ -23,7 +23,7 @@ ARG DOCKER_IMAGE
 RUN apt-get update && apt-get install build-essential -y \
     && pip install uv \
     && uv pip install --system --no-cache-dir -U flytekit==$VERSION \
-        flytekitplugins-deck-standard==$VERSION \
+        kubernetes \
     && apt-get clean autoclean \
     && apt-get autoremove --yes \
     && rm -rf /var/lib/{apt,dpkg,cache,log}/ \

dev-requirements.txt

@@ -44,7 +44,7 @@ botocore==1.34.106
     # via aiobotocore
 cachetools==5.3.3
     # via google-auth
-certifi==2024.2.2
+certifi==2024.7.4
     # via
     #   kubernetes
     #   requests
@@ -76,6 +76,7 @@ cryptography==42.0.7
     #   azure-storage-blob
     #   msal
     #   pyjwt
+    #   secretstorage
 dataclasses-json==0.5.9
     # via flytekit
 decorator==5.1.1
@@ -181,8 +182,7 @@ iniconfig==2.0.0
 ipython==8.25.0
     # via -r dev-requirements.in
 isodate==0.6.1
-    # via
-    #   azure-storage-blob
+    # via azure-storage-blob
 jaraco-classes==3.4.0
     # via
     #   keyring
@@ -195,6 +195,10 @@ jaraco-functools==4.0.1
     # via keyring
 jedi==0.19.1
     # via ipython
+jeepney==0.8.0
+    # via
+    #   keyring
+    #   secretstorage
 jmespath==1.0.1
     # via botocore
 joblib==1.4.2
@@ -307,6 +311,7 @@ proto-plus==1.23.0
     #   google-cloud-bigquery-storage
 protobuf==4.25.3
     # via
+    #   -r dev-requirements.in
     #   flyteidl
     #   flytekit
     #   google-api-core
@@ -416,6 +421,8 @@ scikit-learn==1.5.0
     # via -r dev-requirements.in
 scipy==1.13.1
     # via scikit-learn
+secretstorage==3.3.3
+    # via keyring
 setuptools-scm==8.1.0
     # via -r dev-requirements.in
 six==1.16.0
@@ -443,7 +450,7 @@ types-decorator==5.1.8.20240310
     # via -r dev-requirements.in
 types-mock==5.1.0.20240425
     # via -r dev-requirements.in
-types-protobuf==5.26.0.20240422
+types-protobuf==4.25.0.20240417
     # via -r dev-requirements.in
 types-requests==2.32.0.20240523
     # via -r dev-requirements.in

docs/source/_templates/file_types.rst

@@ -2,7 +2,7 @@
 
 .. currentmodule:: {{ module }}
 
-{% if objname == 'FlyteFile' %}
+{% if objname == 'FlyteFile' or  objname == 'FlyteDirectory' %}
 
 .. autoclass:: {{ objname }}
 

docs/source/plugins/index.rst

@@ -32,6 +32,7 @@ Plugin API reference
 * :ref:`DuckDB <duckdb>` - DuckDB API reference
 * :ref:`SageMaker Inference <awssagemaker_inference>` - SageMaker Inference API reference
 * :ref:`OpenAI <openai>` - OpenAI API reference
+* :ref:`Inference <inference>` - Inference API reference
 
 .. toctree::
    :maxdepth: 2
@@ -65,3 +66,4 @@ Plugin API reference
    DuckDB <duckdb>
    SageMaker Inference <awssagemaker_inference>
    OpenAI <openai>
+   Inference <inference>

docs/source/plugins/inference.rst

@@ -0,0 +1,12 @@
+.. _inference:
+
+#########################
+Model Inference reference
+#########################
+
+.. tags:: Integration, Serving, Inference
+
+.. automodule:: flytekitplugins.inference
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:

flytekit/_ast/parser.py

@@ -0,0 +1,25 @@
+import ast
+import inspect
+import typing
+
+
+def get_function_param_location(func: typing.Callable, param_name: str) -> (int, int):
+    """
+    Get the line and column number of the parameter in the source code of the function definition.
+    """
+    # Get source code of the function
+    source_lines, start_line = inspect.getsourcelines(func)
+    source_code = "".join(source_lines)
+
+    # Parse the source code into an AST
+    module = ast.parse(source_code)
+
+    # Traverse the AST to find the function definition
+    for node in ast.walk(module):
+        if isinstance(node, ast.FunctionDef) and node.name == func.__name__:
+            for i, arg in enumerate(node.args.args):
+                if arg.arg == param_name:
+                    # Calculate the line and column number of the parameter
+                    line_number = start_line + node.lineno - 1
+                    column_offset = arg.col_offset
+                    return line_number, column_offset

flytekit/clients/auth/authenticator.py

@@ -179,12 +179,15 @@ def refresh_credentials(self):
         This function is used when the configuration value for AUTH_MODE is set to 'external_process'.
         It reads an id token generated by an external process started by running the 'command'.
         """
-        logging.debug("Starting external process to generate id token. Command {}".format(self._cmd))
+        cmd_joined = " ".join(self._cmd)
+        logging.debug("Starting external process to generate id token. Command `{}`".format(" ".join(cmd_joined)))
         try:
             output = subprocess.run(self._cmd, capture_output=True, text=True, check=True)
-        except subprocess.CalledProcessError as e:
-            logging.error("Failed to generate token from command {}".format(self._cmd))
-            raise AuthenticationError("Problems refreshing token with command: " + str(e))
+        except subprocess.CalledProcessError:
+            logging.error("Failed to generate token from command `{}`".format(cmd_joined))
+            raise AuthenticationError(
+                f"Failed to refresh token with command `{cmd_joined}`. Please execute this command in your terminal to debug."
+            )
         self._creds = Credentials(output.stdout.strip())
 
 

flytekit/clients/grpc_utils/auth_interceptor.py

@@ -61,6 +61,8 @@ def intercept_unary_unary(
         fut: grpc.Future = continuation(updated_call_details, request)
         e = fut.exception()
         if e:
+            if not hasattr(e, "code"):
+                raise e
             if e.code() == grpc.StatusCode.UNAUTHENTICATED or e.code() == grpc.StatusCode.UNKNOWN:
                 self._authenticator.refresh_credentials()
                 updated_call_details = self._call_details_with_auth_metadata(client_call_details)

flytekit/clis/sdk_in_container/run.py

@@ -7,11 +7,12 @@
 import tempfile
 import typing
 from dataclasses import dataclass, field, fields
-from typing import cast, get_args
+from typing import Iterator, get_args
 
 import rich_click as click
-from dataclasses_json import DataClassJsonMixin
+from mashumaro.codecs.json import JSONEncoder
 from rich.progress import Progress
+from typing_extensions import get_origin
 
 from flytekit import Annotations, FlyteContext, FlyteContextManager, Labels, Literal
 from flytekit.clis.sdk_in_container.helpers import patch_image_config
@@ -395,7 +396,8 @@ def to_click_option(
             if type(default_val) == dict or type(default_val) == list:
                 default_val = json.dumps(default_val)
             else:
-                default_val = cast(DataClassJsonMixin, default_val).to_json()
+                encoder = JSONEncoder(python_type)
+                default_val = encoder.encode(default_val)
         if literal_var.type.metadata:
             description_extra = f": {json.dumps(literal_var.type.metadata)}"
 
@@ -537,10 +539,21 @@ def _run(*args, **kwargs):
             for input_name, v in entity.python_interface.inputs_with_defaults.items():
                 processed_click_value = kwargs.get(input_name)
                 optional_v = False
+
+                skip_default_value_selection = False
                 if processed_click_value is None and isinstance(v, typing.Tuple):
-                    optional_v = is_optional(v[0])
-                    if len(v) == 2:
-                        processed_click_value = v[1]
+                    if entity_type == "workflow" and hasattr(v[0], "__args__"):
+                        origin_base_type = get_origin(v[0])
+                        if inspect.isclass(origin_base_type) and issubclass(origin_base_type, Iterator):  # Iterator
+                            args = getattr(v[0], "__args__")
+                            if isinstance(args, tuple) and get_origin(args[0]) is typing.Union:  # Iterator[JSON]
+                                logger.debug(f"Detected Iterator[JSON] in {entity.name} input annotations...")
+                                skip_default_value_selection = True
+
+                    if not skip_default_value_selection:
+                        optional_v = is_optional(v[0])
+                        if len(v) == 2:
+                            processed_click_value = v[1]
                 if isinstance(processed_click_value, ArtifactQuery):
                     if run_level_params.is_remote:
                         click.secho(

flytekit/clis/sdk_in_container/utils.py

@@ -13,7 +13,8 @@
 
 from flytekit.core.constants import SOURCE_CODE
 from flytekit.exceptions.base import FlyteException
-from flytekit.exceptions.user import FlyteInvalidInputException
+from flytekit.exceptions.user import FlyteCompilationException, FlyteInvalidInputException
+from flytekit.exceptions.utils import annotate_exception_with_code
 from flytekit.loggers import get_level_from_cli_verbosity, logger
 
 project_option = click.Option(
@@ -130,12 +131,14 @@ def pretty_print_traceback(e: Exception, verbosity: int = 1):
     else:
         raise ValueError(f"Verbosity level must be between 0 and 2. Got {verbosity}")
 
-    if hasattr(e, SOURCE_CODE):
-        # TODO: Use other way to check if the background is light or dark
-        theme = "emacs" if "LIGHT_BACKGROUND" in os.environ else "monokai"
-        syntax = Syntax(getattr(e, SOURCE_CODE), "python", theme=theme, background_color="default")
-        panel = Panel(syntax, border_style="red", title=type(e).__name__, title_align="left")
-        console.print(panel, no_wrap=False)
+    if isinstance(e, FlyteCompilationException):
+        e = annotate_exception_with_code(e, e.fn, e.param_name)
+        if hasattr(e, SOURCE_CODE):
+            # TODO: Use other way to check if the background is light or dark
+            theme = "emacs" if "LIGHT_BACKGROUND" in os.environ else "monokai"
+            syntax = Syntax(getattr(e, SOURCE_CODE), "python", theme=theme, background_color="default")
+            panel = Panel(syntax, border_style="red", title=e._ERROR_CODE, title_align="left")
+            console.print(panel, no_wrap=False)
 
 
 def pretty_print_exception(e: Exception, verbosity: int = 1):
@@ -161,20 +164,14 @@ def pretty_print_exception(e: Exception, verbosity: int = 1):
                 pretty_print_grpc_error(cause)
             else:
                 pretty_print_traceback(e, verbosity)
+        else:
+            pretty_print_traceback(e, verbosity)
         return
 
     if isinstance(e, grpc.RpcError):
         pretty_print_grpc_error(e)
         return
 
-    if isinstance(e, AssertionError):
-        click.secho(f"Assertion Error: {e}", fg="red")
-        return
-
-    if isinstance(e, ValueError):
-        click.secho(f"Value Error: {e}", fg="red")
-        return
-
     pretty_print_traceback(e, verbosity)
 
 

flytekit/core/array_node_map_task.py

@@ -8,6 +8,8 @@
 from contextlib import contextmanager
 from typing import Any, Dict, List, Optional, Set, Union, cast
 
+from flyteidl.core import tasks_pb2
+
 from flytekit.configuration import SerializationSettings
 from flytekit.core import tracker
 from flytekit.core.base_task import PythonTask, TaskResolverMixin
@@ -152,6 +154,9 @@ def python_function_task(self) -> Union[PythonFunctionTask, PythonInstanceTask]:
     def bound_inputs(self) -> Set[str]:
         return self._bound_inputs
 
+    def get_extended_resources(self, settings: SerializationSettings) -> Optional[tasks_pb2.ExtendedResources]:
+        return self.python_function_task.get_extended_resources(settings)
+
     @contextmanager
     def prepare_target(self):
         """

flytekit/core/interface.py

@@ -12,11 +12,16 @@
 
 from flytekit.core import context_manager
 from flytekit.core.artifact import Artifact, ArtifactIDSpecification, ArtifactQuery
+from flytekit.core.context_manager import FlyteContextManager
 from flytekit.core.docstring import Docstring
 from flytekit.core.sentinel import DYNAMIC_INPUT_BINDING
 from flytekit.core.type_engine import TypeEngine, UnionTransformer
-from flytekit.exceptions.user import FlyteValidationException
-from flytekit.exceptions.utils import annotate_exception_with_code
+from flytekit.core.utils import has_return_statement
+from flytekit.exceptions.user import (
+    FlyteMissingReturnValueException,
+    FlyteMissingTypeException,
+    FlyteValidationException,
+)
 from flytekit.loggers import developer_logger, logger
 from flytekit.models import interface as _interface_models
 from flytekit.models.literals import Literal, Scalar, Void
@@ -375,15 +380,26 @@ def transform_function_to_interface(fn: typing.Callable, docstring: Optional[Doc
     signature = inspect.signature(fn)
     return_annotation = type_hints.get("return", None)
 
+    ctx = FlyteContextManager.current_context()
+    # Only check if the task/workflow has a return statement at compile time locally.
+    if (
+        ctx.execution_state
+        and ctx.execution_state.mode is None
+        and return_annotation
+        and type(None) not in get_args(return_annotation)
+        and return_annotation is not type(None)
+        and has_return_statement(fn) is False
+    ):
+        raise FlyteMissingReturnValueException(fn=fn)
+
     outputs = extract_return_annotation(return_annotation)
     for k, v in outputs.items():
         outputs[k] = v  # type: ignore
     inputs: Dict[str, Tuple[Type, Any]] = OrderedDict()
     for k, v in signature.parameters.items():  # type: ignore
         annotation = type_hints.get(k, None)
         if annotation is None:
-            err_msg = f"'{k}' has no type. Please add a type annotation to the input parameter."
-            raise annotate_exception_with_code(TypeError(err_msg), fn, k)
+            raise FlyteMissingTypeException(fn=fn, param_name=k)
         default = v.default if v.default is not inspect.Parameter.empty else None
         # Inputs with default values are currently ignored, we may want to look into that in the future
         inputs[k] = (annotation, default)  # type: ignore
@@ -491,7 +507,9 @@ def t(a: int, b: str) -> Dict[str, int]: ...
 
     # This statement results in true for typing.Namedtuple, single and void return types, so this
     # handles Options 1, 2. Even though NamedTuple for us is multi-valued, it's a single value for Python
-    if isinstance(return_annotation, Type) or isinstance(return_annotation, TypeVar):  # type: ignore
+    if hasattr(return_annotation, "__bases__") and (
+        isinstance(return_annotation, Type) or isinstance(return_annotation, TypeVar)  # type: ignore
+    ):
         # isinstance / issubclass does not work for Namedtuple.
         # Options 1 and 2
         bases = return_annotation.__bases__  # type: ignore