Merge branch 'master' of github.com:flyteorg/flytekit into 6277-map-container

-e
Signed-off-by: machichima <nary12321@gmail.com>

Author: machichima

flytekit/clients/auth/token_client.py

@@ -159,7 +159,7 @@ def poll_token_endpoint(
     token_endpoint: str,
     client_id: str,
     audience: typing.Optional[str] = None,
-    scopes: typing.Optional[str] = None,
+    scopes: typing.Optional[typing.List[str]] = None,
     http_proxy_url: typing.Optional[str] = None,
     verify: typing.Optional[typing.Union[bool, str]] = None,
 ) -> typing.Tuple[str, str, int]:

flytekit/clis/sdk_in_container/run.py

@@ -472,10 +472,8 @@ def to_click_option(
         python_type=python_type,
         is_remote=run_level_params.is_remote,
     )
-
     if literal_converter.is_bool() and not default_val:
         default_val = False
-
     description_extra = ""
     if literal_var.type.simple == SimpleType.STRUCT:
         if default_val and not isinstance(default_val, ArtifactQuery):
@@ -517,6 +515,7 @@ def to_click_option(
         required=required,
         help=literal_var.description + description_extra,
         callback=literal_converter.convert,
+        is_flag=literal_converter.is_bool(),
     )
 
 
@@ -870,12 +869,17 @@ def invoke(self, ctx: click.Context) -> typing.Any:
         run_level_params: RunLevelParams = ctx.obj
         r = run_level_params.remote_instance()
         entity = self._fetch_entity(ctx)
+
+        param_defaults = {param.name: param.default for param in ctx.command.params}
+
+        filtered_inputs = {k: param_defaults.get(k, v) if v is None else v for k, v in ctx.params.items()}
+
         run_remote(
             r,
             entity,
             run_level_params.project,
             run_level_params.domain,
-            ctx.params,
+            filtered_inputs,
             run_level_params,
             type_hints=entity.python_interface.inputs if entity.python_interface else None,
         )

flytekit/core/type_engine.py

@@ -16,7 +16,7 @@
 import typing
 from abc import ABC, abstractmethod
 from collections import OrderedDict
-from functools import lru_cache
+from functools import lru_cache, reduce
 from types import GenericAlias
 from typing import Any, Dict, List, NamedTuple, Optional, Type, cast
 
@@ -1089,56 +1089,78 @@ def assert_type(self, t: Type[enum.Enum], v: T):
             raise TypeTransformerFailedError(f"Value {v} is not in Enum {t}")
 
 
-def generate_attribute_list_from_dataclass_json_mixin(schema: dict, schema_name: typing.Any):
+def _handle_json_schema_property(
+    property_key: str,
+    property_val: dict,
+) -> typing.Tuple[str, typing.Any]:
+    """
+    A helper to handle the properties of a JSON schema and returns their equivalent Flyte attribute name and type.
+    """
+
+    # Handle Optional[T] or Union[T1, T2, ...] at the top level for proper recursion
+    if property_val.get("anyOf"):
+        # Sanity check 'anyOf' is not empty
+        assert len(property_val["anyOf"]) > 0
+        # Check that there are no nested Optional or Union types - no need to support that pattern
+        # as it would just add complexity without much benefit
+        # A few examples: Optional[Optional[T]] or Union[T1, T2, Union[T3, T4], etc...]
+        if any(item.get("anyOf") for item in property_val["anyOf"]):
+            raise ValueError(
+                f"The property with name {property_key} has a nested Optional or Union type, this is not allowed for dataclass JSON deserialization."
+            )
+        attr_types = []
+        for item in property_val["anyOf"]:
+            _, attr_type = _handle_json_schema_property(property_key, item)
+            attr_types.append(attr_type)
+
+        # Gather all the types and return a Union[T1, T2, ...]
+        attr_union_type = reduce(lambda x, y: typing.Union[x, y], attr_types)
+        return (property_key, attr_union_type)  # type: ignore
+
+    # Handle enum
+    if property_val.get("enum"):
+        property_type = "enum"
+    else:
+        property_type = property_val["type"]
+
+    # Handle list
+    if property_type == "array":
+        return (property_key, typing.List[_get_element_type(property_val["items"])])  # type: ignore
+    # Handle null types (i.e. None)
+    elif property_type == "null":
+        return (property_key, type(None))  # type: ignore
+    # Handle dataclass and dict
+    elif property_type == "object":
+        # NOTE: No need to handle optional dataclasses here (i.e. checking for property_val.get("anyOf"))
+        # those are handled in the top level of the function with recursion.
+        if property_val.get("additionalProperties"):
+            # For typing.Dict type
+            elem_type = _get_element_type(property_val["additionalProperties"])
+            return (property_key, typing.Dict[str, elem_type])  # type: ignore
+        elif property_val.get("title"):
+            # For nested dataclass
+            sub_schema_name = property_val["title"]
+            return (
+                property_key,
+                typing.cast(GenericAlias, convert_mashumaro_json_schema_to_python_class(property_val, sub_schema_name)),
+            )
+        else:
+            # For untyped dict
+            return (property_key, dict)  # type: ignore
+    elif property_type == "enum":
+        return (property_key, str)  # type: ignore
+    # Handle None, int, float, bool or str
+    else:
+        return (property_key, _get_element_type(property_val))  # type: ignore
+
+
+def generate_attribute_list_from_dataclass_json_mixin(
+    schema: dict,
+    schema_name: typing.Any,
+):
     attribute_list: typing.List[typing.Tuple[Any, Any]] = []
     for property_key, property_val in schema["properties"].items():
-        property_type = ""
-        if property_val.get("anyOf"):
-            property_type = property_val["anyOf"][0]["type"]
-        elif property_val.get("enum"):
-            property_type = "enum"
-        else:
-            property_type = property_val["type"]
-        # Handle list
-        if property_type == "array":
-            attribute_list.append((property_key, typing.List[_get_element_type(property_val["items"])]))  # type: ignore
-        # Handle dataclass and dict
-        elif property_type == "object":
-            if property_val.get("anyOf"):
-                # For optional with dataclass
-                sub_schemea = property_val["anyOf"][0]
-                sub_schemea_name = sub_schemea["title"]
-                attribute_list.append(
-                    (
-                        property_key,
-                        typing.cast(
-                            GenericAlias, convert_mashumaro_json_schema_to_python_class(sub_schemea, sub_schemea_name)
-                        ),
-                    )
-                )
-            elif property_val.get("additionalProperties"):
-                # For typing.Dict type
-                elem_type = _get_element_type(property_val["additionalProperties"])
-                attribute_list.append((property_key, typing.Dict[str, elem_type]))  # type: ignore
-            elif property_val.get("title"):
-                # For nested dataclass
-                sub_schemea_name = property_val["title"]
-                attribute_list.append(
-                    (
-                        property_key,
-                        typing.cast(
-                            GenericAlias, convert_mashumaro_json_schema_to_python_class(property_val, sub_schemea_name)
-                        ),
-                    )
-                )
-            else:
-                # For untyped dict
-                attribute_list.append((property_key, dict))  # type: ignore
-        elif property_type == "enum":
-            attribute_list.append([property_key, str])  # type: ignore
-        # Handle int, float, bool or str
-        else:
-            attribute_list.append([property_key, _get_element_type(property_val)])  # type: ignore
+        attribute_list.append(_handle_json_schema_property(property_key, property_val))
     return attribute_list
 
 

flytekit/types/directory/types.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import json
 import os
 import pathlib
@@ -22,15 +23,24 @@
 
 from flytekit.core.constants import MESSAGEPACK
 from flytekit.core.context_manager import FlyteContext, FlyteContextManager
-from flytekit.core.type_engine import AsyncTypeTransformer, TypeEngine, TypeTransformerFailedError, get_batch_size
+from flytekit.core.type_engine import (
+    AsyncTypeTransformer,
+    TypeEngine,
+    TypeTransformerFailedError,
+    get_batch_size,
+)
 from flytekit.exceptions.user import FlyteAssertion
-from flytekit.extras.pydantic_transformer.decorator import model_serializer, model_validator
+from flytekit.extras.pydantic_transformer.decorator import (
+    model_serializer,
+    model_validator,
+)
 from flytekit.models import types as _type_models
 from flytekit.models.core import types as _core_types
 from flytekit.models.core.types import BlobType
 from flytekit.models.literals import Binary, Blob, BlobMetadata, Literal, Scalar
 from flytekit.models.types import LiteralType
 from flytekit.types.file import FileExt, FlyteFile
+from flytekit.utils.asyn import loop_manager
 
 T = typing.TypeVar("T")
 PathType = typing.Union[str, os.PathLike]
@@ -193,7 +203,19 @@ def __fspath__(self):
         This function should be called by os.listdir as well.
         """
         if not self._downloaded:
-            self._downloader()
+            if isinstance(self._downloader, partial):
+                underlying_func = self._downloader.func
+            else:
+                underlying_func = self._downloader
+
+            if asyncio.iscoroutinefunction(underlying_func):
+                # If the downloader is a coroutine function, we need to run it in the event loop.
+                # This is required when using the Elastic task config with the start method set to 'spawn'.
+                # This is possibly due to pickling and unpickling the task function, which may not properly rebind to the synced version.
+                loop_manager.synced(self._downloader)()
+            else:
+                # If the downloader is not a coroutine, we can just call it directly.
+                self._downloader()
             self._downloaded = True
         return self.path
 

plugins/flytekit-onnx-pytorch/dev-requirements.txt

@@ -22,7 +22,7 @@ humanfriendly==10.0
     # via coloredlogs
 idna==3.4
     # via requests
-jinja2==3.1.5
+jinja2==3.1.6
     # via torch
 markupsafe==2.1.3
     # via jinja2

plugins/flytekit-onnx-scikitlearn/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit", "skl2onnx>=1.10.3", "networkx<3.2; python_version<'3.9'"]
+plugin_requires = ["flytekit", "skl2onnx>=1.10.3,<1.19.0", "networkx<3.2; python_version<'3.9'", "onnx<1.18.0"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-spark/setup.py

@@ -4,7 +4,8 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=1.15.1", "pyspark>=3.4.0", "aiohttp", "flyteidl>=1.11.0b1", "pandas"]
+# TODO: Add support spark 4.0.0, https://github.com/flyteorg/flyte/issues/6478
+plugin_requires = ["flytekit>=1.15.1", "pyspark>=3.4.0,<4.0.0", "aiohttp", "flyteidl>=1.11.0b1", "pandas"]
 
 __version__ = "0.0.0+develop"
 

pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
     "docstring-parser>=0.9.0",
     "flyteidl>=1.15.1",
     "fsspec>=2023.3.0",
-    "gcsfs>=2023.3.0",
+    "gcsfs>=2023.3.0,!=2025.5.0,!=2025.5.0post1",  # Bug in 2025.5.0, 2025.5.0post1 https://github.com/fsspec/gcsfs/issues/687
     "googleapis-common-protos>=1.57",
     "grpcio",
     "grpcio-status",

tests/flytekit/integration/remote/test_remote.py

@@ -16,7 +16,7 @@
 import uuid
 import pytest
 from unittest import mock
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 
 from flytekit import LaunchPlan, kwtypes, WorkflowExecutionPhase, task, workflow
 from flytekit.configuration import Config, ImageConfig, SerializationSettings
@@ -1148,6 +1148,88 @@ def test_execute_workflow_with_dataclass():
     )
     assert out.outputs["o0"] == ""
 
+def test_execute_wf_out_dataclass_with_optional():
+    """Test remote execution of a workflow outputting a dataclass where optional fields are present."""
+    from tests.flytekit.integration.remote.workflows.basic.dataclass_with_optional_wf import MyDataClassWithOptional, MyParentDataClass
+
+    remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
+
+    # Simple case where dataclasses are not nested
+    in_dataclass = MyDataClassWithOptional(foo={"a": 1.0, "b": 2.0}, bar={"c": 3.0, "d": 4.0})
+
+    execution_id = run(
+        "dataclass_with_optional_wf.py",
+        "wf",
+        "--in_dataclass",
+        json.dumps(asdict(in_dataclass)),
+    )
+
+    execution = remote.fetch_execution(name=execution_id, project=PROJECT, domain=DOMAIN)
+    execution = remote.wait(execution=execution, timeout=datetime.timedelta(minutes=10))
+
+    assert asdict(execution.outputs["o0"]) == {
+        "foo": {"a": 1.0, "b": 2.0},
+        "bar": {"c": 3.0, "d": 4.0},
+        "baz": None,
+        "qux": None,
+    }
+
+    # Case where dataclasses are nested (all fields are populated)
+    in_dataclass = MyParentDataClass(
+        child=MyDataClassWithOptional(foo={"a": 1.0, "b": 2.0}, bar={"c": 3.0, "d": 4.0}),
+        a={"a": 1.0, "b": 2.0},
+        b=MyDataClassWithOptional(foo={"a": 1.0, "b": 2.0}, bar={"c": 3.0, "d": 4.0}),
+    )
+
+    execution_id = run(
+        "dataclass_with_optional_wf.py",
+        "wf_nested_dc",
+        "--in_dataclass",
+        json.dumps(asdict(in_dataclass)),
+    )
+
+    execution = remote.fetch_execution(name=execution_id, project=PROJECT, domain=DOMAIN)
+    execution = remote.wait(execution=execution, timeout=datetime.timedelta(minutes=10))
+    assert asdict(execution.outputs["o0"]) == {
+        "child": {
+            "foo": {"a": 1.0, "b": 2.0},
+            "bar": {"c": 3.0, "d": 4.0},
+            "baz": None,
+            "qux": None,
+        },
+        "a": {"a": 1.0, "b": 2.0},
+        "b": {
+            "foo": {"a": 1.0, "b": 2.0},
+            "bar": {"c": 3.0, "d": 4.0},
+            "baz": None,
+            "qux": None,
+        }
+    }
+
+    # Case where dataclasses are nested (optionals are left as None)
+    in_dataclass = MyParentDataClass(
+        child=MyDataClassWithOptional(foo={"a": 1.0, "b": 2.0}, bar={"c": 3.0, "d": 4.0}),
+    )
+
+    execution_id = run(
+        "dataclass_with_optional_wf.py",
+        "wf_nested_dc",
+        "--in_dataclass",
+        json.dumps(asdict(in_dataclass)),
+    )
+
+    execution = remote.fetch_execution(name=execution_id, project=PROJECT, domain=DOMAIN)
+    execution = remote.wait(execution=execution, timeout=datetime.timedelta(minutes=10))
+    assert asdict(execution.outputs["o0"]) == {
+        "child": {
+            "foo": {"a": 1.0, "b": 2.0},
+            "bar": {"c": 3.0, "d": 4.0},
+            "baz": None,
+            "qux": None,
+        },
+        "a": None,
+        "b": None,
+    }
 
 def test_register_wf_twice(register):
     # Register the same workflow again should not raise an error